Repository: snowflakedb/gosnowflake
Pull request #967 (2023-11-13 14:57:27)

SNOW-968719: Remember queryId for failed queries
### Description
Remember queryId for failed queries.

### Checklist
- [x] Code compiles correctly
- [x] Run ``make fmt`` to fix inconsistent formats
- [x] Run ``make lint`` to get lint errors and fix all of them
- [x] Created tests which fail without the change (if possible)
- [x] All tests passing
- [x] Extended the README / documentation, if necessary


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/967?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) Report
> Merging [#967](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/967?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) (56b8ed8) into [master](https://app.codecov.io/gh/snowflakedb/gosnowflake/commit/730c03ce4ec317bc05f071367772af899f2d4b11?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) (730c03c) will **decrease** coverage by `0.09%`.
> Report is 2 commits behind head on master.
> The diff coverage is `83.09%`.

```diff
@@            Coverage Diff             @@
##           master     #967      +/-   ##
==========================================
- Coverage   82.15%   82.06%   -0.09%     
==========================================
  Files          50       49       -1     
  Lines        9570     9626      +56     
==========================================
+ Hits         7862     7900      +38     
- Misses       1708     1726      +18     
```


| [Files](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/967?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) | Coverage Δ | |
|---|---|---|
| [gcs\_storage\_client.go](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/967?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb#diff-Z2NzX3N0b3JhZ2VfY2xpZW50Lmdv) | `87.29% <100.00%> (+12.63%)` | :arrow_up: |
| [authexternalbrowser.go](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/967?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb#diff-YXV0aGV4dGVybmFsYnJvd3Nlci5nbw==) | `57.54% <71.42%> (+0.15%)` | :arrow_up: |
| [dsn.go](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/967?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb#diff-ZHNuLmdv) | `87.95% <60.00%> (+0.07%)` | :arrow_up: |
| [statement.go](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/967?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb#diff-c3RhdGVtZW50Lmdv) | `88.57% <82.22%> (-11.43%)` | :arrow_down: |

... and [9 files with indirect coverage changes](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/967/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb)

:mega: Codecov offers a browser extension for seamless coverage viewing on GitHub. Try it in [Chrome](https://chrome.google.com/webstore/detail/codecov/gedikamndpbemklijjkncpnolildpbgo) or [Firefox](https://addons.mozilla.org/en-US/firefox/addon/codecov/) today!


Patch:
diff --git a/statement.go b/statement.go
index 9ebf941..3647a76 100644
--- a/statement.go
+++ b/statement.go
@@ -1,10 +1,11 @@
-// Copyright (c) 2017-2022 Snowflake Computing Inc. All rights reserved.
+// Copyright (c) 2017-2023 Snowflake Computing Inc. All rights reserved.
 
 package gosnowflake
 
 import (
 	"context"
 	"database/sql/driver"
+	"errors"
 	"fmt"
 )
 
@@ -35,6 +36,7 @@ func (stmt *snowflakeStmt) ExecContext(ctx context.Context, args []driver.NamedV
 	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)
@@ -49,6 +51,7 @@ func (stmt *snowflakeStmt) QueryContext(ctx context.Context, args []driver.Named
 	logger.WithContext(stmt.sc.ctx).Infoln("Stmt.QueryContext")
 	rows, err := stmt.sc.QueryContext(ctx, stmt.query, args)
 	if err != nil {
+		stmt.setQueryIDFromError(err)
 		return nil, err
 	}
 	r, ok := rows.(SnowflakeRows)
@@ -63,6 +66,7 @@ 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)
 	if err != nil {
+		stmt.setQueryIDFromError(err)
 		return nil, err
 	}
 	r, ok := result.(SnowflakeResult)
@@ -77,6 +81,7 @@ func (stmt *snowflakeStmt) Query(args []driver.Value) (driver.Rows, error) {
 	logger.WithContext(stmt.sc.ctx).Infoln("Stmt.Query")
 	rows, err := stmt.sc.Query(stmt.query, args)
 	if err != nil {
+		stmt.setQueryIDFromError(err)
 		return nil, err
 	}
 	r, ok := rows.(SnowflakeRows)
@@ -90,3 +95,10 @@ func (stmt *snowflakeStmt) Query(args []driver.Value) (driver.Rows, error) {
 func (stmt *snowflakeStmt) GetQueryID() string {
 	return stmt.lastQueryID
 }
+
+func (stmt *snowflakeStmt) setQueryIDFromError(err error) {
+	var snowflakeError *SnowflakeError
+	if errors.As(err, &snowflakeError) {
+		stmt.lastQueryID = snowflakeError.QueryID
+	}
+}


Test patch:
diff --git a/statement_test.go b/statement_test.go
index f1b6090..5fe5058 100644
--- a/statement_test.go
+++ b/statement_test.go
@@ -1,4 +1,4 @@
-// Copyright (c) 2020-2022 Snowflake Computing Inc. All rights reserved.
+// Copyright (c) 2020-2023 Snowflake Computing Inc. All rights reserved.
 //lint:file-ignore SA1019 Ignore deprecated methods. We should leave them as-is to keep backward compatibility.
 
 package gosnowflake
@@ -7,6 +7,7 @@ import (
 	"context"
 	"database/sql"
 	"database/sql/driver"
+	"errors"
 	"fmt"
 	"net/http"
 	"net/url"
@@ -39,6 +40,149 @@ func openConn(t *testing.T) *sql.Conn {
 	return conn
 }
 
+func TestFailedQueryIdInSnowflakeError(t *testing.T) {
+	failingQuery := "SELECTT 1"
+	failingExec := "INSERT 1 INTO NON_EXISTENT_TABLE"
+
+	runDBTest(t, func(dbt *DBTest) {
+		testcases := []struct {
+			name  string
+			query string
+			f     func(dbt *DBTest) (any, error)
+		}{
+			{
+				name: "query",
+				f: func(dbt *DBTest) (any, error) {
+					return dbt.query(failingQuery)
+				},
+			},
+			{
+				name: "exec",
+				f: func(dbt *DBTest) (any, error) {
+					return dbt.exec(failingExec)
+				},
+			},
+		}
+
+		for _, tc := range testcases {
+			t.Run(tc.name, func(t *testing.T) {
+				_, err := tc.f(dbt)
+				if err == nil {
+					t.Error("should have failed")
+				}
+				var snowflakeError *SnowflakeError
+				if !errors.As(err, &snowflakeError) {
+					t.Error("should be a SnowflakeError")
+				}
+				if snowflakeError.QueryID == "" {
+					t.Error("QueryID should be set")
+				}
+			})
+		}
+	})
+}
+
+func TestSetFailedQueryId(t *testing.T) {
+	ctx := context.Background()
+	failingQuery := "SELECTT 1"
+	failingExec := "INSERT 1 INTO NON_EXISTENT_TABLE"
+
+	runDBTest(t, func(dbt *DBTest) {
+		testcases := []struct {
+			name  string
+			query string
+			f     func(stmt driver.Stmt) (any, error)
+		}{
+			{
+				name:  "query",
+				query: failingQuery,
+				f: func(stmt driver.Stmt) (any, error) {
+					return stmt.Query(nil)
+				},
+			},
+			{
+				name:  "exec",
+				query: failingExec,
+				f: func(stmt driver.Stmt) (any, error) {
+					return stmt.Exec(nil)
+				},
+			},
+			{
+				name:  "queryContext",
+				query: failingQuery,
+				f: func(stmt driver.Stmt) (any, error) {
+					return stmt.(driver.StmtQueryContext).QueryContext(ctx, nil)
+				},
+			},
+			{
+				name:  "execContext",
+				query: failingExec,
+				f: func(stmt driver.Stmt) (any, error) {
+					return stmt.(driver.StmtExecContext).ExecContext(ctx, 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.Error("should have failed to execute the query")
+					}
+					if stmt.(SnowflakeStmt).GetQueryID() == "" {
+						t.Error("should have set the query id")
+					}
+					return nil
+				})
+				if err != nil {
+					t.Fatal(err)
+				}
+			})
+		}
+	})
+}
+
+func TestAsyncFailQueryId(t *testing.T) {
+	ctx := WithAsyncMode(context.Background())
+	runDBTest(t, func(dbt *DBTest) {
+		err := dbt.conn.Raw(func(x any) error {
+			stmt, err := x.(driver.ConnPrepareContext).PrepareContext(ctx, "SELECTT 1")
+			if err != nil {
+				t.Error(err)
+			}
+			if stmt.(SnowflakeStmt).GetQueryID() != "" {
+				t.Error("queryId should be empty before executing any query")
+			}
+			rows, err := stmt.(driver.StmtQueryContext).QueryContext(ctx, nil)
+			if err != nil {
+				t.Error("should not fail the initial request")
+			}
+			if rows.(SnowflakeRows).GetStatus() != QueryStatusInProgress {
+				t.Error("should be in progress")
+			}
+			// Wait for the query to complete
+			rows.Next(nil)
+			if rows.(SnowflakeRows).GetStatus() != QueryFailed {
+				t.Error("should have failed")
+			}
+			if rows.(SnowflakeRows).GetQueryID() != stmt.(SnowflakeStmt).GetQueryID() {
+				t.Error("last query id should be the same as rows query id")
+			}
+			return nil
+		})
+		if err != nil {
+			t.Fatal(err)
+		}
+	})
+}
+
 func TestGetQueryID(t *testing.T) {
 	ctx := context.Background()
 	conn := openConn(t)