github.com/igungor/gosnowflake@v0.0.0-20170920230020-f7fa719c23e6/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  	"fmt"
     8  	"log"
     9  	_ "net/http/pprof"
    10  	"os"
    11  	"testing"
    12  
    13  	"database/sql"
    14  
    15  	"context"
    16  	_ "github.com/snowflakedb/gosnowflake"
    17  	"os/signal"
    18  )
    19  
    20  func TestLargeResultSet(t *testing.T) {
    21  	runLargeResultSet()
    22  }
    23  
    24  func BenchmarkLargeResultSet(*testing.B) {
    25  	runLargeResultSet()
    26  }
    27  
    28  func runLargeResultSet() {
    29  	if !flag.Parsed() {
    30  		// enable glog for Go Snowflake Driver
    31  		flag.Parse()
    32  	}
    33  
    34  	// handler interrupt signal
    35  	ctx, cancel := context.WithCancel(context.Background())
    36  	c := make(chan os.Signal, 1)
    37  	signal.Notify(c, os.Interrupt)
    38  	defer func() {
    39  		signal.Stop(c)
    40  	}()
    41  	go func() {
    42  		<-c
    43  		log.Println("Caught signal, canceling...")
    44  		cancel()
    45  	}()
    46  
    47  	// get environment variables
    48  	env := func(k string) string {
    49  		if value := os.Getenv(k); value != "" {
    50  			return value
    51  		}
    52  		log.Fatalf("%v environment variable is not set.", k)
    53  		return ""
    54  	}
    55  
    56  	account := env("SNOWFLAKE_TEST_ACCOUNT")
    57  	user := env("SNOWFLAKE_TEST_USER")
    58  	password := env("SNOWFLAKE_TEST_PASSWORD")
    59  
    60  	dsn := fmt.Sprintf("%v:%v@%v", user, password, account)
    61  	db, err := sql.Open("snowflake", dsn)
    62  	defer db.Close()
    63  	if err != nil {
    64  		log.Fatalf("failed to connect. %v, err: %v", dsn, err)
    65  	}
    66  
    67  	query := "SELECT seq8(), randstr(100, random()) FROM table(generator(rowcount=>100000))"
    68  	rows, err := db.QueryContext(ctx, query)
    69  	if err != nil {
    70  		log.Fatalf("failed to run a query. %v, err: %v", query, err)
    71  	}
    72  	defer rows.Close()
    73  	var v int
    74  	var s string
    75  	for rows.Next() {
    76  		err := rows.Scan(&v, &s)
    77  		if err != nil {
    78  			log.Fatalf("failed to get result. err: %v", err)
    79  		}
    80  	}
    81  }