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

     1  package main
     2  
     3  import (
     4  	"context"
     5  	"database/sql"
     6  	"database/sql/driver"
     7  	"flag"
     8  	"log"
     9  	"strings"
    10  
    11  	sf "github.com/snowflakedb/gosnowflake"
    12  )
    13  
    14  func main() {
    15  	if !flag.Parsed() {
    16  		flag.Parse()
    17  	}
    18  
    19  	cfg, err := sf.GetConfigFromEnv([]*sf.ConfigParam{
    20  		{Name: "Account", EnvName: "SNOWFLAKE_TEST_ACCOUNT", FailOnMissing: true},
    21  		{Name: "User", EnvName: "SNOWFLAKE_TEST_USER", FailOnMissing: true},
    22  		{Name: "Password", EnvName: "SNOWFLAKE_TEST_PASSWORD", FailOnMissing: true},
    23  		{Name: "Host", EnvName: "SNOWFLAKE_TEST_HOST", FailOnMissing: false},
    24  		{Name: "Port", EnvName: "SNOWFLAKE_TEST_PORT", FailOnMissing: false},
    25  		{Name: "Protocol", EnvName: "SNOWFLAKE_TEST_PROTOCOL", FailOnMissing: false},
    26  	})
    27  	if err != nil {
    28  		log.Fatalf("failed to create Config, err: %v", err)
    29  	}
    30  
    31  	dsn, err := sf.DSN(cfg)
    32  	if err != nil {
    33  		log.Fatalf("failed to create DSN from Config: %v, err: %v", cfg, err)
    34  	}
    35  
    36  	db, err := sql.Open("snowflake", dsn)
    37  	if err != nil {
    38  		log.Fatalf("failed to connect. %v, err: %v", dsn, err)
    39  	}
    40  	defer db.Close()
    41  
    42  	log.Println("Lets simulate running synchronous query and fetching the result by the query ID using the WithFetchResultByID context")
    43  	sqlRows := fetchResultByIDSync(db, "SELECT 1")
    44  	printSQLRowsResult(sqlRows)
    45  
    46  	log.Println("Lets simulate running long query asynchronously and fetching result by query ID using a channel provided in the WithQueryIDChan context")
    47  	sqlRows = fetchResultByIDAsync(db, "CALL SYSTEM$WAIT(10, 'SECONDS')")
    48  	printSQLRowsResult(sqlRows)
    49  }
    50  
    51  func fetchResultByIDSync(db *sql.DB, query string) *sql.Rows {
    52  	ctx := context.Background()
    53  	conn, err := db.Conn(ctx)
    54  	if err != nil {
    55  		log.Fatalf("failed to get Conn. err: %v", err)
    56  	}
    57  	defer conn.Close()
    58  
    59  	var rows1 driver.Rows
    60  	var queryID string
    61  
    62  	// Get the query ID using raw connection
    63  	err = conn.Raw(func(x any) error {
    64  		log.Printf("Executing query: %v\n", query)
    65  		rows1, err = x.(driver.QueryerContext).QueryContext(ctx, query, nil)
    66  		if err != nil {
    67  			return err
    68  		}
    69  
    70  		queryID = rows1.(sf.SnowflakeRows).GetQueryID()
    71  		log.Printf("Query ID retrieved from GetQueryID(): %v\n", queryID)
    72  		return nil
    73  	})
    74  	if err != nil {
    75  		log.Fatalf("unable to run the query. err: %v", err)
    76  	}
    77  
    78  	// Update the Context object to specify the query ID
    79  	ctx = sf.WithFetchResultByID(ctx, queryID)
    80  
    81  	// Execute an empty string query
    82  	rows2, err := db.QueryContext(ctx, "")
    83  	if err != nil {
    84  		log.Fatal(err)
    85  	}
    86  
    87  	return rows2
    88  }
    89  
    90  func fetchResultByIDAsync(db *sql.DB, query string) *sql.Rows {
    91  	// Make a channel to receive the query ID
    92  	queryIDChan := make(chan string, 1)
    93  
    94  	// Enable asynchronous mode
    95  	ctx := sf.WithAsyncMode(context.Background())
    96  
    97  	// Pass the channel to receive the query ID
    98  	ctx = sf.WithQueryIDChan(ctx, queryIDChan)
    99  
   100  	// Run a long running query asynchronously and without retrieving the result
   101  	log.Printf("Executing query: %v\n", query)
   102  	go db.ExecContext(ctx, query)
   103  
   104  	// Get the query ID without waiting for the query to finish
   105  	queryID := <-queryIDChan
   106  	log.Printf("Query ID retrieved from the channel: %v\n", queryID)
   107  
   108  	// Update the Context object to specify the query ID
   109  	ctx = sf.WithFetchResultByID(ctx, queryID)
   110  
   111  	// Execute an empty string query
   112  	rows, err := db.QueryContext(ctx, "")
   113  	if err != nil {
   114  		log.Fatal(err)
   115  	}
   116  
   117  	return rows
   118  }
   119  
   120  func printSQLRowsResult(rows *sql.Rows) {
   121  	log.Print("Printing the results: \n")
   122  
   123  	cols, err := rows.Columns()
   124  	if err != nil {
   125  		log.Fatalf("failed to get columns. err: %v", err)
   126  	}
   127  	log.Println(strings.Join(cols, ", "))
   128  
   129  	var val string
   130  	for rows.Next() {
   131  		err := rows.Scan(&val)
   132  		if err != nil {
   133  			log.Fatalf("failed to scan rows. err: %v", err)
   134  		}
   135  		log.Printf("%v\n", val)
   136  	}
   137  }