Add or edit documentation above methods

Signed-off-by: Matthew Kim <[email protected]>

Author: mattdeekay

connection.go

@@ -21,16 +21,19 @@ type conn struct {
 	session *cli_service.TOpenSessionResp
 }
 
-// The driver does not really implement prepared statements.
+// Prepare prepares a statement with the query bound to this connection.
 func (c *conn) Prepare(query string) (driver.Stmt, error) {
 	return &stmt{conn: c, query: query}, nil
 }
 
-// The driver does not really implement prepared statements.
+// PrepareContext prepares a statement with the query bound to this connection.
+// Currently, PrepareContext does not use context and is functionally equivalent to Prepare.
 func (c *conn) PrepareContext(ctx context.Context, query string) (driver.Stmt, error) {
 	return &stmt{conn: c, query: query}, nil
 }
 
+// Close closes the session.
+// sql package maintains a free pool of connections and only calls Close when there's a surplus of idle connections.
 func (c *conn) Close() error {
 	log := logger.WithContext(c.id, "", "")
 	ctx := driverctx.NewContextWithConnId(context.Background(), c.id)
@@ -46,16 +49,18 @@ func (c *conn) Close() error {
 	return nil
 }
 
-// Not supported in Databricks
+// Not supported in Databricks.
 func (c *conn) Begin() (driver.Tx, error) {
 	return nil, errors.New(ErrTransactionsNotSupported)
 }
 
-// Not supported in Databricks
+// Not supported in Databricks.
 func (c *conn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
 	return nil, errors.New(ErrTransactionsNotSupported)
 }
 
+// Ping attempts to verify that the server is accessible.
+// Returns ErrBadConn if ping fails and consequently DB.Ping will remove the conn from the pool.
 func (c *conn) Ping(ctx context.Context) error {
 	log := logger.WithContext(c.id, driverctx.CorrelationIdFromContext(ctx), "")
 	ctx = driverctx.NewContextWithConnId(ctx, c.id)
@@ -69,12 +74,13 @@ func (c *conn) Ping(ctx context.Context) error {
 	return nil
 }
 
-// Implementation of SessionResetter
+// ResetSession is called prior to executing a query on the connection.
+// The session with this driver does not have any important state to reset before re-use.
 func (c *conn) ResetSession(ctx context.Context) error {
-	// For now our session does not have any important state to reset before re-use
 	return nil
 }
 
+// IsValid signals whether a connection is valid or if it should be discarded.
 func (c *conn) IsValid() bool {
 	return c.session.GetStatus().StatusCode == cli_service.TStatusCode_SUCCESS_STATUS
 }

connector.go

@@ -18,6 +18,7 @@ type connector struct {
 	cfg *config.Config
 }
 
+// Connect returns a connection to the Databricks database from a connection pool.
 func (c *connector) Connect(ctx context.Context) (driver.Conn, error) {
 	var catalogName *cli_service.TIdentifier
 	var schemaName *cli_service.TIdentifier
@@ -69,6 +70,7 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) {
 	return conn, nil
 }
 
+// Driver returns underlying databricksDriver for compatibility with sql.DB Driver method
 func (c *connector) Driver() driver.Driver {
 	return &databricksDriver{}
 }
@@ -77,7 +79,7 @@ var _ driver.Connector = (*connector)(nil)
 
 type connOption func(*config.Config)
 
-// NewConnector creates a connection that can be used with sql.OpenDB().
+// 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

driver.go

@@ -15,6 +15,8 @@ func init() {
 
 type databricksDriver struct{}
 
+// Open returns a new connection to Databricks database with a DSN string.
+// Use sql.Open("databricks", <dsn string>) after importing this driver package.
 func (d *databricksDriver) Open(dsn string) (driver.Conn, error) {
 	cfg := config.WithDefaults()
 	userCfg, err := config.ParseDSN(dsn)
@@ -28,6 +30,8 @@ func (d *databricksDriver) Open(dsn string) (driver.Conn, error) {
 	return c.Connect(context.Background())
 }
 
+// OpenConnector returns a new Connector.
+// Used by sql.DB to obtain a Connector and invoke its Connect method to obtain each needed connection.
 func (d *databricksDriver) OpenConnector(dsn string) (driver.Connector, error) {
 	cfg := config.WithDefaults()
 	ucfg, err := config.ParseDSN(dsn)

internal/client/client.go

@@ -16,14 +16,16 @@ import (
 	"github.com/pkg/errors"
 )
 
-// this is used to generate test data. Developer should change this manually
+// RecordResults is used to generate test data. Developer should change this manually
 var RecordResults bool
 var resultIndex int
 
 type ThriftServiceClient struct {
 	*cli_service.TCLIServiceClient
 }
 
+// OpenSession is a wrapper around the thrift operation OpenSession
+// If RecordResults is true, the results will be marshalled to JSON format and written to OpenSession<index>.json
 func (tsc *ThriftServiceClient) OpenSession(ctx context.Context, req *cli_service.TOpenSessionReq) (*cli_service.TOpenSessionResp, error) {
 	msg, start := logger.Track("OpenSession")
 	resp, err := tsc.TCLIServiceClient.OpenSession(ctx, req)
@@ -40,6 +42,8 @@ func (tsc *ThriftServiceClient) OpenSession(ctx context.Context, req *cli_servic
 	return resp, CheckStatus(resp)
 }
 
+// CloseSession is a wrapper around the thrift operation CloseSession
+// If RecordResults is true, the results will be marshalled to JSON format and written to CloseSession<index>.json
 func (tsc *ThriftServiceClient) CloseSession(ctx context.Context, req *cli_service.TCloseSessionReq) (*cli_service.TCloseSessionResp, error) {
 	log := logger.WithContext(driverctx.ConnIdFromContext(ctx), driverctx.CorrelationIdFromContext(ctx), "")
 	defer log.Duration(logger.Track("CloseSession"))
@@ -55,6 +59,8 @@ func (tsc *ThriftServiceClient) CloseSession(ctx context.Context, req *cli_servi
 	return resp, CheckStatus(resp)
 }
 
+// FetchResults is a wrapper around the thrift operation FetchResults
+// If RecordResults is true, the results will be marshalled to JSON format and written to FetchResults<index>.json
 func (tsc *ThriftServiceClient) FetchResults(ctx context.Context, req *cli_service.TFetchResultsReq) (*cli_service.TFetchResultsResp, error) {
 	log := logger.WithContext(driverctx.ConnIdFromContext(ctx), driverctx.CorrelationIdFromContext(ctx), SprintGuid(req.OperationHandle.OperationId.GUID))
 	defer log.Duration(logger.Track("FetchResults"))
@@ -70,6 +76,8 @@ func (tsc *ThriftServiceClient) FetchResults(ctx context.Context, req *cli_servi
 	return resp, CheckStatus(resp)
 }
 
+// GetResultSetMetadata is a wrapper around the thrift operation GetResultSetMetadata
+// If RecordResults is true, the results will be marshalled to JSON format and written to GetResultSetMetadata<index>.json
 func (tsc *ThriftServiceClient) GetResultSetMetadata(ctx context.Context, req *cli_service.TGetResultSetMetadataReq) (*cli_service.TGetResultSetMetadataResp, error) {
 	log := logger.WithContext(driverctx.ConnIdFromContext(ctx), driverctx.CorrelationIdFromContext(ctx), SprintGuid(req.OperationHandle.OperationId.GUID))
 	defer log.Duration(logger.Track("GetResultSetMetadata"))
@@ -85,6 +93,8 @@ func (tsc *ThriftServiceClient) GetResultSetMetadata(ctx context.Context, req *c
 	return resp, CheckStatus(resp)
 }
 
+// ExecuteStatement is a wrapper around the thrift operation ExecuteStatement
+// If RecordResults is true, the results will be marshalled to JSON format and written to ExecuteStatement<index>.json
 func (tsc *ThriftServiceClient) ExecuteStatement(ctx context.Context, req *cli_service.TExecuteStatementReq) (*cli_service.TExecuteStatementResp, error) {
 	msg, start := logger.Track("ExecuteStatement")
 	resp, err := tsc.TCLIServiceClient.ExecuteStatement(context.Background(), req)
@@ -106,6 +116,8 @@ func (tsc *ThriftServiceClient) ExecuteStatement(ctx context.Context, req *cli_s
 	return resp, CheckStatus(resp)
 }
 
+// GetOperationStatus is a wrapper around the thrift operation GetOperationStatus
+// If RecordResults is true, the results will be marshalled to JSON format and written to GetOperationStatus<index>.json
 func (tsc *ThriftServiceClient) GetOperationStatus(ctx context.Context, req *cli_service.TGetOperationStatusReq) (*cli_service.TGetOperationStatusResp, error) {
 	log := logger.WithContext(driverctx.ConnIdFromContext(ctx), driverctx.CorrelationIdFromContext(ctx), SprintGuid(req.OperationHandle.OperationId.GUID))
 	defer log.Duration(logger.Track("GetOperationStatus"))
@@ -121,6 +133,8 @@ func (tsc *ThriftServiceClient) GetOperationStatus(ctx context.Context, req *cli
 	return resp, CheckStatus(resp)
 }
 
+// CloseOperation is a wrapper around the thrift operation CloseOperation
+// If RecordResults is true, the results will be marshalled to JSON format and written to CloseOperation<index>.json
 func (tsc *ThriftServiceClient) CloseOperation(ctx context.Context, req *cli_service.TCloseOperationReq) (*cli_service.TCloseOperationResp, error) {
 	log := logger.WithContext(driverctx.ConnIdFromContext(ctx), driverctx.CorrelationIdFromContext(ctx), SprintGuid(req.OperationHandle.OperationId.GUID))
 	defer log.Duration(logger.Track("CloseOperation"))
@@ -136,6 +150,8 @@ func (tsc *ThriftServiceClient) CloseOperation(ctx context.Context, req *cli_ser
 	return resp, CheckStatus(resp)
 }
 
+// CancelOperation is a wrapper around the thrift operation CancelOperation
+// If RecordResults is true, the results will be marshalled to JSON format and written to CancelOperation<index>.json
 func (tsc *ThriftServiceClient) CancelOperation(ctx context.Context, req *cli_service.TCancelOperationReq) (*cli_service.TCancelOperationResp, error) {
 	log := logger.WithContext(driverctx.ConnIdFromContext(ctx), driverctx.CorrelationIdFromContext(ctx), SprintGuid(req.OperationHandle.OperationId.GUID))
 	defer log.Duration(logger.Track("CancelOperation"))
@@ -157,9 +173,8 @@ func (tsc *ThriftServiceClient) CancelOperation(ctx context.Context, req *cli_se
 // log.Debug().Msg(c.transport.response.Header.Get("x-thriftserver-error-message"))
 // log.Debug().Msg(c.transport.response.Header.Get("x-databricks-reason-phrase"))
 
-// This is a wrapper of the http transport so we can have access to response code and headers
+// InitThriftClient is a wrapper of the http transport, so we can have access to response code and headers.
 // It is important to know the code and headers to know if we need to retry or not
-
 func InitThriftClient(cfg *config.Config) (*ThriftServiceClient, error) {
 	endpoint := cfg.ToEndpointURL()
 	tcfg := &thrift.TConfiguration{
@@ -224,11 +239,13 @@ func InitThriftClient(cfg *config.Config) (*ThriftServiceClient, error) {
 	return tsClient, nil
 }
 
-// ThriftResponse respresents thrift rpc response
+// ThriftResponse represents the thrift rpc response
 type ThriftResponse interface {
 	GetStatus() *cli_service.TStatus
 }
 
+// CheckStatus checks the status code after a thrift operation.
+// Returns nil if the operation is successful or still executing, otherwise returns an error.
 func CheckStatus(resp interface{}) error {
 	rpcresp, ok := resp.(ThriftResponse)
 	if ok {
@@ -247,6 +264,7 @@ func CheckStatus(resp interface{}) error {
 	return errors.New("thrift: invalid response")
 }
 
+// SprintGuid is a convenience function to format a byte array into GUID.
 func SprintGuid(bts []byte) string {
 	if len(bts) == 16 {
 		return fmt.Sprintf("%x-%x-%x-%x-%x", bts[0:4], bts[4:6], bts[6:8], bts[8:10], bts[10:16])

internal/config/config.go

@@ -13,7 +13,7 @@ import (
 	"github.com/pkg/errors"
 )
 
-// Driver Configurations
+// Driver Configurations.
 // Only UserConfig are currently exposed to users
 type Config struct {
 	UserConfig
@@ -34,6 +34,7 @@ type Config struct {
 	ThriftDebugClientProtocol bool
 }
 
+// ToEndpointURL generates the endpoint URL from Config that a Thrift client will connect to
 func (c *Config) ToEndpointURL() string {
 	var userInfo string
 	if c.AccessToken != "" {
@@ -43,6 +44,7 @@ func (c *Config) ToEndpointURL() string {
 	return endpointUrl
 }
 
+// DeepCopy returns a true deep copy of Config
 func (c *Config) DeepCopy() *Config {
 	if c == nil {
 		return nil
@@ -84,6 +86,7 @@ type UserConfig struct {
 	SessionParams  map[string]string
 }
 
+// DeepCopy returns a true deep copy of UserConfig
 func (ucfg UserConfig) DeepCopy() UserConfig {
 	var sessionParams map[string]string
 	if ucfg.SessionParams != nil {
@@ -120,6 +123,7 @@ func (ucfg UserConfig) DeepCopy() UserConfig {
 
 var defaultMaxRows = 100000
 
+// WithDefaults provides default settings for optional fields in UserConfig
 func (ucfg UserConfig) WithDefaults() UserConfig {
 	if ucfg.MaxRows <= 0 {
 		ucfg.MaxRows = defaultMaxRows
@@ -131,6 +135,7 @@ func (ucfg UserConfig) WithDefaults() UserConfig {
 	return ucfg
 }
 
+// WithDefaults provides default settings for Config
 func WithDefaults() *Config {
 	return &Config{
 		UserConfig:                UserConfig{}.WithDefaults(),
@@ -152,6 +157,7 @@ func WithDefaults() *Config {
 
 }
 
+// ParseDSN constructs UserConfig by parsing DSN string supplied to `sql.Open()`
 func ParseDSN(dsn string) (UserConfig, error) {
 	fullDSN := dsn
 	if !strings.HasPrefix(dsn, "https://") && !strings.HasPrefix(dsn, "http://") {

logger/logger.go

@@ -14,10 +14,24 @@ type DBSQLLogger struct {
 	zerolog.Logger
 }
 
+// Track is a simple utility function to use with logger to log a message with a timestamp.
+// Recommended to use in conjunction with Duration.
+//
+// For example:
+//
+//	msg, start := log.Track("Run operation")
+//	defer log.Duration(msg, start)
 func (l *DBSQLLogger) Track(msg string) (string, time.Time) {
 	return msg, time.Now()
 }
 
+// Duration logs a debug message with the time elapsed between the provided start and the current time.
+// Use in conjunction with Track.
+//
+// For example:
+//
+//	msg, start := log.Track("Run operation")
+//	defer log.Duration(msg, start)
 func (l *DBSQLLogger) Duration(msg string, start time.Time) {
 	l.Debug().Msgf("%v elapsed time: %v", msg, time.Since(start))
 }
@@ -26,7 +40,7 @@ var Logger = &DBSQLLogger{
 	zerolog.New(os.Stderr).With().Timestamp().Logger(),
 }
 
-// enable pretty printing for interactive terminals and json for production.
+// Enable pretty printing for interactive terminals and json for production.
 func init() {
 	// for tty terminal enable pretty logs
 	if isatty.IsTerminal(os.Stdout.Fd()) && runtime.GOOS != "windows" {
@@ -114,7 +128,7 @@ func Err(err error) *zerolog.Event {
 	return Logger.Err(err)
 }
 
-// WithContext sets connectionId, correlationId, and queryID to be used as fields.
+// WithContext sets connectionId, correlationId, and queryId to be used as fields.
 func WithContext(connectionId string, correlationId string, queryId string) *DBSQLLogger {
 	return &DBSQLLogger{Logger.With().Str("connId", connectionId).Str("corrId", correlationId).Str("queryId", queryId).Logger()}
 }

result.go

@@ -9,10 +9,13 @@ type result struct {
 
 var _ driver.Result = (*result)(nil)
 
+// LastInsertId returns the database's auto-generated ID after an insert into a table.
+// This is currently not really implemented for this driver and will always return 0.
 func (res *result) LastInsertId() (int64, error) {
 	return res.InsertId, nil
 }
 
+// RowsAffected returns the number of rows affected by the query.
 func (res *result) RowsAffected() (int64, error) {
 	return res.AffectedRows, nil
 }

rows.go

@@ -43,6 +43,8 @@ var errRowsNoClient = "databricks: instance of Rows missing client"
 var errRowsNilRows = "databricks: nil Rows instance"
 var errRowsParseValue = "databricks: unable to parse %s value '%s' from column %s"
 
+// NewRows generates a new rows object given the rows' fields.
+// NewRows will also parse directResults if it is available for some rows' fields.
 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 {
 	r := &rows{
 		connId:        connID,

statement.go

@@ -17,15 +17,20 @@ func (s *stmt) Close() error {
 	return nil
 }
 
+// NumInput returns -1 and the sql package will not sanity check Exec or Query argument counts.
 func (s *stmt) NumInput() int {
 	return -1
 }
 
+// Exec is not implemented.
+//
 // Deprecated: Use StmtExecContext instead.
 func (s *stmt) Exec(args []driver.Value) (driver.Result, error) {
 	return nil, errors.New(ErrNotImplemented)
 }
 
+// Query is not implemented.
+//
 // Deprecated: Use StmtQueryContext instead.
 func (s *stmt) Query(args []driver.Value) (driver.Rows, error) {
 	return nil, errors.New(ErrNotImplemented)