Skip to content

Commit

Permalink
SNOW-898296: Enable HTAP optimisations in tests
Browse files Browse the repository at this point in the history
  • Loading branch information
sfc-gh-pfus committed Sep 7, 2023
1 parent eecf8bf commit 3135b3c
Show file tree
Hide file tree
Showing 4 changed files with 180 additions and 4 deletions.
16 changes: 12 additions & 4 deletions connection.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand Down
2 changes: 2 additions & 0 deletions driver_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}

Expand Down
163 changes: 163 additions & 0 deletions htap_test.go
Original file line number Diff line number Diff line change
@@ -1,10 +1,13 @@
package gosnowflake

import (
"context"
"encoding/json"
"fmt"
"reflect"
"strings"
"testing"
"time"
)

func TestMarshallAndDecodeOpaqueContext(t *testing.T) {
Expand Down Expand Up @@ -416,3 +419,163 @@ func htapTestSnowflakeConn() *snowflakeConn {
},
}
}

func TestHTAPOptimizations(t *testing.T) {
runSnowflakeConnTest(t, func(sc *snowflakeConn) {
runId := time.Now().UnixMilli()

Check failure on line 425 in htap_test.go

View workflow job for this annotation

GitHub Actions / AZURE Go 1.19 on Ubuntu

var runId should be runID

Check failure on line 425 in htap_test.go

View workflow job for this annotation

GitHub Actions / AWS Go 1.19 on Ubuntu

var runId should be runID

Check failure on line 425 in htap_test.go

View workflow job for this annotation

GitHub Actions / AWS Go 1.20 on Ubuntu

var runId should be runID

Check failure on line 425 in htap_test.go

View workflow job for this annotation

GitHub Actions / AZURE Go 1.20 on Ubuntu

var runId should be runID

Check failure on line 425 in htap_test.go

View workflow job for this annotation

GitHub Actions / GCP Go 1.19 on Ubuntu

var runId should be runID

Check failure on line 425 in htap_test.go

View workflow job for this annotation

GitHub Actions / GCP Go 1.20 on Ubuntu

var runId should be runID

Check failure on line 425 in htap_test.go

View workflow job for this annotation

GitHub Actions / AWS Go 1.20 on Mac

var runId should be runID

Check failure on line 425 in htap_test.go

View workflow job for this annotation

GitHub Actions / AZURE Go 1.20 on Mac

var runId should be runID

Check failure on line 425 in htap_test.go

View workflow job for this annotation

GitHub Actions / GCP Go 1.19 on Mac

var runId should be runID

Check failure on line 425 in htap_test.go

View workflow job for this annotation

GitHub Actions / GCP Go 1.20 on Mac

var runId should be runID

Check failure on line 425 in htap_test.go

View workflow job for this annotation

GitHub Actions / AZURE Go 1.19 on Mac

var runId should be runID

Check failure on line 425 in htap_test.go

View workflow job for this annotation

GitHub Actions / AWS Go 1.19 on Mac

var runId should be runID

Check failure on line 425 in htap_test.go

View workflow job for this annotation

GitHub Actions / AWS Go 1.19 on Windows

var runId should be runID

Check failure on line 425 in htap_test.go

View workflow job for this annotation

GitHub Actions / AWS Go 1.20 on Windows

var runId should be runID

Check failure on line 425 in htap_test.go

View workflow job for this annotation

GitHub Actions / AZURE Go 1.19 on Windows

var runId should be runID

Check failure on line 425 in htap_test.go

View workflow job for this annotation

GitHub Actions / GCP Go 1.20 on Windows

var runId should be runID

Check failure on line 425 in htap_test.go

View workflow job for this annotation

GitHub Actions / AZURE Go 1.20 on Windows

var runId should be runID

Check failure on line 425 in htap_test.go

View workflow job for this annotation

GitHub Actions / GCP Go 1.19 on Windows

var runId should be runID
t.Run("Schema", func(t *testing.T) {
newSchema := fmt.Sprintf("test_schema_%v", runId)
if strings.EqualFold(sc.cfg.Schema, newSchema) {
t.Errorf("schema should not be switched")
}

// TODO replace with mustExec
sc.Exec(fmt.Sprintf("CREATE SCHEMA %v", newSchema), nil)
defer sc.Exec(fmt.Sprintf("DROP SCHEMA %v", newSchema), nil)

if !strings.EqualFold(sc.cfg.Schema, newSchema) {
t.Errorf("schema should be switched, expected %v, got %v", newSchema, sc.cfg.Schema)
}

query, _ := sc.Query("SELECT 1", nil)
query.Close()

if !strings.EqualFold(sc.cfg.Schema, newSchema) {
t.Errorf("schema should be switched, expected %v, got %v", newSchema, sc.cfg.Schema)
}
})
t.Run("Database", func(t *testing.T) {
newDatabase := fmt.Sprintf("test_database_%v", runId)
if strings.EqualFold(sc.cfg.Database, newDatabase) {
t.Errorf("database should not be switched")
}

// TODO replace with mustExec
sc.Exec(fmt.Sprintf("CREATE DATABASE %v", newDatabase), nil)
defer sc.Exec(fmt.Sprintf("DROP DATABASE %v", newDatabase), nil)

if !strings.EqualFold(sc.cfg.Database, newDatabase) {
t.Errorf("database should be switched, expected %v, got %v", newDatabase, sc.cfg.Database)
}

query, _ := sc.Query("SELECT 1", nil)
query.Close()

if !strings.EqualFold(sc.cfg.Database, newDatabase) {
t.Errorf("database should be switched, expected %v, got %v", newDatabase, sc.cfg.Database)
}
})
t.Run("Warehouse", func(t *testing.T) {
newWarehouse := fmt.Sprintf("test_warehouse_%v", runId)
if strings.EqualFold(sc.cfg.Warehouse, newWarehouse) {
t.Errorf("warehouse should not be switched")
}

// TODO replace with mustExec
sc.Exec(fmt.Sprintf("CREATE WAREHOUSE %v", newWarehouse), nil)
defer sc.Exec(fmt.Sprintf("DROP WAREHOUSE %v", newWarehouse), nil)

if !strings.EqualFold(sc.cfg.Warehouse, newWarehouse) {
t.Errorf("warehouse should be switched, expected %v, got %v", newWarehouse, sc.cfg.Warehouse)
}

query, _ := sc.Query("SELECT 1", nil)
query.Close()

if !strings.EqualFold(sc.cfg.Warehouse, newWarehouse) {
t.Errorf("warehouse should be switched, expected %v, got %v", newWarehouse, sc.cfg.Warehouse)
}
})
t.Run("Role", func(t *testing.T) {
if strings.EqualFold(sc.cfg.Role, "PUBLIC") {
t.Errorf("role should not be public for this test")
}

// TODO replace with mustExec
sc.Exec("USE ROLE public", nil)

if !strings.EqualFold(sc.cfg.Role, "PUBLIC") {
t.Errorf("role should be switched, expected public, got %v", sc.cfg.Warehouse)
}

query, _ := sc.Query("SELECT 1", nil)
query.Close()

if !strings.EqualFold(sc.cfg.Role, "PUBLIC") {
t.Errorf("role should be switched, expected public, got %v", sc.cfg.Warehouse)
}
})
t.Run("Session param - DATE_OUTPUT_FORMAT", func(t *testing.T) {
if !strings.EqualFold(*sc.cfg.Params["date_output_format"], "YYYY-MM-DD") {
t.Errorf("should use default date_output_format, but got: %v", *sc.cfg.Params["date_output_format"])
}

// TODO replace with mustExec
sc.Exec("ALTER SESSION SET DATE_OUTPUT_FORMAT = 'DD-MM-YYYY'", nil)
defer sc.Exec("ALTER SESSION SET DATE_OUTPUT_FORMAT = 'YYYY-MM-DD'", nil)

if !strings.EqualFold(*sc.cfg.Params["date_output_format"], "DD-MM-YYYY") {
t.Errorf("role should be switched, expected public, got %v", sc.cfg.Warehouse)
}

query, _ := sc.Query("SELECT 1", nil)
query.Close()

if !strings.EqualFold(*sc.cfg.Params["date_output_format"], "DD-MM-YYYY") {
t.Errorf("role should be switched, expected public, got %v", sc.cfg.Warehouse)
}
})
})
}

func TestConnIsCleanAfterClose(t *testing.T) {
t.Skip("Fails, because connection is returned to a pool dirty")
ctx := context.Background()
runId := time.Now().UnixMilli()

Check failure on line 534 in htap_test.go

View workflow job for this annotation

GitHub Actions / AZURE Go 1.19 on Ubuntu

var runId should be runID

Check failure on line 534 in htap_test.go

View workflow job for this annotation

GitHub Actions / AWS Go 1.19 on Ubuntu

var runId should be runID

Check failure on line 534 in htap_test.go

View workflow job for this annotation

GitHub Actions / AWS Go 1.20 on Ubuntu

var runId should be runID

Check failure on line 534 in htap_test.go

View workflow job for this annotation

GitHub Actions / AZURE Go 1.20 on Ubuntu

var runId should be runID

Check failure on line 534 in htap_test.go

View workflow job for this annotation

GitHub Actions / GCP Go 1.19 on Ubuntu

var runId should be runID

Check failure on line 534 in htap_test.go

View workflow job for this annotation

GitHub Actions / GCP Go 1.20 on Ubuntu

var runId should be runID

Check failure on line 534 in htap_test.go

View workflow job for this annotation

GitHub Actions / AWS Go 1.20 on Mac

var runId should be runID

Check failure on line 534 in htap_test.go

View workflow job for this annotation

GitHub Actions / AZURE Go 1.20 on Mac

var runId should be runID

Check failure on line 534 in htap_test.go

View workflow job for this annotation

GitHub Actions / GCP Go 1.19 on Mac

var runId should be runID

Check failure on line 534 in htap_test.go

View workflow job for this annotation

GitHub Actions / GCP Go 1.20 on Mac

var runId should be runID

Check failure on line 534 in htap_test.go

View workflow job for this annotation

GitHub Actions / AZURE Go 1.19 on Mac

var runId should be runID

Check failure on line 534 in htap_test.go

View workflow job for this annotation

GitHub Actions / AWS Go 1.19 on Mac

var runId should be runID

Check failure on line 534 in htap_test.go

View workflow job for this annotation

GitHub Actions / AWS Go 1.19 on Windows

var runId should be runID

Check failure on line 534 in htap_test.go

View workflow job for this annotation

GitHub Actions / AWS Go 1.20 on Windows

var runId should be runID

Check failure on line 534 in htap_test.go

View workflow job for this annotation

GitHub Actions / AZURE Go 1.19 on Windows

var runId should be runID

Check failure on line 534 in htap_test.go

View workflow job for this annotation

GitHub Actions / GCP Go 1.20 on Windows

var runId should be runID

Check failure on line 534 in htap_test.go

View workflow job for this annotation

GitHub Actions / AZURE Go 1.20 on Windows

var runId should be runID

Check failure on line 534 in htap_test.go

View workflow job for this annotation

GitHub Actions / GCP Go 1.19 on Windows

var runId should be runID

db := openDB(t)
defer db.Close()
db.SetMaxOpenConns(1)

conn, err := db.Conn(ctx)
if err != nil {
t.Fatal(err)
}
defer conn.Close()
conn.ExecContext(ctx, forceJSON)

var dbName string
rows1, err := conn.QueryContext(ctx, "SELECT CURRENT_DATABASE()")
if err != nil {
t.Fatal(err)
}
rows1.Next()
rows1.Scan(&dbName)

newDbName := fmt.Sprintf("test_database_%v", runId)
_, err = conn.ExecContext(ctx, fmt.Sprintf("CREATE DATABASE %v", newDbName))
if err != nil {
t.Fatal(err)
}

rows1.Close()
conn.Close()

conn2, err := db.Conn(ctx)
if err != nil {
t.Fatal(err)
}

var dbName2 string
rows2, err := conn2.QueryContext(ctx, "SELECT CURRENT_DATABASE()")
if err != nil {
t.Fatal(err)
}
defer rows2.Close()
rows2.Next()
rows2.Scan(&dbName2)

if !strings.EqualFold(dbName, dbName2) {
t.Errorf("fresh connection from pool should have original database")
}
}
3 changes: 3 additions & 0 deletions statement_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}

Expand Down

0 comments on commit 3135b3c

Please sign in to comment.