dbsql

README

Databricks SQL Driver for Go

Description

This repo contains a Databricks SQL Driver for Go's database/sql package. It can be used to connect and query Databricks clusters and SQL Warehouses.

Documentation

See doc.go for full documentation or the Databrick's documentation for SQL Driver for Go.

Usage

import (
  "context"
  "database/sql"
  _ "github.com/databricks/databricks-sql-go"
)

db, err := sql.Open("databricks", "token:********@********.databricks.com:443/sql/1.0/endpoints/********")
if err != nil {
  panic(err)
}
defer db.Close()


rows, err := db.QueryContext(context.Background(), "SELECT 1")
defer rows.Close()

Additional usage examples are available here.

Connecting with DSN (Data Source Name)

The DSN format is:

token:[your token]@[Workspace hostname]:[Port number][Endpoint HTTP Path]?param=value

You can set query timeout value by appending a timeout query parameter (in seconds) and you can set max rows to retrieve per network request by setting the maxRows query parameter:

token:[your token]@[Workspace hostname]:[Port number][Endpoint HTTP Path]?timeout=1000&maxRows=1000

Connecting with a new Connector

You can also connect with a new connector object. For example:

import (
"database/sql"
  _ "github.com/databricks/databricks-sql-go"
)

connector, err := dbsql.NewConnector(
  dbsql.WithServerHostname(<Workspace hostname>),
  dbsql.WithPort(<Port number>),
  dbsql.WithHTTPPath(<Endpoint HTTP Path>),
  dbsql.WithAccessToken(<your token>)
)
if err != nil {
  log.Fatal(err)
}
db := sql.OpenDB(connector)
defer db.Close()

View doc.go or connector.go to understand all the functional options available when creating a new connector object.

Develop

Lint

We use golangci-lint as the lint tool. If you use vs code, just add the following settings:

{
    "go.lintTool": "golangci-lint",
    "go.lintFlags": [
        "--fast"
    ]
}

Unit Tests

go test

Issues

If you find any issues, feel free to create an issue or send a pull request directly.

Contributing

See CONTRIBUTING.md

License

Apache 2.0

Documentation

Overview

Package dbsql implements the go driver to Databricks SQL

Usage

Clients should use the database/sql package in conjunction with the driver:

import (
	"database/sql"

	_ "github.com/databricks/databricks-sql-go"
)

func main() {
	db, err := sql.Open("databricks", "token:<token>@<hostname>:<port>/<endpoint_path>")

	if err != nil {
		log.Fatal(err)
	}
	defer db.Close()
}

Connection via DSN (Data Source Name)

Use sql.Open() to create a database handle via a data source name string:

db, err := sql.Open("databricks", "<dsn_string>")

The DSN format is:

token:[my_token]@[hostname]:[port]/[endpoint http path]?param=value

Supported optional connection parameters can be specified in param=value and include:

• catalog: Sets the initial catalog name in the session
• schema: Sets the initial schema name in the session
• maxRows: Sets up the max rows fetched per request. Default is 100000
• timeout: Adds timeout (in seconds) for the server query execution. Default is no timeout
• userAgentEntry: Used to identify partners. Set as a string with format <isv-name+product-name>

Supported optional session parameters can be specified in param=value and include:

• ansi_mode: (Boolean string). Session statements will adhere to rules defined by ANSI SQL specification.
• timezone: (e.g. "America/Los_Angeles"). Sets the timezone of the session

Connection via new connector object

Use sql.OpenDB() to create a database handle via a new connector object created with dbsql.NewConnector():

import (
	"database/sql"
	dbsql "github.com/databricks/databricks-sql-go"
)

func main() {
	connector, err := dbsql.NewConnector(
		dbsql.WithServerHostname(<hostname>),
		dbsql.WithPort(<port>),
		dbsql.WithHTTPPath(<http_path>),
		dbsql.WithAccessToken(<my_token>)
	)
	if err != nil {
		log.Fatal(err)
	}

	db := sql.OpenDB(connector)
	defer db.Close()
	...
}

Supported functional options include:

• WithServerHostname(<hostname> string): Sets up the server hostname. Mandatory
• WithPort(<port> int): Sets up the server port. Mandatory
• WithAccessToken(<my_token> string): Sets up the Personal Access Token. Mandatory
• WithHTTPPath(<http_path> string): Sets up the endpoint to the warehouse. Mandatory
• WithInitialNamespace(<catalog> string, <schema> string): Sets up the catalog and schema name in the session. Optional
• WithMaxRows(<max_rows> int): Sets up the max rows fetched per request. Default is 100000. Optional
• WithSessionParams(<params_map> map[string]string): Sets up session parameters including "timezone" and "ansi_mode". Optional
• WithTimeout(<timeout> Duration). Adds timeout (in time.Duration) for the server query execution. Default is no timeout. Optional
• WithUserAgentEntry(<isv-name+product-name> string). Used to identify partners. Optional

Query cancellation and timeout

Cancelling a query via context cancellation or timeout is supported.

// Set up context timeout
ctx, cancel := context.WithTimeout(context.Background(), 30 * time.Second)
defer cancel()

// Execute query. Query will be cancelled after 30 seconds if still running
res, err := db.ExecContext(ctx, "CREATE TABLE example(id int, message string)")

CorrelationId and ConnId

Use the driverctx package under driverctx/ctx.go to add CorrelationId and ConnId to the context. CorrelationId and ConnId makes it convenient to parse and create metrics in logging.

**Connection Id** Internal id to track what happens under a connection. Connections can be reused so this would track across queries.

**Query Id** Internal id to track what happens under a query. Useful because the same query can be used with multiple connections.

**Correlation Id** External id, such as request ID, to track what happens under a request. Useful to track multiple connections in the same request.

ctx := dbsqlctx.NewContextWithCorrelationId(context.Background(), "workflow-example")

Logging

Use the logger package under logger.go to set up logging (from zerolog). By default, logging level is `warn`. If you want to disable logging, use `disabled`. The user can also utilize Track() and Duration() to custom log the elapsed time of anything tracked.

import (
	dbsqllog "github.com/databricks/databricks-sql-go/logger"
	dbsqlctx "github.com/databricks/databricks-sql-go/driverctx"
)

func main() {
	// Optional. Set the logging level with SetLogLevel()
	if err := dbsqllog.SetLogLevel("debug"); err != nil {
		log.Fatal(err)
	}

	// Optional. Set logging output with SetLogOutput()
	// Default is os.Stderr. If running in terminal, logger will use ConsoleWriter to prettify logs
	dbsqllog.SetLogOutput(os.Stdout)

	// Optional. Set correlation id with NewContextWithCorrelationId
	ctx := dbsqlctx.NewContextWithCorrelationId(context.Background(), "workflow-example")

	// Optional. Track time spent and log elapsed time
	msg, start := logger.Track("Run Main")
	defer log.Duration(msg, start)

	db, err := sql.Open("databricks", "<dsn_string>")
	...
}

The result log may look like this:

{"level":"debug","connId":"01ed6545-5669-1ec7-8c7e-6d8a1ea0ab16","corrId":"workflow-example","queryId":"01ed6545-57cc-188a-bfc5-d9c0eaf8e189","time":1668558402,"message":"Run Main elapsed time: 1.298712292s"}

Supported Data Types

==================================

Databricks Type --> Golang Type

==================================

BOOLEAN --> bool

TINYINT --> int8

SMALLINT --> int16

INT --> int32

BIGINT --> int64

FLOAT --> float32

DOUBLE --> float64

VOID --> nil

STRING --> string

DATE --> time.Time

TIMESTAMP --> time.Time

DECIMAL(p,s) --> sql.RawBytes

BINARY --> sql.RawBytes

ARRAY<elementType> --> sql.RawBytes

STRUCT --> sql.RawBytes

MAP<keyType, valueType> --> sql.RawBytes

INTERVAL (year-month) --> string

INTERVAL (day-time) --> string

For ARRAY, STRUCT, and MAP types, sql.Scan can cast sql.RawBytes to JSON string, which can be unmarshalled to Golang arrays, maps, and structs. For example:

type structVal struct {
		StringField string `json:"string_field"`
		ArrayField  []int  `json:"array_field"`
}
type row struct {
	arrayVal  []int
	mapVal    map[string]int
	structVal structVal
}
res := []row{}

for rows.Next() {
	r := row{}
	tempArray := []byte{}
	tempStruct := []byte{}
	tempMap := []byte{}
	if err := rows.Scan(&tempArray, &tempMap, &tempStruct); err != nil {
		log.Fatal(err)
	}
	if err := json.Unmarshal(tempArray, &r.arrayVal); err != nil {
		log.Fatal(err)
	}
	if err := json.Unmarshal(tempMap, &r.mapVal); err != nil {
		log.Fatal(err)
	}
	if err := json.Unmarshal(tempStruct, &r.structVal); err != nil {
		log.Fatal(err)
	}
	res = append(res, r)
}

May generate the following row:

{arrayVal:[1,2,3] mapVal:{"key1":1} structVal:{"string_field":"string_val","array_field":[4,5,6]}}

Index

• Variables
• func NewConnector(options ...connOption) (driver.Connector, error)
• func NewRows(connID string, corrId string, client cli_service.TCLIService, ...) driver.Rows
• func WithAccessToken(token string) connOption
• func WithHTTPPath(path string) connOption
• func WithInitialNamespace(catalog, schema string) connOption
• func WithMaxRows(n int) connOption
• func WithPort(port int) connOption
• func WithServerHostname(host string) connOption
• func WithSessionParams(params map[string]string) connOption
• func WithTimeout(n time.Duration) connOption
• func WithUserAgentEntry(entry string) connOption

Variables

var ErrNotImplemented = "databricks: not implemented"

var ErrParametersNotSupported = "databricks: query parameters are not supported"

var ErrTransactionsNotSupported = "databricks: transactions are not supported"

Functions

func NewConnector

func NewConnector(options ...connOption) (driver.Connector, error)

NewConnector creates a connection that can be used with `sql.OpenDB()`. This is an easier way to set up the DB instead of having to construct a DSN string.

func NewRows

func NewRows(connID string, corrId string, client cli_service.TCLIService, opHandle *cli_service.TOperationHandle, pageSize int64, location *time.Location, directResults *cli_service.TSparkDirectResults) driver.Rows

NewRows generates a new rows object given the rows' fields. NewRows will also parse directResults if it is available for some rows' fields.

func WithAccessToken

func WithAccessToken(token string) connOption

WithAccessToken sets up the Personal Access Token. Mandatory for now.

func WithHTTPPath

func WithHTTPPath(path string) connOption

WithHTTPPath sets up the endpoint to the warehouse. Mandatory.

func WithInitialNamespace

func WithInitialNamespace(catalog, schema string) connOption

Sets the initial catalog name and schema name in the session. Use <select * from foo> instead of <select * from catalog.schema.foo>

func WithMaxRows

func WithMaxRows(n int) connOption

WithMaxRows sets up the max rows fetched per request. Default is 10000

func WithPort

func WithPort(port int) connOption

WithPort sets up the server port. Mandatory.

func WithServerHostname

func WithServerHostname(host string) connOption

WithServerHostname sets up the server hostname. Mandatory.

func WithSessionParams

func WithSessionParams(params map[string]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 ..."

func WithTimeout

func WithTimeout(n time.Duration) connOption

WithTimeout adds timeout for the server query execution. Default is no timeout.

func WithUserAgentEntry

func WithUserAgentEntry(entry string) connOption

Used to identify partners. Set as a string with format <isv-name+product-name>.