SNOW-911181: Add parameter to disable query context cache (#912)

connection.go

@@ -141,7 +141,14 @@ func (sc *snowflakeConn) exec(
 		return nil, err
 	}
 
-	sc.queryContextCache.add(sc, data.Data.QueryContext.Entries...)
+	if !sc.cfg.DisableQueryContextCache && data.Data.QueryContext != nil {
+		queryContext, err := extractQueryContext(data)
+		if err != nil {
+			logger.Errorf("error while decoding query context: ", err)
+		} else {
+			sc.queryContextCache.add(sc, queryContext.Entries...)
+		}
+	}
 
 	// handle PUT/GET commands
 	if isFileTransfer(query) {
@@ -160,6 +167,12 @@ func (sc *snowflakeConn) exec(
 	return data, err
 }
 
+func extractQueryContext(data *execResponse) (queryContext, error) {
+	var queryContext queryContext
+	err := json.Unmarshal(data.Data.QueryContext, &queryContext)
+	return queryContext, err
+}
+
 func (sc *snowflakeConn) Begin() (driver.Tx, error) {
 	return sc.BeginTx(sc.ctx, driver.TxOptions{})
 }

doc.go

@@ -118,6 +118,9 @@ The following connection parameters are supported:
   - tracing: Specifies the logging level to be used. Set to error by default.
     Valid values are trace, debug, info, print, warning, error, fatal, panic.
 
+  - disableQueryContextCache: disables parsing of query context returned from server and resending it to server as well.
+    Default value is false.
+
 All other parameters are interpreted as session parameters (https://docs.snowflake.com/en/sql-reference/parameters.html).
 For example, the TIMESTAMP_OUTPUT_FORMAT session parameter can be set by adding:
 

dsn.go

@@ -97,6 +97,8 @@ type Config struct {
 	IDToken                        string     // Internally used to cache the Id Token for external browser
 	ClientRequestMfaToken          ConfigBool // When true the MFA token is cached in the credential manager. True by default in Windows/OSX. False for Linux.
 	ClientStoreTemporaryCredential ConfigBool // When true the ID token is cached in the credential manager. True by default in Windows/OSX. False for Linux.
+
+	DisableQueryContextCache bool // Should HTAP query context cache be disabled
 }
 
 // Validate enables testing if config is correct.
@@ -230,6 +232,9 @@ func DSN(cfg *Config) (dsn string, err error) {
 	if cfg.TmpDirPath != "" {
 		params.Add("tmpDirPath", cfg.TmpDirPath)
 	}
+	if cfg.DisableQueryContextCache {
+		params.Add("disableQueryContextCache", "true")
+	}
 
 	params.Add("ocspFailOpen", strconv.FormatBool(cfg.OCSPFailOpen != OCSPFailOpenFalse))
 
@@ -702,6 +707,13 @@ func parseDSNParams(cfg *Config, params string) (err error) {
 			cfg.Tracing = value
 		case "tmpDirPath":
 			cfg.TmpDirPath = value
+		case "disableQueryContextCache":
+			var b bool
+			b, err = strconv.ParseBool(value)
+			if err != nil {
+				return
+			}
+			cfg.DisableQueryContextCache = b
 		default:
 			if cfg.Params == nil {
 				cfg.Params = make(map[string]*string)

dsn_test.go

@@ -572,9 +572,10 @@ func TestParseDSN(t *testing.T) {
 				Account: "a", User: "u", Password: "p",
 				Protocol: "https", Host: "a.r.c.snowflakecomputing.com", Port: 443,
 				Database: "db", Schema: "s", ValidateDefaultParameters: ConfigBoolTrue, OCSPFailOpen: OCSPFailOpenTrue,
-				ClientTimeout:          300 * time.Second,
-				JWTClientTimeout:       45 * time.Second,
-				ExternalBrowserTimeout: defaultExternalBrowserTimeout,
+				ClientTimeout:            300 * time.Second,
+				JWTClientTimeout:         45 * time.Second,
+				ExternalBrowserTimeout:   defaultExternalBrowserTimeout,
+				DisableQueryContextCache: false,
 			},
 			ocspMode: ocspModeFailOpen,
 			err:      nil,
@@ -593,6 +594,20 @@ func TestParseDSN(t *testing.T) {
 			ocspMode: ocspModeFailOpen,
 			err:      nil,
 		},
+		{
+			dsn: "u:p@a.r.c.snowflakecomputing.com/db/s?account=a.r.c&disableQueryContextCache=true",
+			config: &Config{
+				Account: "a", User: "u", Password: "p",
+				Protocol: "https", Host: "a.r.c.snowflakecomputing.com", Port: 443,
+				Database: "db", Schema: "s", ValidateDefaultParameters: ConfigBoolTrue, OCSPFailOpen: OCSPFailOpenTrue,
+				ClientTimeout:            defaultClientTimeout,
+				JWTClientTimeout:         defaultJWTClientTimeout,
+				ExternalBrowserTimeout:   defaultExternalBrowserTimeout,
+				DisableQueryContextCache: true,
+			},
+			ocspMode: ocspModeFailOpen,
+			err:      nil,
+		},
 	}
 
 	for _, at := range []AuthType{AuthTypeExternalBrowser, AuthTypeOAuth} {
@@ -743,6 +758,9 @@ func TestParseDSN(t *testing.T) {
 				if test.config.TmpDirPath != cfg.TmpDirPath {
 					t.Fatalf("%v: Failed to match TmpDirPatch. expected: %v, got: %v", i, test.config.TmpDirPath, cfg.TmpDirPath)
 				}
+				if test.config.DisableQueryContextCache != cfg.DisableQueryContextCache {
+					t.Fatalf("%v: Failed to match DisableQueryContextCache. expected: %v, got: %v", i, test.config.DisableQueryContextCache, cfg.DisableQueryContextCache)
+				}
 			case test.err != nil:
 				driverErrE, okE := test.err.(*SnowflakeError)
 				driverErrG, okG := err.(*SnowflakeError)
@@ -1133,6 +1151,15 @@ func TestDSN(t *testing.T) {
 			},
 			dsn: "u:p@a.b.c.snowflakecomputing.com:443?ocspFailOpen=true&region=b.c&tmpDirPath=%2Ftmp&validateDefaultParameters=true",
 		},
+		{
+			cfg: &Config{
+				User:                     "u",
+				Password:                 "p",
+				Account:                  "a.b.c",
+				DisableQueryContextCache: true,
+			},
+			dsn: "u:p@a.b.c.snowflakecomputing.com:443?disableQueryContextCache=true&ocspFailOpen=true&region=b.c&validateDefaultParameters=true",
+		},
 	}
 	for _, test := range testcases {
 		t.Run(test.dsn, func(t *testing.T) {

htap.go

@@ -11,6 +11,10 @@ const (
 	defaultQueryContextCacheSize   = 5
 )
 
+type queryContext struct {
+	Entries []queryContextEntry `json:"entries,omitempty"`
+}
+
 type queryContextEntry struct {
 	ID        int   `json:"id"`
 	Timestamp int64 `json:"timestamp"`

htap_test.go

@@ -412,3 +412,17 @@ func htapTestSnowflakeConn() *snowflakeConn {
 		},
 	}
 }
+
+func TestQueryContextCacheDisabled(t *testing.T) {
+	origDsn := dsn
+	defer func() {
+		dsn = origDsn
+	}()
+	dsn += "&disableQueryContextCache=true"
+	runSnowflakeConnTest(t, func(sct *SCTest) {
+		sct.mustExec("SELECT 1", nil)
+		if len(sct.sc.queryContextCache.entries) > 0 {
+			t.Error("should not contain any entries")
+		}
+	})
+}

query.go

@@ -3,6 +3,7 @@
 package gosnowflake
 
 import (
+	"encoding/json"
 	"time"
 )
 
@@ -120,9 +121,7 @@ type execResponseData struct {
 	Operation               string                `json:"operation,omitempty"`
 
 	// HTAP
-	QueryContext struct {
-		Entries []queryContextEntry `json:"entries,omitempty"`
-	} `json:"queryContext,omitempty"`
+	QueryContext json.RawMessage `json:"queryContext,omitempty"`
 }
 
 type execResponse struct {