1// This code is to profile a large result set query. It is basically similar to selectmany example code but 2// leverages benchmark framework. 3package largesetresult 4 5import ( 6 "flag" 7 "log" 8 _ "net/http/pprof" 9 "os" 10 "testing" 11 12 "database/sql" 13 14 "context" 15 "os/signal" 16 17 "runtime/debug" 18 19 "strconv" 20 21 sf "github.com/snowflakedb/gosnowflake" 22) 23 24func TestLargeResultSet(t *testing.T) { 25 runLargeResultSet() 26} 27 28func BenchmarkLargeResultSet(*testing.B) { 29 runLargeResultSet() 30} 31 32// getDSN constructs a DSN based on the test connection parameters 33func getDSN() (dsn string, cfg *sf.Config, err error) { 34 env := func(k string, failOnMissing bool) string { 35 if value := os.Getenv(k); value != "" { 36 return value 37 } 38 if failOnMissing { 39 log.Fatalf("%v environment variable is not set.", k) 40 } 41 return "" 42 } 43 44 account := env("SNOWFLAKE_TEST_ACCOUNT", true) 45 user := env("SNOWFLAKE_TEST_USER", true) 46 password := env("SNOWFLAKE_TEST_PASSWORD", true) 47 host := env("SNOWFLAKE_TEST_HOST", false) 48 port := env("SNOWFLAKE_TEST_PORT", false) 49 protocol := env("SNOWFLAKE_TEST_PROTOCOL", false) 50 51 portStr, _ := strconv.Atoi(port) 52 cfg = &sf.Config{ 53 Account: account, 54 User: user, 55 Password: password, 56 Host: host, 57 Port: portStr, 58 Protocol: protocol, 59 } 60 61 dsn, err = sf.DSN(cfg) 62 return dsn, cfg, err 63} 64 65func runLargeResultSet() { 66 if !flag.Parsed() { 67 flag.Parse() 68 } 69 70 // handler interrupt signal 71 ctx, cancel := context.WithCancel(context.Background()) 72 c := make(chan os.Signal, 1) 73 defer close(c) 74 signal.Notify(c, os.Interrupt) 75 defer func() { 76 signal.Stop(c) 77 }() 78 go func() { 79 select { 80 case <-c: 81 cancel() 82 case <-ctx.Done(): 83 } 84 }() 85 86 dsn, cfg, err := getDSN() 87 if err != nil { 88 log.Fatalf("failed to create DSN from Config: %v, err: %v", cfg, err) 89 } 90 91 db, err := sql.Open("snowflake", dsn) 92 defer db.Close() 93 if err != nil { 94 log.Fatalf("failed to connect. %v, err: %v", dsn, err) 95 } 96 97 query := `select * from 98 (select 0 a union select 1 union select 2 union select 3 union select 4 union select 5 union select 6 union select 7 union select 8 union select 9) A, 99 (select 0 b union select 1 union select 2 union select 3 union select 4 union select 5 union select 6 union select 7 union select 8 union select 9) B, 100 (select 0 c union select 1 union select 2 union select 3 union select 4 union select 5 union select 6 union select 7 union select 8 union select 9) C, 101 (select 0 d union select 1 union select 2 union select 3 union select 4 union select 5 union select 6 union select 7 union select 8 union select 9) E, 102 (select 0 e union select 1 union select 2 union select 3 union select 4 union select 5 union select 6 union select 7 union select 8 union select 9) F, 103 (select 0 f union select 1 union select 2 union select 3 union select 4 union select 5 union select 6 union select 7 union select 8 union select 9) G, 104 (select 0 f union select 1 union select 2 union select 3 union select 4 union select 5 union select 6 union select 7 union select 8 union select 9) H` 105 rows, err := db.QueryContext(ctx, query) 106 if err != nil { 107 log.Fatalf("failed to run a query. %v, err: %v", query, err) 108 } 109 defer rows.Close() 110 var v1 int 111 var v2 int 112 var v3 int 113 var v4 int 114 var v5 int 115 var v6 int 116 var v7 int 117 counter := 0 118 for rows.Next() { 119 err := rows.Scan(&v1, &v2, &v3, &v4, &v5, &v6, &v7) 120 if err != nil { 121 log.Fatalf("failed to get result. err: %v", err) 122 } 123 if counter%1000000 == 0 { 124 debug.FreeOSMemory() 125 } 126 counter++ 127 } 128} 129