Skip to content

Commit ceaf079

Browse files
committed
passing session params in open session request
Signed-off-by: Sreekanth Vadigi <[email protected]>
1 parent 631401e commit ceaf079

File tree

3 files changed

+56
-11
lines changed

3 files changed

+56
-11
lines changed

.gitignore

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
.vscode
2+
.idea
23

34
# Binaries for programs and plugins
45
*.exe

connector.go

Lines changed: 2 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) {
4444
protocolVersion := int64(c.cfg.ThriftProtocolVersion)
4545
session, err := tclient.OpenSession(ctx, &cli_service.TOpenSessionReq{
4646
ClientProtocolI64: &protocolVersion,
47-
Configuration: make(map[string]string),
47+
Configuration: c.cfg.SessionParams,
4848
InitialNamespace: &cli_service.TNamespace{
4949
CatalogName: catalogName,
5050
SchemaName: schemaName,
@@ -65,14 +65,6 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) {
6565

6666
log.Info().Msgf("connect: host=%s port=%d httpPath=%s serverProtocolVersion=0x%X", c.cfg.Host, c.cfg.Port, c.cfg.HTTPPath, session.ServerProtocolVersion)
6767

68-
for k, v := range c.cfg.SessionParams {
69-
setStmt := fmt.Sprintf("SET `%s` = `%s`;", k, v)
70-
_, err := conn.ExecContext(ctx, setStmt, []driver.NamedValue{})
71-
if err != nil {
72-
return nil, dbsqlerrint.NewExecutionError(ctx, fmt.Sprintf("error setting session param: %s", setStmt), err, nil)
73-
}
74-
log.Info().Msgf("set session parameter: param=%s value=%s", k, v)
75-
}
7668
return conn, nil
7769
}
7870

@@ -215,8 +207,7 @@ func WithUserAgentEntry(entry string) ConnOption {
215207
}
216208
}
217209

218-
// Sessions params will be set upon opening the session by calling SET function.
219-
// If using connection pool, session params can avoid successive calls of "SET ..."
210+
// Session parameters are passed directly in TOpenSessionReq.Configuration during session creation.
220211
func WithSessionParams(params map[string]string) ConnOption {
221212
return func(c *config.Config) {
222213
for k, v := range params {

examples/query_tags/main.go

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
package main
2+
3+
import (
4+
"context"
5+
"database/sql"
6+
"fmt"
7+
"log"
8+
"os"
9+
"strconv"
10+
11+
dbsql "github.com/databricks/databricks-sql-go"
12+
"github.com/joho/godotenv"
13+
)
14+
15+
func main() {
16+
17+
_ = godotenv.Load()
18+
19+
port, err := strconv.Atoi(os.Getenv("DATABRICKS_PORT"))
20+
if err != nil {
21+
log.Fatal(err.Error())
22+
}
23+
24+
connector, err := dbsql.NewConnector(
25+
dbsql.WithServerHostname(os.Getenv("DATABRICKS_HOST")),
26+
dbsql.WithPort(port),
27+
dbsql.WithHTTPPath(os.Getenv("DATABRICKS_HTTPPATH")),
28+
dbsql.WithAccessToken(os.Getenv("DATABRICKS_ACCESSTOKEN")),
29+
dbsql.WithSessionParams(map[string]string{
30+
"QUERY_TAGS": "team:engineering,test:query-tags,driver:go",
31+
"ansi_mode": "false",
32+
}),
33+
)
34+
if err != nil {
35+
log.Fatal(err)
36+
}
37+
38+
db := sql.OpenDB(connector)
39+
defer db.Close()
40+
41+
ctx := context.Background()
42+
var result int
43+
err = db.QueryRowContext(ctx, "SELECT 1").Scan(&result)
44+
if err != nil {
45+
if err == sql.ErrNoRows {
46+
fmt.Println("not found")
47+
return
48+
} else {
49+
fmt.Printf("err: %v\n", err)
50+
}
51+
}
52+
fmt.Println(result)
53+
}

0 commit comments

Comments
 (0)