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