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

     1  package main
     2  
     3  import (
     4  	"context"
     5  	"database/sql"
     6  	"flag"
     7  	"fmt"
     8  	"log"
     9  	"time"
    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  		{Name: "Database", EnvName: "SNOWFLAKE_TEST_DATABASE", FailOnMissing: true},
    27  		{Name: "Schema", EnvName: "SNOWFLAKE_TEST_SCHEMA", FailOnMissing: true},
    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, error: %v", cfg, err)
    35  	}
    36  
    37  	db, err := sql.Open("snowflake", dsn)
    38  	if err != nil {
    39  		log.Fatalf("Failed to connect, dsn: %v, error: %v", dsn, err)
    40  	}
    41  
    42  	defer db.Close()
    43  
    44  	printSelectDemo(db)
    45  	printModifyingDemo(db)
    46  }
    47  
    48  func printSelectDemo(db *sql.DB) {
    49  	fmt.Println("SELECTs only")
    50  	numberOfQueries := 2
    51  	query := `
    52  		WITH table1 AS (SELECT 'table 1, row 1', 11 UNION SELECT 'table 1, row 2', 12)
    53  		SELECT * FROM table1;
    54  
    55  		WITH table2 AS (SELECT 'table 2, row 1', 21 UNION SELECT 'table 2, row 2', 22)
    56  		SELECT * FROM table2;
    57  	`
    58  
    59  	fmt.Println(query)
    60  
    61  	context := createMultistatementContext(numberOfQueries)
    62  
    63  	result, err := db.QueryContext(context, query)
    64  	if err != nil {
    65  		log.Fatalf("Error while querying snowflake: %v", err)
    66  	}
    67  
    68  	defer result.Close()
    69  
    70  	for hasNextResultSet := true; hasNextResultSet; hasNextResultSet = result.NextResultSet() {
    71  		for result.Next() {
    72  			var str string
    73  			var int int
    74  			err := result.Scan(&str, &int)
    75  			if err != nil {
    76  				log.Fatalf("Error while scanning row: %v", err)
    77  			}
    78  			fmt.Println(str, " | ", int)
    79  		}
    80  	}
    81  	if result.Err() != nil {
    82  		log.Fatalf("Error while reading results: %v", err)
    83  	}
    84  }
    85  
    86  func printModifyingDemo(db *sql.DB) {
    87  	fmt.Println("Modifications only")
    88  	numberOfQueries := 4
    89  	tableSuffix := time.Now().UnixMilli()
    90  	query := fmt.Sprintf(`
    91  		CREATE TABLE multistatement_test_%d (id integer);
    92  		INSERT INTO multistatement_test_%d VALUES (1);
    93  		INSERT INTO multistatement_test_%d VALUES (2), (3);
    94  		DROP TABLE multistatement_test_%d
    95  	`, tableSuffix, tableSuffix, tableSuffix, tableSuffix)
    96  
    97  	fmt.Println(query)
    98  
    99  	context := createMultistatementContext(numberOfQueries)
   100  
   101  	result, err := db.ExecContext(context, query)
   102  	if err != nil {
   103  		log.Fatalf("Error while querying snowflake: %v", err)
   104  	}
   105  
   106  	rowsAffected, err := result.RowsAffected()
   107  	if err != nil {
   108  		log.Fatalf("Error while reading rows affected: %v", rowsAffected)
   109  	}
   110  
   111  	fmt.Printf("Rows affected: %d, expected: %d\n", rowsAffected, 3)
   112  }
   113  
   114  func createMultistatementContext(numberOfQueries int) context.Context {
   115  	context, err := sf.WithMultiStatement(context.Background(), numberOfQueries)
   116  	if err != nil {
   117  		log.Fatalf("Error while creating multi statement context: %v", err)
   118  	}
   119  	return context
   120  }