Skip to content

Commit 8d0e384

Browse files
feat(go): add databricks:// uri support, tests, and docs (#102)
## What's Changed Support databricks:// URI scheme Basic connection test with URI Document URI formats
1 parent bc9c54a commit 8d0e384

File tree

3 files changed

+162
-14
lines changed

3 files changed

+162
-14
lines changed

go/database.go

Lines changed: 50 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,7 @@ type databaseImpl struct {
5555
needsRefresh bool // Whether we need to re-initialize
5656

5757
// Connection parameters
58+
uri string
5859
serverHostname string
5960
httpPath string
6061
accessToken string
@@ -160,20 +161,29 @@ func (d *databaseImpl) resolveConnectionOptions() ([]dbsql.ConnOption, error) {
160161
}
161162

162163
func (d *databaseImpl) initializeConnectionPool(ctx context.Context) (*sql.DB, error) {
163-
opts, err := d.resolveConnectionOptions()
164+
var db *sql.DB
164165

165-
if err != nil {
166-
return nil, err
167-
}
166+
// Use URI if provided
167+
if d.uri != "" {
168+
var err error
169+
db, err = sql.Open("databricks", d.uri)
170+
if err != nil {
171+
return nil, err
172+
}
173+
} else {
174+
opts, err := d.resolveConnectionOptions()
175+
if err != nil {
176+
return nil, err
177+
}
168178

169-
connector, err := dbsql.NewConnector(opts...)
179+
connector, err := dbsql.NewConnector(opts...)
180+
if err != nil {
181+
return nil, err
182+
}
170183

171-
if err != nil {
172-
return nil, err
184+
db = sql.OpenDB(connector)
173185
}
174186

175-
db := sql.OpenDB(connector)
176-
177187
// Test the connection
178188
if err := db.PingContext(ctx); err != nil {
179189
err = errors.Join(db.Close())
@@ -238,6 +248,8 @@ func (d *databaseImpl) Close() error {
238248

239249
func (d *databaseImpl) GetOption(key string) (string, error) {
240250
switch key {
251+
case adbc.OptionKeyURI:
252+
return d.uri, nil
241253
case OptionServerHostname:
242254
return d.serverHostname, nil
243255
case OptionHTTPPath:
@@ -288,6 +300,25 @@ func (d *databaseImpl) GetOption(key string) (string, error) {
288300
func (d *databaseImpl) SetOptions(options map[string]string) error {
289301
// We need to re-initialize the db/connection pool if options change
290302
d.needsRefresh = true
303+
304+
hasURI := false
305+
hasOtherOptions := false
306+
307+
if _, ok := options[adbc.OptionKeyURI]; ok {
308+
hasURI = true
309+
}
310+
311+
if len(options) > 1 || (len(options) == 1 && !hasURI) {
312+
hasOtherOptions = true
313+
}
314+
315+
if hasURI && hasOtherOptions {
316+
return adbc.Error{
317+
Code: adbc.StatusInvalidArgument,
318+
Msg: "cannot specify both URI and individual connection options",
319+
}
320+
}
321+
291322
for k, v := range options {
292323
err := d.SetOption(k, v)
293324
if err != nil {
@@ -301,6 +332,16 @@ func (d *databaseImpl) SetOption(key, value string) error {
301332
// We need to re-initialize the db/connection pool if options change
302333
d.needsRefresh = true
303334
switch key {
335+
case adbc.OptionKeyURI:
336+
// Strip the databricks:// scheme since databricks-sql-go expects raw DSN format
337+
if after, ok := strings.CutPrefix(value, "databricks://"); ok {
338+
d.uri = after
339+
} else {
340+
return adbc.Error{
341+
Code: adbc.StatusInvalidArgument,
342+
Msg: fmt.Sprintf("invalid URI scheme: expected 'databricks://', got '%s'", value),
343+
}
344+
}
304345
case OptionServerHostname:
305346
d.serverHostname = value
306347
case OptionHTTPPath:

go/docs/databricks.md

Lines changed: 59 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -33,8 +33,6 @@ dbc install databricks
3333

3434
## Connecting
3535

36-
TODO: This section once https://github.com/apache/arrow-adbc/pull/3771 is merged here.
37-
3836
To connect, edit the `uri` option below to match your environment and run the following:
3937

4038
```python
@@ -43,7 +41,7 @@ from adbc_driver_manager import dbapi
4341
conn = dbapi.connect(
4442
driver="databricks",
4543
db_kwargs = {
46-
"uri": "TODO"
44+
"uri": "databricks://token:dapi1234abcd5678efgh@dbc-a1b2345c-d6e7.cloud.databricks.com:443/sql/protocolv1/o/1234567890123456/1234-567890-abcdefgh"
4745
}
4846
)
4947
```
@@ -52,7 +50,64 @@ Note: The example above is for Python using the [adbc-driver-manager](https://py
5250

5351
### Connection String Format
5452

55-
TODO: This section once https://github.com/apache/arrow-adbc/pull/3771 is merged here.
53+
Databricks's URI syntax supports three primary forms:
54+
55+
1. Databricks personal access token authentication:
56+
57+
```
58+
databricks://token:<personal-access-token>@<server-hostname>:<port-number>/<http-path>?<param1=value1>&<param2=value2>
59+
```
60+
61+
Components:
62+
- `scheme`: `databricks://` (required)
63+
- `<personal-access-token>`: (required) Databricks personal access token.
64+
- `<server-hostname>`: (required) Server Hostname value.
65+
- `port-number`: (required) Port value, which is typically 443.
66+
- `http-path`: (required) HTTP Path value.
67+
- Query params: Databricks connection attributes. For complete list of optional parameters, see [Databricks Optional Parameters](https://docs.databricks.com/aws/en/dev-tools/go-sql-driver#optional-parameters)
68+
69+
70+
2. OAuth user-to-machine (U2M) authentication:
71+
72+
```
73+
databricks://<server-hostname>:<port-number>/<http-path>?authType=OauthU2M&<param1=value1>&<param2=value2>
74+
```
75+
76+
Components:
77+
- `scheme`: `databricks://` (required)
78+
- `<server-hostname>`: (required) Server Hostname value.
79+
- `port-number`: (required) Port value, which is typically 443.
80+
- `http-path`: (required) HTTP Path value.
81+
- `authType=OauthU2M`: (required) Specifies OAuth user-to-machine authentication.
82+
- Query params: Additional Databricks connection attributes. For complete list of optional parameters, see [Databricks Optional Parameters](https://docs.databricks.com/aws/en/dev-tools/go-sql-driver#optional-parameters)
83+
84+
3. OAuth machine-to-machine (M2M) authentication:
85+
86+
```
87+
databricks://<server-hostname>:<port-number>/<http-path>?authType=OAuthM2M&clientID=<client-id>&clientSecret=<client-secret>&<param1=value1>&<param2=value2>
88+
```
89+
90+
Components:
91+
- `scheme`: `databricks://` (required)
92+
- `<server-hostname>`: (required) Server Hostname value.
93+
- `port-number`: (required) Port value, which is typically 443.
94+
- `http-path`: (required) HTTP Path value.
95+
- `authType=OAuthM2M`: (required) Specifies OAuth machine-to-machine authentication.
96+
- `<client-id>`: (required) Service principal's UUID or Application ID value.
97+
- `<client-secret>`: (required) Secret value for the service principal's OAuth secret.
98+
- Query params: Additional Databricks connection attributes. For complete list of optional parameters, see [Databricks Optional Parameters](https://docs.databricks.com/aws/en/dev-tools/go-sql-driver#optional-parameters)
99+
100+
This follows the [Databricks SQL Driver for Go](https://docs.databricks.com/aws/en/dev-tools/go-sql-driver#connect-with-a-dsn-connection-string) format with the addition of the `databricks://` scheme.
101+
102+
:::{note}
103+
Reserved characters in URI elements must be URI-encoded. For example, `@` becomes `%40`.
104+
:::
105+
106+
Examples:
107+
108+
- `databricks://token:dapi1234abcd5678efgh@dbc-a1b2345c-d6e7.cloud.databricks.com:443/sql/protocolv1/o/1234567890123456/1234-567890-abcdefgh`
109+
- `databricks://myworkspace.cloud.databricks.com:443/sql/1.0/warehouses/abc123def456?authType=OauthU2M`
110+
- `databricks://myworkspace.cloud.databricks.com:443/sql/1.0/warehouses/abc123def456?authType=OAuthM2M&clientID=12345678-1234-1234-1234-123456789012&clientSecret=mysecret123`
56111

57112
## Feature & Type Support
58113

go/driver_test.go

Lines changed: 53 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,8 +29,8 @@ import (
2929
"strings"
3030
"testing"
3131

32+
"github.com/adbc-drivers/databricks/go/databricks"
3233
"github.com/apache/arrow-adbc/go/adbc"
33-
"github.com/apache/arrow-adbc/go/adbc/driver/databricks"
3434
"github.com/apache/arrow-adbc/go/adbc/validation"
3535
"github.com/apache/arrow-go/v18/arrow"
3636
"github.com/apache/arrow-go/v18/arrow/memory"
@@ -47,6 +47,7 @@ type DatabricksQuirks struct {
4747
httpPath string
4848
token string
4949
port string
50+
uri string // The URI to use for the test if set
5051
}
5152

5253
func (d *DatabricksQuirks) SetupDriver(t *testing.T) adbc.Driver {
@@ -59,6 +60,12 @@ func (d *DatabricksQuirks) TearDownDriver(t *testing.T, _ adbc.Driver) {
5960
}
6061

6162
func (d *DatabricksQuirks) DatabaseOptions() map[string]string {
63+
if d.uri != "" {
64+
return map[string]string{
65+
adbc.OptionKeyURI: d.uri,
66+
}
67+
}
68+
6269
opts := map[string]string{
6370
databricks.OptionServerHostname: d.hostname,
6471
databricks.OptionHTTPPath: d.httpPath,
@@ -327,6 +334,18 @@ func withQuirks(t *testing.T, fn func(*DatabricksQuirks)) {
327334
fn(q)
328335
}
329336

337+
func withQuirksURI(t *testing.T, fn func(*DatabricksQuirks)) {
338+
uri := os.Getenv("DATABRICKS_URI")
339+
if uri == "" {
340+
t.Skip("DATABRICKS_URI not defined, skipping URI tests")
341+
}
342+
343+
q := &DatabricksQuirks{
344+
uri: uri,
345+
}
346+
fn(q)
347+
}
348+
330349
func TestValidation(t *testing.T) {
331350
withQuirks(t, func(q *DatabricksQuirks) {
332351
suite.Run(t, &validation.DatabaseTests{Quirks: q})
@@ -341,6 +360,39 @@ func TestDatabricks(t *testing.T) {
341360
})
342361
}
343362

363+
func TestDatabricksWithURI(t *testing.T) {
364+
withQuirksURI(t, func(q *DatabricksQuirks) {
365+
drv := q.SetupDriver(t)
366+
defer q.TearDownDriver(t, drv)
367+
368+
db, err := drv.NewDatabase(q.DatabaseOptions())
369+
require.NoError(t, err)
370+
defer validation.CheckedClose(t, db)
371+
372+
ctx := context.Background()
373+
cnxn, err := db.Open(ctx)
374+
require.NoError(t, err)
375+
defer validation.CheckedClose(t, cnxn)
376+
377+
stmt, err := cnxn.NewStatement()
378+
require.NoError(t, err)
379+
defer validation.CheckedClose(t, stmt)
380+
381+
require.NoError(t, stmt.SetSqlQuery("SELECT 1 as test_col"))
382+
rdr, _, err := stmt.ExecuteQuery(ctx)
383+
require.NoError(t, err)
384+
defer rdr.Release()
385+
386+
assert.True(t, rdr.Next())
387+
rec := rdr.RecordBatch()
388+
assert.Equal(t, int64(1), rec.NumRows())
389+
assert.Equal(t, int64(1), rec.NumCols())
390+
assert.Equal(t, "test_col", rec.ColumnName(0))
391+
assert.False(t, rdr.Next())
392+
require.NoError(t, rdr.Err())
393+
})
394+
}
395+
344396
// ---- Additional Tests --------------------
345397

346398
type DatabricksTests struct {

0 commit comments

Comments
 (0)