SNOW-978164: Fix stmt.Exec for DML

connection.go

@@ -333,9 +333,20 @@ 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 isSelect(data.Data.StatementTypeID) {
+		logger.WithContext(ctx).Debugf("SELECT")
+		return &snowflakeResult{
+			affectedRows: 0,
+			insertID:     -1,
+			queryID:      data.Data.QueryID,
+		}, nil
 	}
 	logger.Debug("DDL")
-	return driver.ResultNoRows, nil
+	return &snowflakeResult{
+		affectedRows: 0,
+		insertID:     -1,
+		queryID:      data.Data.QueryID,
+	}, nil
 }
 
 func (sc *snowflakeConn) QueryContext(

connection_util.go

@@ -197,7 +197,14 @@ func isDml(v int64) bool {
 	return statementTypeIDDml <= v && v <= statementTypeIDMultiTableInsert
 }
 
+func isSelect(v int64) bool {
+	return v == statementTypeIDSelect
+}
+
 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)

rows_test.go

@@ -74,22 +74,6 @@ func TestRowsClose(t *testing.T) {
 	})
 }
 
-func TestResultNoRows(t *testing.T) {
-	// DDL
-	runDBTest(t, func(dbt *DBTest) {
-		row, err := dbt.exec("CREATE OR REPLACE TABLE test(c1 int)")
-		if err != nil {
-			t.Fatalf("failed to execute DDL. err: %v", err)
-		}
-		if _, err = row.RowsAffected(); err == nil {
-			t.Fatal("should have failed to get RowsAffected")
-		}
-		if _, err = row.LastInsertId(); err == nil {
-			t.Fatal("should have failed to get LastInsertID")
-		}
-	})
-}
-
 func TestRowsWithoutChunkDownloader(t *testing.T) {
 	sts1 := "1"
 	sts2 := "Test1"

statement_test.go

@@ -40,6 +40,68 @@ func openConn(t *testing.T) *sql.Conn {
 	return conn
 }
 
+func TestDMLExec(t *testing.T) {
+	query := "SELECT 1"
+	runDBTest(t, func(dbt *DBTest) {
+		testcases := []struct {
+			name string
+			f    func(dbt *DBTest) (any, error)
+		}{
+			{
+				name: "Exec",
+				f: func(dbt *DBTest) (any, error) {
+					stmt, _ := dbt.prepare(query)
+					return stmt.Exec()
+				},
+			},
+			{
+				name: "ExecContext",
+				f: func(dbt *DBTest) (any, error) {
+					stmt, _ := dbt.prepare(query)
+					return stmt.ExecContext(context.Background())
+				},
+			},
+		}
+		for _, tc := range testcases {
+			_, err := tc.f(dbt)
+			if err != nil {
+				t.Error(err)
+			}
+		}
+	})
+}
+
+func TestDDLExec(t *testing.T) {
+	query := "CREATE OR REPLACE TABLE TestDDLExec (num NUMBER)"
+	runDBTest(t, func(dbt *DBTest) {
+		testcases := []struct {
+			name string
+			f    func(dbt *DBTest) (any, error)
+		}{
+			{
+				name: "Exec",
+				f: func(dbt *DBTest) (any, error) {
+					stmt, _ := dbt.prepare(query)
+					return stmt.Exec()
+				},
+			},
+			{
+				name: "ExecContext",
+				f: func(dbt *DBTest) (any, error) {
+					stmt, _ := dbt.prepare(query)
+					return stmt.ExecContext(context.Background())
+				},
+			},
+		}
+		for _, tc := range testcases {
+			_, err := tc.f(dbt)
+			if err != nil {
+				t.Error(err)
+			}
+		}
+	})
+}
+
 func TestFailedQueryIdInSnowflakeError(t *testing.T) {
 	failingQuery := "SELECTT 1"
 	failingExec := "INSERT 1 INTO NON_EXISTENT_TABLE"