SNOW-949224 Introduce conn in async example

cmd/async/async.go

@@ -5,10 +5,10 @@ import (
 	"database/sql"
 	"database/sql/driver"
 	"flag"
-	"fmt"
 	"io"
 	"log"
 	"strings"
+	"time"
 
 	sf "github.com/snowflakedb/gosnowflake"
 )
@@ -40,30 +40,36 @@ func main() {
 	}
 	defer db.Close()
 
-	fmt.Println("Lets simulate long running query by passing execution logic as a function")
-	driverRows := runAsyncDriverQuery(db, "CALL SYSTEM$WAIT(10, 'SECONDS')")
-	fmt.Println("The query is running asynchronously - you can continue your workflow after starting the query")
+	conn, err := db.Conn(context.Background())
+	if err != nil {
+		log.Fatalf("failed to create connection. %v", err)
+	}
+	defer conn.Close()
+
+	log.Println("Lets simulate long running query by passing execution logic as a function")
+	driverRows := runAsyncDriverQuery(conn, "CALL SYSTEM$WAIT(10, 'SECONDS')")
+	log.Println("The query is running asynchronously - you can continue your workflow after starting the query")
 	printDriverRowsResult(driverRows)
 
-	fmt.Println("Lets simulate long running query using the standard sql package")
-	sqlRows := runAsyncSQLQuery(db, "CALL SYSTEM$WAIT(10, 'SECONDS')")
-	fmt.Println("The query is running asynchronously - you can continue your workflow after starting the query")
+	log.Println("Lets simulate long running query using the standard sql package")
+	sqlRows := runAsyncSQLQuery(conn, "CALL SYSTEM$WAIT(10, 'SECONDS')")
+	log.Println("The query is running asynchronously - you can continue your workflow after starting the query")
 	printSQLRowsResult(sqlRows)
 }
 
-func runAsyncDriverQuery(db *sql.DB, query string) driver.Rows {
+func runAsyncDriverQuery(conn *sql.Conn, query string) driver.Rows {
 	// Enable asynchronous mode
 	ctx := sf.WithAsyncMode(context.Background())
 
-	// Establish a connection
-	conn, _ := db.Conn(ctx)
 	var rows driver.Rows
 
 	// Unwrap connection
 	err := conn.Raw(func(x interface{}) error {
 		var err error
 		// Execute asynchronous query
-		rows, err = x.(driver.QueryerContext).QueryContext(ctx, query, nil)
+		rows, err = withTimePrinted(func() (driver.Rows, error) {
+			return x.(driver.QueryerContext).QueryContext(ctx, query, nil)
+		})
 
 		return err
 	})
@@ -75,12 +81,14 @@ func runAsyncDriverQuery(db *sql.DB, query string) driver.Rows {
 	return rows
 }
 
-func runAsyncSQLQuery(db *sql.DB, query string) *sql.Rows {
+func runAsyncSQLQuery(conn *sql.Conn, query string) *sql.Rows {
 	// Enable asynchronous mode
 	ctx := sf.WithAsyncMode(context.Background())
 
 	// Execute asynchronous query
-	rows, err := db.QueryContext(ctx, query)
+	rows, err := withTimePrinted(func() (*sql.Rows, error) {
+		return conn.QueryContext(ctx, query)
+	})
 
 	if err != nil {
 		log.Fatalf("unable to run the query. err: %v", err)
@@ -90,12 +98,12 @@ func runAsyncSQLQuery(db *sql.DB, query string) *sql.Rows {
 }
 
 func printDriverRowsResult(rows driver.Rows) {
-	fmt.Println(strings.Join(rows.Columns(), ", "))
+	log.Println(strings.Join(rows.Columns(), ", "))
 
 	dest := make([]driver.Value, 1)
 	for rows.Next(dest) != io.EOF {
 		for val := range dest {
-			fmt.Printf("%v\n", dest[val])
+			log.Printf("%v\n", dest[val])
 		}
 	}
 }
@@ -105,14 +113,22 @@ func printSQLRowsResult(rows *sql.Rows) {
 	if err != nil {
 		log.Fatalf("failed to get columns. err: %v", err)
 	}
-	fmt.Println(strings.Join(cols, ", "))
+	log.Println(strings.Join(cols, ", "))
 
 	var val string
 	for rows.Next() {
 		err := rows.Scan(&val)
 		if err != nil {
 			log.Fatalf("failed to scan rows. err: %v", err)
 		}
-		fmt.Printf("%v\n", val)
+		log.Printf("%v\n", val)
 	}
 }
+
+func withTimePrinted[T any](f func() (T, error)) (T, error) {
+	log.Printf("Start millis: %v", time.Now().UnixMilli())
+	defer func() {
+		log.Printf("End millis  : %v", time.Now().UnixMilli())
+	}()
+	return f()
+}