From 213a5ac4a5e9caa0fb00f802f000deb6f1f2d159 Mon Sep 17 00:00:00 2001 From: Piotr Fus Date: Mon, 11 Sep 2023 11:40:18 +0200 Subject: [PATCH] SNOW-911181: Add parameter to disable query context cache (#912) --- connection.go | 15 ++++++++++++++- doc.go | 3 +++ dsn.go | 12 ++++++++++++ dsn_test.go | 33 ++++++++++++++++++++++++++++++--- htap.go | 4 ++++ htap_test.go | 14 ++++++++++++++ query.go | 5 ++--- 7 files changed, 79 insertions(+), 7 deletions(-) diff --git a/connection.go b/connection.go index a48af3600..4ea8dbd77 100644 --- a/connection.go +++ b/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{}) } diff --git a/doc.go b/doc.go index 84b859803..93def4075 100644 --- a/doc.go +++ b/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: diff --git a/dsn.go b/dsn.go index bf6710cc7..d5fa2ab73 100644 --- a/dsn.go +++ b/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) diff --git a/dsn_test.go b/dsn_test.go index 362ad21cc..0086e2cf5 100644 --- a/dsn_test.go +++ b/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®ion=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®ion=b.c&validateDefaultParameters=true", + }, } for _, test := range testcases { t.Run(test.dsn, func(t *testing.T) { diff --git a/htap.go b/htap.go index 6c33d4f48..d8730be36 100644 --- a/htap.go +++ b/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"` diff --git a/htap_test.go b/htap_test.go index aac102994..8aae85810 100644 --- a/htap_test.go +++ b/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") + } + }) +} diff --git a/query.go b/query.go index 5d7dff053..ffd3983cd 100644 --- a/query.go +++ b/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 {