Skip to content

Commit 5f3e28b

Browse files
authored
[PECOBLR-676] passing session params in open session request instead of SET commands (#283)
<!-- We welcome contributions. All patches must include a sign-off. Please see CONTRIBUTING.md for details --> ## Description This PR **refactors session parameter handling** in the Go SQL Connector to align with other Databricks drivers (Java, Python). Previously, session parameters were set using `SET` statements after connection establishment, which was inefficient and inconsistent with other drivers. This new implementation also supports usage of QUERY_TAGS session parameter. Query Tags allow users to attach key-value pairs to SQL executions that appear in the `system.query.history` table. ### Changes Made: 1. **Refactored Session Parameter Flow** (`connector.go`): - **Before**: Session parameters were sent via `SET` statements after `OpenSession` - **After**: Session parameters are passed directly to `TOpenSessionReq.Configuration` during session creation - **Benefit**: Aligns with Java and Python drivers, and improves performance 2. **Created Example** (`examples/query_tags/main.go`): - Demonstrates the usage of Query Tags 4. **Enhanced Unit Tests** (`internal/config/config_test.go`): - Added `QUERY_TAGS` to DSN parsing test to verify parameter extraction ### Change from User Perspective: The key difference is in **error timing**: Previously, invalid session parameters would fail during connection establishment (SET command execution), but now they fail during the first query execution. This happens because the Databricks backend uses lazy validation - it accepts session parameters during connection but validates them only when the first query is executed. This matches Python driver behavior. ### Query tags usage example: ```go connector, err := dbsql.NewConnector( dbsql.WithSessionParams(map[string]string{ "timezone": "America/Sao_Paulo", "ansi_mode": "true", "QUERY_TAGS": "team:engineering", }), ) ``` ## How is this tested? - [x] Unit tests - [x] Manually ### Testing Details: **Unit Tests**: - Enhanced DSN parsing tests to verify session parameter extraction for query tags - All existing tests pass with the new implementation **Manual Testing**: - Verified session parameters including query tags work correctly with real backend using the created example
2 parents 631401e + d7ebf01 commit 5f3e28b

File tree

4 files changed

+58
-13
lines changed

4 files changed

+58
-13
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+
}

internal/config/config_test.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -129,7 +129,7 @@ func TestParseConfig(t *testing.T) {
129129
},
130130
{
131131
name: "with query params and session params",
132-
args: args{dsn: "token:supersecret@example.cloud.databricks.com:8000/sql/1.0/endpoints/12346a5b5b0e123a?timeout=100&maxRows=1000&timezone=America/Vancouver"},
132+
args: args{dsn: "token:supersecret@example.cloud.databricks.com:8000/sql/1.0/endpoints/12346a5b5b0e123a?timeout=100&maxRows=1000&timezone=America/Vancouver&QUERY_TAGS=team:testing,driver:go"},
133133
wantCfg: UserConfig{
134134
Protocol: "https",
135135
Host: "example.cloud.databricks.com",
@@ -140,7 +140,7 @@ func TestParseConfig(t *testing.T) {
140140
QueryTimeout: 100 * time.Second,
141141
MaxRows: 1000,
142142
Location: tz,
143-
SessionParams: map[string]string{"timezone": "America/Vancouver"},
143+
SessionParams: map[string]string{"timezone": "America/Vancouver", "QUERY_TAGS": "team:testing,driver:go"},
144144
RetryMax: 4,
145145
RetryWaitMin: 1 * time.Second,
146146
RetryWaitMax: 30 * time.Second,

0 commit comments

Comments
 (0)