github.com/snowflakedb/gosnowflake@v1.9.0/benchmark/jsonresultset/jsonresultset_test.go (about)

     1  // This code is to profile a large json result set query. It is basically similar to selectmany example code but
     2  // leverages benchmark framework.
     3  package jsonresultset
     4  
     5  import (
     6  	"flag"
     7  	"log"
     8  	_ "net/http/pprof"
     9  	"os"
    10  	"strconv"
    11  	"strings"
    12  	"testing"
    13  
    14  	"database/sql"
    15  
    16  	"context"
    17  	"os/signal"
    18  
    19  	"runtime/debug"
    20  
    21  	sf "github.com/snowflakedb/gosnowflake"
    22  )
    23  
    24  func TestJsonResultSet(t *testing.T) {
    25  	runJSONResultSet()
    26  }
    27  
    28  func BenchmarkJsonResultSet(*testing.B) {
    29  	runJSONResultSet()
    30  }
    31  
    32  func runJSONResultSet() {
    33  	if !flag.Parsed() {
    34  		flag.Parse()
    35  	}
    36  
    37  	// handler interrupt signal
    38  	ctx, cancel := context.WithCancel(context.Background())
    39  	c := make(chan os.Signal, 1)
    40  	defer close(c)
    41  	signal.Notify(c, os.Interrupt)
    42  	defer func() {
    43  		signal.Stop(c)
    44  	}()
    45  	go func() {
    46  		select {
    47  		case <-c:
    48  			cancel()
    49  		case <-ctx.Done():
    50  		}
    51  	}()
    52  	setCustomJSONDecoder()
    53  	setMaxChunkDownloadWorkers()
    54  	cfg, err := sf.GetConfigFromEnv([]*sf.ConfigParam{
    55  		{Name: "Account", EnvName: "SNOWFLAKE_TEST_ACCOUNT", FailOnMissing: true},
    56  		{Name: "User", EnvName: "SNOWFLAKE_TEST_USER", FailOnMissing: true},
    57  		{Name: "Password", EnvName: "SNOWFLAKE_TEST_PASSWORD", FailOnMissing: true},
    58  		{Name: "Role", EnvName: "SNOWFLAKE_TEST_ROLE", FailOnMissing: false},
    59  		{Name: "Host", EnvName: "SNOWFLAKE_TEST_HOST", FailOnMissing: false},
    60  		{Name: "Port", EnvName: "SNOWFLAKE_TEST_PORT", FailOnMissing: false},
    61  		{Name: "Protocol", EnvName: "SNOWFLAKE_TEST_PROTOCOL", FailOnMissing: false},
    62  	})
    63  	if err != nil {
    64  		log.Fatalf("failed to create Config, err: %v", err)
    65  	}
    66  	dsn, err := sf.DSN(cfg)
    67  	if err != nil {
    68  		log.Fatalf("failed to create DSN from Config: %v, err: %v", cfg, err)
    69  	}
    70  
    71  	db, err := sql.Open("snowflake", dsn)
    72  	defer db.Close()
    73  	if err != nil {
    74  		log.Fatalf("failed to connect. %v, err: %v", dsn, err)
    75  	}
    76  
    77  	query := `SELECT V FROM SNOWFLAKE_SAMPLE_DATA.WEATHER.HOURLY_14_TOTAL LIMIT 100000`
    78  	rows, err := db.QueryContext(ctx, query)
    79  	if err != nil {
    80  		log.Fatalf("failed to run a query. %v, err: %v", query, err)
    81  	}
    82  	defer rows.Close()
    83  	var v1 string
    84  	counter := 0
    85  	for rows.Next() {
    86  		err := rows.Scan(&v1)
    87  		if err != nil {
    88  			log.Fatalf("failed to get result. err: %v", err)
    89  		}
    90  		if counter%1000000 == 0 {
    91  			debug.FreeOSMemory()
    92  		}
    93  		counter++
    94  	}
    95  }
    96  
    97  func setMaxChunkDownloadWorkers() {
    98  	maxChunkDownloadWorkersStr, err := sf.GetFromEnv("SNOWFLAKE_TEST_MAX_CHUNK_DOWNLOAD_WORKERS", false)
    99  	if err != nil {
   100  		log.Fatal(err)
   101  	}
   102  	if maxChunkDownloadWorkersStr != "" {
   103  		maxChunkDownloadWorkers, err := strconv.Atoi(maxChunkDownloadWorkersStr)
   104  		if err != nil {
   105  			log.Fatalf("invalid value for SNOWFLAKE_TEST_MAX_CHUNK_DOWNLOAD_WORKERS: %v", maxChunkDownloadWorkers)
   106  		}
   107  		sf.MaxChunkDownloadWorkers = maxChunkDownloadWorkers
   108  	}
   109  }
   110  
   111  func setCustomJSONDecoder() {
   112  	customJSONDecoderEnabledStr, err := sf.GetFromEnv("SNOWFLAKE_TEST_CUSTOME_JSON_DECODER_ENABLED", true)
   113  	if err != nil {
   114  		log.Fatal(err)
   115  	}
   116  	sf.CustomJSONDecoderEnabled = strings.EqualFold("true", customJSONDecoderEnabledStr)
   117  }