passing session params in open session request

Signed-off-by: Sreekanth Vadigi <[email protected]>

Author: sreekanth-db

.gitignore

@@ -1,4 +1,5 @@
 .vscode
+.idea
 
 # Binaries for programs and plugins
 *.exe

connector.go

@@ -44,7 +44,7 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) {
 	protocolVersion := int64(c.cfg.ThriftProtocolVersion)
 	session, err := tclient.OpenSession(ctx, &cli_service.TOpenSessionReq{
 		ClientProtocolI64: &protocolVersion,
-		Configuration:     make(map[string]string),
+		Configuration:     c.cfg.SessionParams,
 		InitialNamespace: &cli_service.TNamespace{
 			CatalogName: catalogName,
 			SchemaName:  schemaName,
@@ -65,14 +65,6 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) {
 
 	log.Info().Msgf("connect: host=%s port=%d httpPath=%s serverProtocolVersion=0x%X", c.cfg.Host, c.cfg.Port, c.cfg.HTTPPath, session.ServerProtocolVersion)
 
-	for k, v := range c.cfg.SessionParams {
-		setStmt := fmt.Sprintf("SET `%s` = `%s`;", k, v)
-		_, err := conn.ExecContext(ctx, setStmt, []driver.NamedValue{})
-		if err != nil {
-			return nil, dbsqlerrint.NewExecutionError(ctx, fmt.Sprintf("error setting session param: %s", setStmt), err, nil)
-		}
-		log.Info().Msgf("set session parameter: param=%s value=%s", k, v)
-	}
 	return conn, nil
 }
 
@@ -215,8 +207,7 @@ func WithUserAgentEntry(entry string) ConnOption {
 	}
 }
 
-// Sessions params will be set upon opening the session by calling SET function.
-// If using connection pool, session params can avoid successive calls of "SET ..."
+// Session parameters are passed directly in TOpenSessionReq.Configuration during session creation.
 func WithSessionParams(params map[string]string) ConnOption {
 	return func(c *config.Config) {
 		for k, v := range params {

examples/query_tags/main.go

@@ -0,0 +1,53 @@
+package main
+
+import (
+	"context"
+	"database/sql"
+	"fmt"
+	"log"
+	"os"
+	"strconv"
+
+	dbsql "github.com/databricks/databricks-sql-go"
+	"github.com/joho/godotenv"
+)
+
+func main() {
+
+	_ = godotenv.Load()
+
+	port, err := strconv.Atoi(os.Getenv("DATABRICKS_PORT"))
+	if err != nil {
+		log.Fatal(err.Error())
+	}
+
+	connector, err := dbsql.NewConnector(
+		dbsql.WithServerHostname(os.Getenv("DATABRICKS_HOST")),
+		dbsql.WithPort(port),
+		dbsql.WithHTTPPath(os.Getenv("DATABRICKS_HTTPPATH")),
+		dbsql.WithAccessToken(os.Getenv("DATABRICKS_ACCESSTOKEN")),
+		dbsql.WithSessionParams(map[string]string{
+			"QUERY_TAGS": "team:engineering,test:query-tags,driver:go",
+			"ansi_mode":  "false",
+		}),
+	)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	db := sql.OpenDB(connector)
+	defer db.Close()
+
+	ctx := context.Background()
+	var result int
+	err = db.QueryRowContext(ctx, "SELECT 1").Scan(&result)
+	if err != nil {
+		if err == sql.ErrNoRows {
+			fmt.Println("not found")
+			return
+		} else {
+			fmt.Printf("err: %v\n", err)
+		}
+	}
+	fmt.Println(result)
+}