From 1ad992d78ed28895339b52db210da63396c3f58b Mon Sep 17 00:00:00 2001
From: Piotr Fus <piotr.fus@snowflake.com>
Date: Tue, 5 Sep 2023 12:23:42 +0200
Subject: [PATCH] SNOW-898296: Enable HTAP optimisations in tests

---
 connection.go     | 16 ++++++++++++----
 driver_test.go    |  2 ++
 statement_test.go |  3 +++
 3 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/connection.go b/connection.go
index a48af3600..d07d7da89 100644
--- a/connection.go
+++ b/connection.go
@@ -152,10 +152,18 @@ func (sc *snowflakeConn) exec(
 	}
 
 	logger.WithContext(ctx).Info("Exec/Query SUCCESS")
-	sc.cfg.Database = data.Data.FinalDatabaseName
-	sc.cfg.Schema = data.Data.FinalSchemaName
-	sc.cfg.Role = data.Data.FinalRoleName
-	sc.cfg.Warehouse = data.Data.FinalWarehouseName
+	if data.Data.FinalDatabaseName != "" {
+		sc.cfg.Database = data.Data.FinalDatabaseName
+	}
+	if data.Data.FinalSchemaName != "" {
+		sc.cfg.Schema = data.Data.FinalSchemaName
+	}
+	if data.Data.FinalWarehouseName != "" {
+		sc.cfg.Warehouse = data.Data.FinalWarehouseName
+	}
+	if data.Data.FinalRoleName != "" {
+		sc.cfg.Role = data.Data.FinalRoleName
+	}
 	sc.populateSessionParameters(data.Data.Parameters)
 	return data, err
 }
diff --git a/driver_test.go b/driver_test.go
index 653c62777..e718cc630 100644
--- a/driver_test.go
+++ b/driver_test.go
@@ -342,6 +342,8 @@ func runSnowflakeConnTest(t *testing.T, test func(sc *snowflakeConn)) {
 		t.Fatal(err)
 	}
 
+	sc.Exec("ALTER SESSION SET ENABLE_SNOW_654741_FOR_TESTING = true", nil)
+
 	test(sc)
 }
 
diff --git a/statement_test.go b/statement_test.go
index 68f1e7a7b..963e0a61c 100644
--- a/statement_test.go
+++ b/statement_test.go
@@ -36,6 +36,9 @@ func openConn(t *testing.T) *sql.Conn {
 	if conn, err = db.Conn(context.Background()); err != nil {
 		t.Fatalf("failed to open connection: %v", err)
 	}
+
+	conn.ExecContext(context.Background(), "ALTER SESSION SET ENABLE_SNOW_654741_FOR_TESTING = true")
+
 	return conn
 }