SNOW-898296: Enable HTAP optimisations in tests

snowflakedb · Sep 11, 2023 · e20a34b · e20a34b
1 parent ede1d51
commit e20a34b
Show file tree

Hide file tree

Showing 4 changed files with 180 additions and 4 deletions.
diff --git a/connection.go b/connection.go
@@ -152,10 +152,23 @@ 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
+	println("query: " + query)
+	println("db: " + data.Data.FinalDatabaseName)
+	if data.Data.FinalDatabaseName != "" {
+		sc.cfg.Database = data.Data.FinalDatabaseName
+	}
+	println("schema: " + data.Data.FinalSchemaName)
+	if data.Data.FinalSchemaName != "" {
+		sc.cfg.Schema = data.Data.FinalSchemaName
+	}
+	println("warehouse: " + data.Data.FinalWarehouseName)
+	if data.Data.FinalWarehouseName != "" {
+		sc.cfg.Warehouse = data.Data.FinalWarehouseName
+	}
+	println("role: " + data.Data.FinalRoleName)
+	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
@@ -387,6 +387,8 @@ func runSnowflakeConnTest(t *testing.T, test func(sct *SCTest)) {
 	}
 
 	sct := &SCTest{t, sc}
+	sct.mustExec("ALTER SESSION SET ENABLE_SNOW_654741_FOR_TESTING = true", nil)
+
 	test(sct)
 }
 

diff --git a/htap_test.go b/htap_test.go
@@ -1,10 +1,13 @@
 package gosnowflake
 
 import (
+	"context"
 	"encoding/json"
+	"fmt"
 	"reflect"
 	"strings"
 	"testing"
+	"time"
 )
 
 func TestMarshallAndDecodeOpaqueContext(t *testing.T) {
@@ -412,3 +415,158 @@ func htapTestSnowflakeConn() *snowflakeConn {
 		},
 	}
 }
+
+func TestHTAPOptimizations(t *testing.T) {
+	runSnowflakeConnTest(t, func(sct *SCTest) {
+		runID := time.Now().UnixMilli()
+		t.Run("Schema", func(t *testing.T) {
+			newSchema := fmt.Sprintf("test_schema_%v", runID)
+			if strings.EqualFold(sct.sc.cfg.Schema, newSchema) {
+				t.Errorf("schema should not be switched")
+			}
+
+			sct.mustExec(fmt.Sprintf("CREATE SCHEMA %v", newSchema), nil)
+			defer sct.mustExec(fmt.Sprintf("DROP SCHEMA %v", newSchema), nil)
+
+			if !strings.EqualFold(sct.sc.cfg.Schema, newSchema) {
+				t.Errorf("schema should be switched, expected %v, got %v", newSchema, sct.sc.cfg.Schema)
+			}
+
+			query := sct.mustQuery("SELECT 1", nil)
+			query.Close()
+
+			if !strings.EqualFold(sct.sc.cfg.Schema, newSchema) {
+				t.Errorf("schema should be switched, expected %v, got %v", newSchema, sct.sc.cfg.Schema)
+			}
+		})
+		t.Run("Database", func(t *testing.T) {
+			newDatabase := fmt.Sprintf("test_database_%v", runID)
+			if strings.EqualFold(sct.sc.cfg.Database, newDatabase) {
+				t.Errorf("database should not be switched")
+			}
+
+			// TODO replace with mustExec
+			sct.mustExec(fmt.Sprintf("CREATE DATABASE %v", newDatabase), nil)
+			defer sct.mustExec(fmt.Sprintf("DROP DATABASE %v", newDatabase), nil)
+
+			if !strings.EqualFold(sct.sc.cfg.Database, newDatabase) {
+				t.Errorf("database should be switched, expected %v, got %v", newDatabase, sct.sc.cfg.Database)
+			}
+
+			query := sct.mustQuery("SELECT 1", nil)
+			query.Close()
+
+			if !strings.EqualFold(sct.sc.cfg.Database, newDatabase) {
+				t.Errorf("database should be switched, expected %v, got %v", newDatabase, sct.sc.cfg.Database)
+			}
+		})
+		t.Run("Warehouse", func(t *testing.T) {
+			newWarehouse := fmt.Sprintf("test_warehouse_%v", runID)
+			if strings.EqualFold(sct.sc.cfg.Warehouse, newWarehouse) {
+				t.Errorf("warehouse should not be switched")
+			}
+
+			// TODO replace with mustExec
+			sct.mustExec(fmt.Sprintf("CREATE WAREHOUSE %v", newWarehouse), nil)
+			defer sct.mustExec(fmt.Sprintf("DROP WAREHOUSE %v", newWarehouse), nil)
+
+			if !strings.EqualFold(sct.sc.cfg.Warehouse, newWarehouse) {
+				t.Errorf("warehouse should be switched, expected %v, got %v", newWarehouse, sct.sc.cfg.Warehouse)
+			}
+
+			query := sct.mustQuery("SELECT 1", nil)
+			query.Close()
+
+			if !strings.EqualFold(sct.sc.cfg.Warehouse, newWarehouse) {
+				t.Errorf("warehouse should be switched, expected %v, got %v", newWarehouse, sct.sc.cfg.Warehouse)
+			}
+		})
+		t.Run("Role", func(t *testing.T) {
+			if strings.EqualFold(sct.sc.cfg.Role, "PUBLIC") {
+				t.Errorf("role should not be public for this test")
+			}
+
+			sct.mustExec("USE ROLE public", nil)
+
+			if !strings.EqualFold(sct.sc.cfg.Role, "PUBLIC") {
+				t.Errorf("role should be switched, expected public, got %v", sct.sc.cfg.Warehouse)
+			}
+
+			query := sct.mustQuery("SELECT 1", nil)
+			query.Close()
+
+			if !strings.EqualFold(sct.sc.cfg.Role, "PUBLIC") {
+				t.Errorf("role should be switched, expected public, got %v", sct.sc.cfg.Warehouse)
+			}
+		})
+		t.Run("Session param - DATE_OUTPUT_FORMAT", func(t *testing.T) {
+			if !strings.EqualFold(*sct.sc.cfg.Params["date_output_format"], "YYYY-MM-DD") {
+				t.Errorf("should use default date_output_format, but got: %v", *sct.sc.cfg.Params["date_output_format"])
+			}
+
+			// TODO replace with mustExec
+			sct.mustExec("ALTER SESSION SET DATE_OUTPUT_FORMAT = 'DD-MM-YYYY'", nil)
+			defer sct.mustExec("ALTER SESSION SET DATE_OUTPUT_FORMAT = 'YYYY-MM-DD'", nil)
+
+			if !strings.EqualFold(*sct.sc.cfg.Params["date_output_format"], "DD-MM-YYYY") {
+				t.Errorf("role should be switched, expected public, got %v", sct.sc.cfg.Warehouse)
+			}
+
+			query := sct.mustQuery("SELECT 1", nil)
+			query.Close()
+
+			if !strings.EqualFold(*sct.sc.cfg.Params["date_output_format"], "DD-MM-YYYY") {
+				t.Errorf("role should be switched, expected public, got %v", sct.sc.cfg.Warehouse)
+			}
+		})
+	})
+}
+
+func TestConnIsCleanAfterClose(t *testing.T) {
+	// We create a new db here to not use the default pool as we can leave it in dirty state.
+	t.Skip("Fails, because connection is returned to a pool dirty")
+	ctx := context.Background()
+	runID := time.Now().UnixMilli()
+
+	db := openDB(t)
+	defer db.Close()
+	db.SetMaxOpenConns(1)
+
+	conn, err := db.Conn(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer conn.Close()
+
+	dbt := DBTest{t, conn}
+
+	dbt.mustExec(forceJSON)
+
+	var dbName string
+	rows1 := dbt.mustQuery("SELECT CURRENT_DATABASE()")
+	rows1.Next()
+	rows1.Scan(&dbName)
+
+	newDbName := fmt.Sprintf("test_database_%v", runID)
+	dbt.mustExec("CREATE DATABASE " + newDbName)
+
+	rows1.Close()
+	conn.Close()
+
+	conn2, err := db.Conn(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	dbt2 := DBTest{t, conn2}
+
+	var dbName2 string
+	rows2 := dbt2.mustQuery("SELECT CURRENT_DATABASE()")
+	defer rows2.Close()
+	rows2.Next()
+	rows2.Scan(&dbName2)
+
+	if !strings.EqualFold(dbName, dbName2) {
+		t.Errorf("fresh connection from pool should have original database")
+	}
+}
diff --git 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
 }
-Original file line number
+Diff line change
@@ Expand Up @@
     	}
     	sct := &SCTest{t, sc}
+    	sct.mustExec("ALTER SESSION SET ENABLE_SNOW_654741_FOR_TESTING = true", nil)
     	test(sct)
     }
@@ Expand Down @@