github.com/snowflakedb/gosnowflake@v1.9.0/benchmark/largesetresult/largesetresult_test.go (about)

     1  // This code is to profile a large result set query. It is basically similar to selectmany example code but
     2  // leverages benchmark framework.
     3  package largesetresult
     4  
     5  import (
     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  	sf "github.com/snowflakedb/gosnowflake"
    20  )
    21  
    22  func TestLargeResultSet(t *testing.T) {
    23  	runLargeResultSet()
    24  }
    25  
    26  func BenchmarkLargeResultSet(*testing.B) {
    27  	runLargeResultSet()
    28  }
    29  
    30  func runLargeResultSet() {
    31  	if !flag.Parsed() {
    32  		flag.Parse()
    33  	}
    34  
    35  	// handler interrupt signal
    36  	ctx, cancel := context.WithCancel(context.Background())
    37  	c := make(chan os.Signal, 1)
    38  	defer close(c)
    39  	signal.Notify(c, os.Interrupt)
    40  	defer func() {
    41  		signal.Stop(c)
    42  	}()
    43  	go func() {
    44  		select {
    45  		case <-c:
    46  			cancel()
    47  		case <-ctx.Done():
    48  		}
    49  	}()
    50  
    51  	cfg, err := sf.GetConfigFromEnv([]*sf.ConfigParam{
    52  		{Name: "Account", EnvName: "SNOWFLAKE_TEST_ACCOUNT", FailOnMissing: true},
    53  		{Name: "User", EnvName: "SNOWFLAKE_TEST_USER", FailOnMissing: true},
    54  		{Name: "Password", EnvName: "SNOWFLAKE_TEST_PASSWORD", FailOnMissing: true},
    55  		{Name: "Host", EnvName: "SNOWFLAKE_TEST_HOST", FailOnMissing: false},
    56  		{Name: "Port", EnvName: "SNOWFLAKE_TEST_PORT", FailOnMissing: false},
    57  		{Name: "Protocol", EnvName: "SNOWFLAKE_TEST_PROTOCOL", FailOnMissing: false},
    58  	})
    59  	if err != nil {
    60  		log.Fatalf("failed to create Config, err: %v", err)
    61  	}
    62  	dsn, err := sf.DSN(cfg)
    63  	if err != nil {
    64  		log.Fatalf("failed to create DSN from Config: %v, err: %v", cfg, err)
    65  	}
    66  
    67  	db, err := sql.Open("snowflake", dsn)
    68  	defer db.Close()
    69  	if err != nil {
    70  		log.Fatalf("failed to connect. err: %v", err)
    71  	}
    72  
    73  	query := `select * from
    74  	  (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,
    75  	  (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,
    76  	  (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,
    77  	  (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,
    78  	  (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,
    79  	  (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,
    80  	  (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`
    81  	rows, err := db.QueryContext(ctx, query)
    82  	if err != nil {
    83  		log.Fatalf("failed to run a query. %v, err: %v", query, err)
    84  	}
    85  	defer rows.Close()
    86  	var v1, v2, v3, v4, v5, v6, v7 int
    87  	counter := 0
    88  	for rows.Next() {
    89  		if err = rows.Scan(&v1, &v2, &v3, &v4, &v5, &v6, &v7); err != nil {
    90  			log.Fatalf("failed to get result. err: %v", err)
    91  		}
    92  		if counter%1000000 == 0 {
    93  			debug.FreeOSMemory()
    94  		}
    95  		counter++
    96  	}
    97  }