From 3cf347f142c20c17b8e1108b00c4e91466bdc81a Mon Sep 17 00:00:00 2001
From: Luke Powlett <luke.powlett@unity3d.com>
Date: Wed, 22 Mar 2023 11:49:27 +0000
Subject: [PATCH] SNOW-752633: honour KeepSessionAlive config in DSN

---
 dsn.go      |  8 ++++++++
 dsn_test.go | 27 +++++++++++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/dsn.go b/dsn.go
index ae4decd66..1913fcdcd 100644
--- a/dsn.go
+++ b/dsn.go
@@ -186,6 +186,9 @@ func DSN(cfg *Config) (dsn string, err error) {
 			params.Add(k, *v)
 		}
 	}
+	if cfg.KeepSessionAlive {
+		params.Add("client_session_keep_alive", strconv.FormatBool(cfg.KeepSessionAlive))
+	}
 	if cfg.PrivateKey != nil {
 		privateKeyInBytes, err := marshalPKCS8PrivateKey(cfg.PrivateKey)
 		if err != nil {
@@ -597,6 +600,11 @@ func parseDSNParams(cfg *Config, params string) (err error) {
 
 		case "token":
 			cfg.Token = value
+		case "client_session_keep_alive":
+			cfg.KeepSessionAlive, err = strconv.ParseBool(value)
+			if err != nil {
+				return
+			}
 		case "privateKey":
 			var decodeErr error
 			block, decodeErr := base64.URLEncoding.DecodeString(value)
diff --git a/dsn_test.go b/dsn_test.go
index 037480809..c453b33be 100644
--- a/dsn_test.go
+++ b/dsn_test.go
@@ -491,6 +491,20 @@ func TestParseDSN(t *testing.T) {
 			ocspMode: ocspModeFailOpen,
 			err:      nil,
 		},
+		{
+			dsn: "u:p@a?database=d?client_session_keep_alive=true",
+			config: &Config{
+				Account: "a", User: "u", Password: "p",
+				Protocol: "https", Host: "a.snowflakecomputing.com", Port: 443,
+				Database: "d", Schema: "",
+				JWTExpireTimeout:          defaultJWTTimeout,
+				OCSPFailOpen:              OCSPFailOpenTrue,
+				ValidateDefaultParameters: ConfigBoolTrue,
+				ClientTimeout:             defaultClientTimeout,
+				KeepSessionAlive:          true,
+			},
+			ocspMode: ocspModeFailOpen,
+		},
 	}
 
 	for i, test := range testcases {
@@ -573,6 +587,10 @@ func TestParseDSN(t *testing.T) {
 				t.Fatalf("%d: Failed to match ClientTimeout. expected: %v, got: %v",
 					i, test.config.ClientTimeout, cfg.ClientTimeout)
 			}
+			if test.config.KeepSessionAlive != cfg.KeepSessionAlive {
+				t.Fatalf("%d: Failed to match KeepSessionAlive. expected: %v, got: %v",
+					i, test.config.KeepSessionAlive, cfg.KeepSessionAlive)
+			}
 		case test.err != nil:
 			driverErrE, okE := test.err.(*SnowflakeError)
 			driverErrG, okG := err.(*SnowflakeError)
@@ -856,6 +874,15 @@ func TestDSN(t *testing.T) {
 			},
 			dsn: "u:p@a.b.c.snowflakecomputing.com:443?clientTimeout=300&ocspFailOpen=true&region=b.c&validateDefaultParameters=true",
 		},
+		{
+			cfg: &Config{
+				User:             "u",
+				Password:         "p",
+				Account:          "a-aofnadsf.somewhere.azure",
+				KeepSessionAlive: true,
+			},
+			dsn: "u:p@a-aofnadsf.somewhere.azure.snowflakecomputing.com:443?ocspFailOpen=true&region=somewhere.azure&validateDefaultParameters=true&client_session_keep_alive=true",
+		},
 	}
 	for _, test := range testcases {
 		dsn, err := DSN(test.cfg)