End2end (#53)

- Added some good end 2 end tests with way to read server responses from
json files. One can now just create an example, set the `recordResults`
flag in client, and all server responses will be saved to JSON.
- Fix a few bugs regarding maxRows (can't be zero) and timezone.
- Added DSN parsing of userAgentEntry 
- Adds some docs
- Improved a few logs
- Added a very comprehensive example under example/workflow

Signed-off-by: Andre Furlan <[email protected]>

Author: andrefurlan-db

connection.go

@@ -89,7 +89,7 @@ func (c *conn) IsValid() bool {
 // Statement ExecContext is the same as connection ExecContext
 func (c *conn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) {
 	log := logger.WithContext(c.id, driverctx.CorrelationIdFromContext(ctx), "")
-	defer log.Duration(logger.Track("ExecContext"))
+	msg, start := logger.Track("ExecContext")
 	ctx = driverctx.NewContextWithConnId(ctx, c.id)
 	if len(args) > 0 {
 		return nil, errors.New(ErrParametersNotSupported)
@@ -99,6 +99,7 @@ func (c *conn) ExecContext(ctx context.Context, query string, args []driver.Name
 	if exStmtResp != nil && exStmtResp.OperationHandle != nil {
 		log = logger.WithContext(c.id, driverctx.CorrelationIdFromContext(ctx), client.SprintGuid(exStmtResp.OperationHandle.OperationId.GUID))
 	}
+	defer log.Duration(msg, start)
 
 	if err != nil {
 		log.Err(err).Msgf("databricks: failed to execute query: query %s", query)
@@ -243,13 +244,15 @@ func logBadQueryState(log *logger.DBSQLLogger, opStatus *cli_service.TGetOperati
 }
 
 func (c *conn) executeStatement(ctx context.Context, query string, args []driver.NamedValue) (*cli_service.TExecuteStatementResp, error) {
+	corrId := driverctx.CorrelationIdFromContext(ctx)
+	log := logger.WithContext(c.id, corrId, "")
 	sentinel := sentinel.Sentinel{
 		OnDoneFn: func(statusResp any) (any, error) {
 			req := cli_service.TExecuteStatementReq{
 				SessionHandle: c.session.SessionHandle,
 				Statement:     query,
 				RunAsync:      c.cfg.RunAsync,
-				QueryTimeout:  int64(c.cfg.QueryTimeoutSeconds),
+				QueryTimeout:  int64(c.cfg.QueryTimeout / time.Second),
 				// this is specific for databricks. It shortcuts server roundtrips
 				GetDirectResults: &cli_service.TSparkGetDirectResults{
 					MaxRows: int64(c.cfg.MaxRows),
@@ -262,8 +265,12 @@ func (c *conn) executeStatement(ctx context.Context, query string, args []driver
 			resp, err := c.client.ExecuteStatement(ctx, &req)
 			return resp, wrapErr(err, "failed to execute statement")
 		},
+		OnCancelFn: func() (any, error) {
+			log.Warn().Msg("databricks: execute statement canceled while creation operation")
+			return nil, nil
+		},
 	}
-	_, res, err := sentinel.Watch(ctx, c.cfg.PollInterval, 0)
+	_, res, err := sentinel.Watch(ctx, c.cfg.PollInterval, c.cfg.QueryTimeout)
 	if err != nil {
 		return nil, err
 	}

connector.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"database/sql/driver"
 	"fmt"
+	"strings"
+	"time"
 
 	"github.com/databricks/databricks-sql-go/driverctx"
 	"github.com/databricks/databricks-sql-go/internal/cli_service"
@@ -48,7 +50,7 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) {
 		},
 	}
 	// default timeout in here in addition to potential context timeout
-	_, res, err := sentinel.Watch(ctx, c.cfg.PollInterval, c.cfg.DefaultTimeout)
+	_, res, err := sentinel.Watch(ctx, c.cfg.PollInterval, c.cfg.ConnectTimeout)
 	if err != nil {
 		return nil, wrapErrf(err, "error connecting: host=%s port=%d, httpPath=%s", c.cfg.Host, c.cfg.Port, c.cfg.HTTPPath)
 	}
@@ -86,6 +88,8 @@ var _ driver.Connector = (*connector)(nil)
 
 type connOption func(*config.Config)
 
+// 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 NewConnector(options ...connOption) (driver.Connector, error) {
 	// config with default options
 	cfg := config.WithDefaults()
@@ -98,30 +102,38 @@ func NewConnector(options ...connOption) (driver.Connector, error) {
 	return &connector{cfg}, nil
 }
 
+// WithServerHostname sets up the server hostname. Mandatory.
 func WithServerHostname(host string) connOption {
 	return func(c *config.Config) {
+		if host == "localhost" {
+			c.Protocol = "http"
+		}
 		c.Host = host
 	}
 }
 
+// WithPort sets up the server port. Mandatory.
 func WithPort(port int) connOption {
 	return func(c *config.Config) {
 		c.Port = port
 	}
 }
 
+// WithAccessToken sets up the Personal Access Token. Mandatory for now.
 func WithAccessToken(token string) connOption {
 	return func(c *config.Config) {
 		c.AccessToken = token
 	}
 }
 
+// WithHTTPPath sets up the endpoint to the warehouse. Mandatory.
 func WithHTTPPath(path string) connOption {
 	return func(c *config.Config) {
 		c.HTTPPath = path
 	}
 }
 
+// WithMaxRows sets up the max rows fetched per request. Default is 10000
 func WithMaxRows(n int) connOption {
 	return func(c *config.Config) {
 		if n != 0 {
@@ -130,32 +142,43 @@ func WithMaxRows(n int) connOption {
 	}
 }
 
-// This will add a timeout for the server execution.
-// In seconds.
-func WithTimeout(n int) connOption {
+// WithTimeout adds timeout for the server query execution. Default is no timeout.
+func WithTimeout(n time.Duration) connOption {
 	return func(c *config.Config) {
-		c.QueryTimeoutSeconds = n
+		c.QueryTimeout = n
 	}
 }
 
+// Sets the initial catalog name and schema name in the session.
+// Use <select * from foo> instead of <select * from catalog.schema.foo>
 func WithInitialNamespace(catalog, schema string) connOption {
 	return func(c *config.Config) {
 		c.Catalog = catalog
 		c.Schema = schema
 	}
 }
 
+// Used to identify partners. Set as a string with format <isv-name+product-name>.
 func WithUserAgentEntry(entry string) connOption {
 	return func(c *config.Config) {
 		c.UserAgentEntry = entry
 	}
-
 }
 
-// Sessions params will be set upon opening the session
+// 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 WithSessionParams(params map[string]string) connOption {
 	return func(c *config.Config) {
+		for k, v := range params {
+			if strings.ToLower(k) == "timezone" {
+				if loc, err := time.LoadLocation(v); err != nil {
+					logger.Error().Msgf("timezone %s is not valid", v)
+				} else {
+					c.Location = loc
+				}
+
+			}
+		}
 		c.SessionParams = params
 	}
 }

doc.go

@@ -0,0 +1,4 @@
+/*
+Package dbsql implements the go driver to Databricks SQL
+*/
+package dbsql