github.com/snowflakedb/gosnowflake@v1.9.0/cmd/async/async.go (about)

     1  package main
     2  
     3  import (
     4  	"context"
     5  	"database/sql"
     6  	"database/sql/driver"
     7  	"flag"
     8  	"io"
     9  	"log"
    10  	"strings"
    11  	"time"
    12  
    13  	sf "github.com/snowflakedb/gosnowflake"
    14  )
    15  
    16  func main() {
    17  	if !flag.Parsed() {
    18  		flag.Parse()
    19  	}
    20  
    21  	cfg, err := sf.GetConfigFromEnv([]*sf.ConfigParam{
    22  		{Name: "Account", EnvName: "SNOWFLAKE_TEST_ACCOUNT", FailOnMissing: true},
    23  		{Name: "User", EnvName: "SNOWFLAKE_TEST_USER", FailOnMissing: true},
    24  		{Name: "Password", EnvName: "SNOWFLAKE_TEST_PASSWORD", FailOnMissing: true},
    25  		{Name: "Host", EnvName: "SNOWFLAKE_TEST_HOST", FailOnMissing: false},
    26  		{Name: "Port", EnvName: "SNOWFLAKE_TEST_PORT", FailOnMissing: false},
    27  		{Name: "Protocol", EnvName: "SNOWFLAKE_TEST_PROTOCOL", FailOnMissing: false},
    28  	})
    29  	if err != nil {
    30  		log.Fatalf("failed to create Config, err: %v", err)
    31  	}
    32  	dsn, err := sf.DSN(cfg)
    33  	if err != nil {
    34  		log.Fatalf("failed to create DSN from Config: %v, err: %v", cfg, err)
    35  	}
    36  
    37  	db, err := sql.Open("snowflake", dsn)
    38  	if err != nil {
    39  		log.Fatalf("failed to connect. %v, err: %v", dsn, err)
    40  	}
    41  	defer db.Close()
    42  
    43  	conn, err := db.Conn(context.Background())
    44  	if err != nil {
    45  		log.Fatalf("failed to create connection. %v", err)
    46  	}
    47  	defer conn.Close()
    48  
    49  	log.Println("Lets simulate long running query by passing execution logic as a function")
    50  	driverRows := runAsyncDriverQuery(conn, "CALL SYSTEM$WAIT(10, 'SECONDS')")
    51  	log.Println("The query is running asynchronously - you can continue your workflow after starting the query")
    52  	printDriverRowsResult(driverRows)
    53  
    54  	log.Println("Lets simulate long running query using the standard sql package")
    55  	sqlRows := runAsyncSQLQuery(conn, "CALL SYSTEM$WAIT(10, 'SECONDS')")
    56  	log.Println("The query is running asynchronously - you can continue your workflow after starting the query")
    57  	printSQLRowsResult(sqlRows)
    58  }
    59  
    60  func runAsyncDriverQuery(conn *sql.Conn, query string) driver.Rows {
    61  	// Enable asynchronous mode
    62  	ctx := sf.WithAsyncMode(context.Background())
    63  
    64  	var rows driver.Rows
    65  
    66  	// Unwrap connection
    67  	err := conn.Raw(func(x interface{}) error {
    68  		var err error
    69  		// Execute asynchronous query
    70  		rows, err = withTimePrinted(func() (driver.Rows, error) {
    71  			return x.(driver.QueryerContext).QueryContext(ctx, query, nil)
    72  		})
    73  
    74  		return err
    75  	})
    76  
    77  	if err != nil {
    78  		log.Fatalf("unable to run the query. err: %v", err)
    79  	}
    80  
    81  	return rows
    82  }
    83  
    84  func runAsyncSQLQuery(conn *sql.Conn, query string) *sql.Rows {
    85  	// Enable asynchronous mode
    86  	ctx := sf.WithAsyncMode(context.Background())
    87  
    88  	// Execute asynchronous query
    89  	rows, err := withTimePrinted(func() (*sql.Rows, error) {
    90  		return conn.QueryContext(ctx, query)
    91  	})
    92  
    93  	if err != nil {
    94  		log.Fatalf("unable to run the query. err: %v", err)
    95  	}
    96  
    97  	return rows
    98  }
    99  
   100  func printDriverRowsResult(rows driver.Rows) {
   101  	log.Println(strings.Join(rows.Columns(), ", "))
   102  
   103  	dest := make([]driver.Value, 1)
   104  	for rows.Next(dest) != io.EOF {
   105  		for val := range dest {
   106  			log.Printf("%v\n", dest[val])
   107  		}
   108  	}
   109  }
   110  
   111  func printSQLRowsResult(rows *sql.Rows) {
   112  	cols, err := rows.Columns()
   113  	if err != nil {
   114  		log.Fatalf("failed to get columns. err: %v", err)
   115  	}
   116  	log.Println(strings.Join(cols, ", "))
   117  
   118  	var val string
   119  	for rows.Next() {
   120  		err := rows.Scan(&val)
   121  		if err != nil {
   122  			log.Fatalf("failed to scan rows. err: %v", err)
   123  		}
   124  		log.Printf("%v\n", val)
   125  	}
   126  }
   127  
   128  func withTimePrinted[T any](f func() (T, error)) (T, error) {
   129  	log.Printf("Start millis: %v", time.Now().UnixMilli())
   130  	defer func() {
   131  		log.Printf("End millis  : %v", time.Now().UnixMilli())
   132  	}()
   133  	return f()
   134  }