SNOW-978164: Fix stmt.Exec for DML

connection.go

@@ -60,6 +60,13 @@ const (
 	queryResultType     resultType = "query"
 )
 
+type execKey string
+
+const (
+	executionType          execKey = "executionType"
+	executionTypeStatement string  = "statement"
+)
+
 const privateLinkSuffix = "privatelink.snowflakecomputing.com"
 
 type snowflakeConn struct {
@@ -333,8 +340,17 @@ func (sc *snowflakeConn) ExecContext(
 		}, nil // last insert id is not supported by Snowflake
 	} else if isMultiStmt(&data.Data) {
 		return sc.handleMultiExec(ctx, data.Data)
+	} else if isDql(&data.Data) {
+		logger.WithContext(ctx).Debugf("DQL")
+		if isStatementContext(ctx) {
+			return &snowflakeResultNoRows{queryID: data.Data.QueryID}, nil
+		}
+		return driver.ResultNoRows, nil
 	}
 	logger.Debug("DDL")
+	if isStatementContext(ctx) {
+		return &snowflakeResultNoRows{queryID: data.Data.QueryID}, nil
+	}
 	return driver.ResultNoRows, nil
 }
 

connection_util.go

@@ -197,7 +197,14 @@
 	return statementTypeIDDml <= v && v <= statementTypeIDMultiTableInsert
 }
 
+func isDql(data *execResponseData) bool {
+	return data.StatementTypeID == statementTypeIDSelect && !isMultiStmt(data)
+}
+
 func updateRows(data execResponseData) (int64, error) {
+	if data.RowSet == nil {
+		return 0, nil
+	}
 	var count int64
 	for i, n := 0, len(data.RowType); i < n; i++ {
 		v, err := strconv.ParseInt(*data.RowSet[0][i], 10, 64)
@@ -292,3 +299,8 @@
 	}
 	return nil
 }
+
+func isStatementContext(ctx context.Context) bool {
+	v := ctx.Value(executionType)
+	return v == executionTypeStatement
+}

go.sum

@@ -134,8 +134,6 @@ github.com/zeebo/xxh3 v1.0.2 h1:xZmwmqxHZA8AI603jOQ0tMqmBr9lPeFwGg6d+xy9DC0=
 github.com/zeebo/xxh3 v1.0.2/go.mod h1:5NWz9Sef7zIDm2JHfFlcQvNekmcEl9ekUZQQKCYaDcA=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
 golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
-golang.org/x/crypto v0.7.0 h1:AvwMYaRytfdeVt3u6mLaxYtErKYjxA2OXjJ1HHq6t3A=
-golang.org/x/crypto v0.7.0/go.mod h1:pYwdfH91IfpZVANVyUOhSIPZaFoJGxTFbZhFTx+dXZU=
 golang.org/x/crypto v0.15.0 h1:frVn1TEaCEaZcn3Tmd7Y2b5KKPaZ+I32Q2OA3kYp5TA=
 golang.org/x/crypto v0.15.0/go.mod h1:4ChreQoLWfG3xLDer1WdlH5NdlQ3+mwnQq1YTKY+72g=
 golang.org/x/exp v0.0.0-20230206171751-46f607a40771 h1:xP7rWLUr1e1n2xkK5YB4LI0hPEy3LJC6Wk+D4pGlOJg=
@@ -145,8 +143,6 @@ golang.org/x/mod v0.8.0 h1:LUYupSeNrTNCGzR/hVBk2NHZO4hXcVaW1k4Qx7rjPx8=
 golang.org/x/mod v0.8.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs=
 golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
 golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
-golang.org/x/net v0.8.0 h1:Zrh2ngAOFYneWTAIAPethzeaQLuHwhuBkuV6ZiRnUaQ=
-golang.org/x/net v0.8.0/go.mod h1:QVkue5JL9kW//ek3r6jTKnTFis1tRmNAW2P1shuFdJc=
 golang.org/x/net v0.18.0 h1:mIYleuAkSbHh0tCv7RvjL3F6ZVbLjq4+R7zbOn3Kokg=
 golang.org/x/net v0.18.0/go.mod h1:/czyP5RqHAH4odGYxBJ1qz0+CE5WZ+2j1YgoEo8F2jQ=
 golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
@@ -158,17 +154,11 @@ golang.org/x/sys v0.0.0-20210616045830-e2b7044e8c71/go.mod h1:oPkhp1MJrh7nUepCBc
 golang.org/x/sys v0.0.0-20210819135213-f52c844e1c1c/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20220704084225-05e143d24a9e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.6.0 h1:MVltZSvRTcU2ljQOhs94SXPftV6DCNnZViHeQps87pQ=
-golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.14.0 h1:Vz7Qs629MkJkGyHxUlRHizWJRG2j8fbQKjELVSNhy7Q=
 golang.org/x/sys v0.14.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
-golang.org/x/term v0.6.0 h1:clScbb1cHjoCkyRbWwBEUZ5H/tIFu5TAXIqaZD0Gcjw=
-golang.org/x/term v0.6.0/go.mod h1:m6U89DPEgQRMq3DNkDClhWw02AUbt2daBVO4cn4Hv9U=
 golang.org/x/term v0.14.0 h1:LGK9IlZ8T9jvdy6cTdfKUCltatMFOehAQo9SRC46UQ8=
 golang.org/x/term v0.14.0/go.mod h1:TySc+nGkYR6qt8km8wUhuFRTVSMIX3XPR58y2lC8vww=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
-golang.org/x/text v0.8.0 h1:57P1ETyNKtuIjB4SRd15iJxuhj8Gc416Y78H3qgMh68=
-golang.org/x/text v0.8.0/go.mod h1:e1OnstbJyHTd6l/uOt8jFFHp6TRDWZR/bV3emEE/zU8=
 golang.org/x/text v0.14.0 h1:ScX5w1eTa3QqT8oi6+ziP7dTV1S2+ALU0bI+0zXKWiQ=
 golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
 golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=

result.go

@@ -2,6 +2,8 @@
 
 package gosnowflake
 
+import "errors"
+
 type queryStatus string
 
 const (
@@ -73,3 +75,19 @@
 	}
 	return nil
 }
+
+type snowflakeResultNoRows struct {
+	queryID string
+}
+
+func (*snowflakeResultNoRows) LastInsertId() (int64, error) {
+	return 0, errors.New("no LastInsertId available after DDL statement")
+}
+
+func (*snowflakeResultNoRows) RowsAffected() (int64, error) {
+	return 0, errors.New("no RowsAffected available after DDL statement")
+}
+
+func (rnr *snowflakeResultNoRows) GetQueryID() string {
+	return rnr.queryID
+}

statement.go

@@ -34,17 +34,7 @@
 
 func (stmt *snowflakeStmt) ExecContext(ctx context.Context, args []driver.NamedValue) (driver.Result, error) {
 	logger.WithContext(stmt.sc.ctx).Infoln("Stmt.ExecContext")
-	result, err := stmt.sc.ExecContext(ctx, stmt.query, args)
-	if err != nil {
-		stmt.setQueryIDFromError(err)
-		return nil, err
-	}
-	r, ok := result.(SnowflakeResult)
-	if !ok {
-		return nil, fmt.Errorf("interface convertion. expected type SnowflakeResult but got %T", result)
-	}
-	stmt.lastQueryID = r.GetQueryID()
-	return result, err
+	return stmt.execInternal(ctx, args)
 }
 
 func (stmt *snowflakeStmt) QueryContext(ctx context.Context, args []driver.NamedValue) (driver.Rows, error) {
@@ -64,11 +54,25 @@
 
 func (stmt *snowflakeStmt) Exec(args []driver.Value) (driver.Result, error) {
 	logger.WithContext(stmt.sc.ctx).Infoln("Stmt.Exec")
-	result, err := stmt.sc.Exec(stmt.query, args)
+	return stmt.execInternal(context.Background(), toNamedValues(args))
+}
+
+func (stmt *snowflakeStmt) execInternal(ctx context.Context, args []driver.NamedValue) (driver.Result, error) {
+	logger.WithContext(stmt.sc.ctx).Debugln("Stmt.execInternal")
+	if ctx == nil {
+		ctx = context.Background()
+	}
+	stmtCtx := context.WithValue(ctx, executionType, executionTypeStatement)
+	result, err := stmt.sc.ExecContext(stmtCtx, stmt.query, args)
 	if err != nil {
 		stmt.setQueryIDFromError(err)
 		return nil, err
 	}
+	rnr, ok := result.(*snowflakeResultNoRows)
+	if ok {
+		stmt.lastQueryID = rnr.GetQueryID()
+		return driver.ResultNoRows, nil
+	}
 	r, ok := result.(SnowflakeResult)
 	if !ok {
 		return nil, fmt.Errorf("interface convertion. expected type SnowflakeResult but got %T", result)

statement_test.go

@@ -40,6 +40,102 @@ func openConn(t *testing.T) *sql.Conn {
 	return conn
 }
 
+func TestExecStmt(t *testing.T) {
+	dqlQuery := "SELECT 1"
+	dmlQuery := "INSERT INTO TestDDLExec VALUES (1)"
+	ddlQuery := "CREATE OR REPLACE TABLE TestDDLExec (num NUMBER)"
+	multiStmtQuery := "DELETE FROM TestDDLExec;\n" +
+		"SELECT 1;\n" +
+		"SELECT 2;"
+	ctx := context.Background()
+	multiStmtCtx, err := WithMultiStatement(ctx, 3)
+	if err != nil {
+		t.Error(err)
+	}
+	runDBTest(t, func(dbt *DBTest) {
+		dbt.mustExec(ddlQuery)
+		defer dbt.mustExec("DROP TABLE IF EXISTS TestDDLExec")
+		testcases := []struct {
+			name  string
+			query string
+			f     func(stmt driver.Stmt) (any, error)
+		}{
+			{
+				name:  "dql Exec",
+				query: dqlQuery,
+				f: func(stmt driver.Stmt) (any, error) {
+					return stmt.Exec(nil)
+				},
+			},
+			{
+				name:  "dql ExecContext",
+				query: dqlQuery,
+				f: func(stmt driver.Stmt) (any, error) {
+					return stmt.(driver.StmtExecContext).ExecContext(ctx, nil)
+				},
+			},
+			{
+				name:  "ddl Exec",
+				query: ddlQuery,
+				f: func(stmt driver.Stmt) (any, error) {
+					return stmt.Exec(nil)
+				},
+			},
+			{
+				name:  "ddl ExecContext",
+				query: ddlQuery,
+				f: func(stmt driver.Stmt) (any, error) {
+					return stmt.(driver.StmtExecContext).ExecContext(ctx, nil)
+				},
+			},
+			{
+				name:  "dml Exec",
+				query: dmlQuery,
+				f: func(stmt driver.Stmt) (any, error) {
+					return stmt.Exec(nil)
+				},
+			},
+			{
+				name:  "dml ExecContext",
+				query: dmlQuery,
+				f: func(stmt driver.Stmt) (any, error) {
+					return stmt.(driver.StmtExecContext).ExecContext(ctx, nil)
+				},
+			},
+			{
+				name:  "multistmt ExecContext",
+				query: multiStmtQuery,
+				f: func(stmt driver.Stmt) (any, error) {
+					return stmt.(driver.StmtExecContext).ExecContext(multiStmtCtx, nil)
+				},
+			},
+		}
+		for _, tc := range testcases {
+			t.Run(tc.name, func(t *testing.T) {
+				err := dbt.conn.Raw(func(x any) error {
+					stmt, err := x.(driver.ConnPrepareContext).PrepareContext(ctx, tc.query)
+					if err != nil {
+						t.Error(err)
+					}
+					if stmt.(SnowflakeStmt).GetQueryID() != "" {
+						t.Error("queryId should be empty before executing any query")
+					}
+					if _, err := tc.f(stmt); err != nil {
+						t.Errorf("should have not failed to execute the query, err: %s\n", err)
+					}
+					if stmt.(SnowflakeStmt).GetQueryID() == "" {
+						t.Error("should have set the query id")
+					}
+					return nil
+				})
+				if err != nil {
+					t.Fatal(err)
+				}
+			})
+		}
+	})
+}
+
 func TestFailedQueryIdInSnowflakeError(t *testing.T) {
 	failingQuery := "SELECTT 1"
 	failingExec := "INSERT 1 INTO NON_EXISTENT_TABLE"