Repository: snowflakedb/gosnowflake
Pull request #908 (2023-09-05 10:24:11)

SNOW-898296: Enable HTAP optimisations in tests
### Description
Enables htap optimizations - not sending db/schema/warehouse/session parameters.

### Checklist
- [x] Code compiles correctly
- [x] Run ``make fmt`` to fix inconsistent formats
- [x] Run ``make lint`` to get lint errors and fix all of them
- [x] Created tests which fail without the change (if possible)
- [x] All tests passing
- [x] Extended the README / documentation, if necessary


sfc-gh-dheyman-1: what is the purpose of `aaa_test.go`? perhaps that is not the best name>
sfc-gh-pfus: > what is the purpose of `aaa_test.go`? perhaps that is not the best name?

The purpose is to print server version - sometimes if we depend on the brand new features on the server side it's good to know current version. Naming is of course to run it as soon as possible.

Patch:
diff --git a/connection.go b/connection.go
index 4638175..7f716df 100644
--- a/connection.go
+++ b/connection.go
@@ -165,10 +165,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
 }


Test patch:
diff --git a/aaa_test.go b/aaa_test.go
new file mode 100644
index 0000000..de86ae1
--- /dev/null
+++ b/aaa_test.go
@@ -0,0 +1,15 @@
+package gosnowflake
+
+import "testing"
+
+func TestShowServerVersion(t *testing.T) {
+	runDBTest(t, func(dbt *DBTest) {
+		rows := dbt.mustQuery("SELECT CURRENT_VERSION()")
+		defer rows.Close()
+
+		var version string
+		rows.Next()
+		rows.Scan(&version)
+		println(version)
+	})
+}
diff --git a/htap_test.go b/htap_test.go
index a724424..b18df01 100644
--- a/htap_test.go
+++ b/htap_test.go
@@ -1,9 +1,12 @@
 package gosnowflake
 
 import (
+	"context"
 	"database/sql/driver"
 	"fmt"
 	"reflect"
+	"strconv"
+	"strings"
 	"testing"
 	"time"
 )
@@ -409,3 +412,165 @@ func TestHybridTablesE2E(t *testing.T) {
 		})
 	})
 }
+
+func TestHTAPOptimizations(t *testing.T) {
+	if runningOnGithubAction() {
+		t.Skip("insufficient permissions")
+	}
+	for _, useHtapOptimizations := range []bool{true, false} {
+		runSnowflakeConnTest(t, func(sct *SCTest) {
+			t.Run("useHtapOptimizations="+strconv.FormatBool(useHtapOptimizations), func(t *testing.T) {
+				if useHtapOptimizations {
+					sct.mustExec("ALTER SESSION SET ENABLE_SNOW_654741_FOR_TESTING = true", nil)
+				}
+				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")
+					}
+
+					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")
+					}
+
+					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.Role)
+					}
+
+					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.Role)
+					}
+				})
+				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"])
+					}
+
+					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("date output format should be switched, expected DD-MM-YYYY, got %v", sct.sc.cfg.Params["date_output_format"])
+					}
+
+					query := sct.mustQuery("SELECT 1", nil)
+					query.Close()
+
+					if !strings.EqualFold(*sct.sc.cfg.Params["date_output_format"], "DD-MM-YYYY") {
+						t.Errorf("date output format should be switched, expected DD-MM-YYYY, got %v", sct.sc.cfg.Params["date_output_format"])
+					}
+				})
+			})
+		})
+	}
+}
+
+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")
+	}
+}