Implementing error types (#117)

Updated driver error handling to standardize on three error types:
driver error (DBDriverError), request error (DBRequestError), and sql
execution error (DBExecutionError). All three types have accessors for
stack trace, correlation id, and connection id. DBExecutionError also
has accessors for query Id and sql state.
The example in examples/error shows how to use errors.Is() and
errors.As() to work with the error chain to determine error type and
extract the state.

---------

Signed-off-by: Matthew Kim <11141331+mattdeekay@users.noreply.github.com>
Signed-off-by: Raymond Cypher <raymond.cypher@databricks.com>
Co-authored-by: Matthew Kim <11141331+mattdeekay@users.noreply.github.com>

Author: rcypher-databricks

connection.go

@@ -6,10 +6,11 @@ import (
 	"time"
 
 	"github.com/databricks/databricks-sql-go/driverctx"
+	dbsqlerr "github.com/databricks/databricks-sql-go/errors"
 	"github.com/databricks/databricks-sql-go/internal/cli_service"
 	"github.com/databricks/databricks-sql-go/internal/client"
 	"github.com/databricks/databricks-sql-go/internal/config"
-	dbsqlerr "github.com/databricks/databricks-sql-go/internal/err"
+	dbsqlerrint "github.com/databricks/databricks-sql-go/internal/errors"
 	"github.com/databricks/databricks-sql-go/internal/rows"
 	"github.com/databricks/databricks-sql-go/internal/sentinel"
 	"github.com/databricks/databricks-sql-go/logger"
@@ -46,19 +47,19 @@ func (c *conn) Close() error {
 
 	if err != nil {
 		log.Err(err).Msg("databricks: failed to close connection")
-		return dbsqlerr.WrapErr(err, "failed to close connection")
+		return dbsqlerrint.NewRequestError(ctx, dbsqlerr.ErrCloseConnection, err)
 	}
 	return nil
 }
 
 // Not supported in Databricks.
 func (c *conn) Begin() (driver.Tx, error) {
-	return nil, errors.New(dbsqlerr.ErrTransactionsNotSupported)
+	return nil, dbsqlerrint.NewDriverError(context.TODO(), dbsqlerr.ErrNotImplemented, nil)
 }
 
 // Not supported in Databricks.
 func (c *conn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
-	return nil, errors.New(dbsqlerr.ErrTransactionsNotSupported)
+	return nil, dbsqlerrint.NewDriverError(context.TODO(), dbsqlerr.ErrNotImplemented, nil)
 }
 
 // Ping attempts to verify that the server is accessible.
@@ -100,7 +101,7 @@ func (c *conn) ExecContext(ctx context.Context, query string, args []driver.Name
 
 	ctx = driverctx.NewContextWithConnId(ctx, c.id)
 	if len(args) > 0 {
-		return nil, errors.New(dbsqlerr.ErrParametersNotSupported)
+		return nil, dbsqlerrint.NewDriverError(ctx, dbsqlerr.ErrParametersNotSupported, nil)
 	}
 	exStmtResp, opStatusResp, err := c.runQuery(ctx, query, args)
 
@@ -122,7 +123,7 @@ func (c *conn) ExecContext(ctx context.Context, query string, args []driver.Name
 	}
 	if err != nil {
 		log.Err(err).Msgf("databricks: failed to execute query: query %s", query)
-		return nil, dbsqlerr.WrapErrf(err, "failed to execute query")
+		return nil, dbsqlerrint.NewExecutionError(ctx, dbsqlerr.ErrQueryExecution, err, opStatusResp)
 	}
 
 	res := result{AffectedRows: opStatusResp.GetNumModifiedRows()}
@@ -142,20 +143,21 @@ func (c *conn) QueryContext(ctx context.Context, query string, args []driver.Nam
 
 	ctx = driverctx.NewContextWithConnId(ctx, c.id)
 	if len(args) > 0 {
-		return nil, errors.New(dbsqlerr.ErrParametersNotSupported)
+		return nil, dbsqlerrint.NewDriverError(ctx, dbsqlerr.ErrParametersNotSupported, nil)
 	}
 	// first we try to get the results synchronously.
 	// at any point in time that the context is done we must cancel and return
-	exStmtResp, _, err := c.runQuery(ctx, query, args)
+	exStmtResp, opStatusResp, err := c.runQuery(ctx, query, args)
 
 	if exStmtResp != nil && exStmtResp.OperationHandle != nil {
+		ctx = driverctx.NewContextWithQueryId(ctx, client.SprintGuid(exStmtResp.OperationHandle.OperationId.GUID))
 		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).Msg("databricks: failed to run query") // To log query we need to redact credentials
-		return nil, dbsqlerr.WrapErrf(err, "failed to run query")
+		return nil, dbsqlerrint.NewExecutionError(ctx, dbsqlerr.ErrQueryExecution, err, opStatusResp)
 	}
 	// hold on to the operation handle
 	opHandle := exStmtResp.OperationHandle
@@ -177,9 +179,10 @@ func (c *conn) runQuery(ctx context.Context, query string, args []driver.NamedVa
 	}
 	opHandle := exStmtResp.OperationHandle
 	if opHandle != nil && opHandle.OperationId != nil {
+		ctx = driverctx.NewContextWithQueryId(ctx, client.SprintGuid(opHandle.OperationId.GUID))
 		log = logger.WithContext(
 			c.id,
-			driverctx.CorrelationIdFromContext(ctx), client.SprintGuid(opHandle.OperationId.GUID),
+			driverctx.CorrelationIdFromContext(ctx), driverctx.QueryIdFromContext(ctx),
 		)
 	}
 
@@ -217,16 +220,16 @@ func (c *conn) runQuery(ctx context.Context, query string, args []driver.NamedVa
 				cli_service.TOperationState_ERROR_STATE,
 				cli_service.TOperationState_TIMEDOUT_STATE:
 				logBadQueryState(log, statusResp)
-				return exStmtResp, statusResp, errors.New(statusResp.GetDisplayMessage())
+				return exStmtResp, statusResp, dbsqlerrint.NewRequestError(ctx, dbsqlerr.ErrInvalidOperationState, nil)
 				// live states
 			default:
 				logBadQueryState(log, statusResp)
-				return exStmtResp, statusResp, errors.New("invalid operation state. This should not have happened")
+				return exStmtResp, statusResp, dbsqlerrint.NewDriverError(ctx, dbsqlerr.ErrInvalidOperationState, nil)
 			}
 		// weird states
 		default:
 			logBadQueryState(log, opStatus)
-			return exStmtResp, opStatus, errors.New("invalid operation state. This should not have happened")
+			return exStmtResp, opStatus, dbsqlerrint.NewDriverError(ctx, dbsqlerr.ErrInvalidOperationState, nil)
 		}
 
 	} else {
@@ -245,11 +248,11 @@ func (c *conn) runQuery(ctx context.Context, query string, args []driver.NamedVa
 			cli_service.TOperationState_ERROR_STATE,
 			cli_service.TOperationState_TIMEDOUT_STATE:
 			logBadQueryState(log, statusResp)
-			return exStmtResp, statusResp, errors.New(statusResp.GetDisplayMessage())
+			return exStmtResp, statusResp, dbsqlerrint.NewDriverError(ctx, dbsqlerr.ErrInvalidOperationState, nil)
 			// live states
 		default:
 			logBadQueryState(log, statusResp)
-			return exStmtResp, statusResp, errors.New("invalid operation state. This should not have happened")
+			return exStmtResp, statusResp, dbsqlerrint.NewDriverError(ctx, dbsqlerr.ErrInvalidOperationState, nil)
 		}
 	}
 }
@@ -311,7 +314,6 @@ func (c *conn) executeStatement(ctx context.Context, query string, args []driver
 			} else {
 				log.Debug().Msgf("databricks: cancel success")
 			}
-
 		} else {
 			log.Debug().Msg("databricks: query did not need cancellation")
 		}
@@ -337,6 +339,7 @@ func (c *conn) pollOperation(ctx context.Context, opHandle *cli_service.TOperati
 			statusResp, err = c.client.GetOperationStatus(newCtx, &cli_service.TGetOperationStatusReq{
 				OperationHandle: opHandle,
 			})
+
 			if statusResp != nil && statusResp.OperationState != nil {
 				log.Debug().Msgf("databricks: status %s", statusResp.GetOperationState().String())
 			}
@@ -363,13 +366,17 @@ func (c *conn) pollOperation(ctx context.Context, opHandle *cli_service.TOperati
 			return ret, err
 		},
 	}
-	_, resp, err := pollSentinel.Watch(ctx, c.cfg.PollInterval, 0)
+	status, resp, err := pollSentinel.Watch(ctx, c.cfg.PollInterval, 0)
 	if err != nil {
-		return nil, dbsqlerr.WrapErr(err, "failed to poll query state")
+		if status == sentinel.WatchTimeout {
+			err = dbsqlerrint.NewRequestError(ctx, dbsqlerr.ErrSentinelTimeout, err)
+		}
+		return nil, err
 	}
+
 	statusResp, ok := resp.(*cli_service.TGetOperationStatusResp)
 	if !ok {
-		return nil, errors.New("could not read query status")
+		return nil, dbsqlerrint.NewDriverError(ctx, dbsqlerr.ErrReadQueryStatus, nil)
 	}
 	return statusResp, nil
 }

connection_test.go

@@ -161,7 +161,7 @@ func TestConn_executeStatement(t *testing.T) {
 			if opTest.err == "" {
 				assert.NoError(t, err)
 			} else {
-				assert.EqualError(t, err, opTest.err)
+				assert.EqualError(t, err, "databricks: execution error: failed to execute query: "+opTest.err)
 			}
 			assert.Equal(t, 1, executeStatementCount)
 			assert.Equal(t, opTest.closeOperationCount, closeOperationCount)
@@ -539,6 +539,7 @@ func TestConn_pollOperation(t *testing.T) {
 		ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
 		defer cancel()
 		res, err := testConn.pollOperation(ctx, &cli_service.TOperationHandle{
+
 			OperationId: &cli_service.THandleIdentifier{
 				GUID:   []byte{1, 2, 3, 4, 2, 23, 4, 2, 3, 1, 2, 4, 4, 223, 34, 54},
 				Secret: []byte("b"),

connector.go

@@ -10,10 +10,11 @@ import (
 
 	"github.com/databricks/databricks-sql-go/auth/pat"
 	"github.com/databricks/databricks-sql-go/driverctx"
+	dbsqlerr "github.com/databricks/databricks-sql-go/errors"
 	"github.com/databricks/databricks-sql-go/internal/cli_service"
 	"github.com/databricks/databricks-sql-go/internal/client"
 	"github.com/databricks/databricks-sql-go/internal/config"
-	dbsqlerr "github.com/databricks/databricks-sql-go/internal/err"
+	dbsqlerrint "github.com/databricks/databricks-sql-go/internal/errors"
 	"github.com/databricks/databricks-sql-go/logger"
 )
 
@@ -35,7 +36,7 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) {
 
 	tclient, err := client.InitThriftClient(c.cfg, c.client)
 	if err != nil {
-		return nil, dbsqlerr.WrapErr(err, "error initializing thrift client")
+		return nil, dbsqlerrint.NewDriverError(ctx, dbsqlerr.ErrThriftClient, err)
 	}
 	protocolVersion := int64(c.cfg.ThriftProtocolVersion)
 	session, err := tclient.OpenSession(ctx, &cli_service.TOpenSessionReq{
@@ -49,7 +50,7 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) {
 	})
 
 	if err != nil {
-		return nil, dbsqlerr.WrapErrf(err, "error connecting: host=%s port=%d, httpPath=%s", c.cfg.Host, c.cfg.Port, c.cfg.HTTPPath)
+		return nil, dbsqlerrint.NewRequestError(ctx, fmt.Sprintf("error connecting: host=%s port=%d, httpPath=%s", c.cfg.Host, c.cfg.Port, c.cfg.HTTPPath), err)
 	}
 
 	conn := &conn{

doc.go

@@ -144,6 +144,52 @@ 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"}
 
+# Errors
+
+There are three error types exposed via dbsql/errors
+
+	DBDriverError - An error in the go driver. Example: unimplemented functionality, invalid driver state, errors processing a server response, etc.
+
+	DBRequestError - An error that is caused by an invalid request. Example: permission denied, invalid http path or other connection parameter, resource not available, etc.
+
+	DBExecutionError - Any error that occurs after the SQL statement has been accepted such as a SQL syntax error, missing table, etc.
+
+Each type has a corresponding sentinel value which can be used with errors.Is() to determine if one of the types is present in an error chain.
+
+	DriverError
+	RequestError
+	ExecutionError
+
+Example usage:
+
+	import (
+		fmt
+		errors
+		dbsqlerr "github.com/databricks/databricks-sql-go/errors"
+	)
+
+	func main() {
+		...
+		_, err := db.ExecContext(ogCtx, `Select id from range(100)`)
+		if err != nil {
+			if errors.Is(err, dbsqlerr.ExecutionError) {
+				var execErr dbsqlerr.DBExecutionError
+				if ok := errors.As(err, &execError); ok {
+						fmt.Printf("%s, corrId: %s, connId: %s, queryId: %s, sqlState: %s",
+						execErr.Error(),
+						execErr.CorrelationId(),
+						execErr.ConnectionId(),
+						execErr.QueryId(),
+						execErr.SqlState())
+				}
+			}
+			...
+		}
+		...
+	}
+
+See the documentation for dbsql/errors for more information.
+
 # Supported Data Types
 
 ==================================

driver_e2e_test.go

@@ -13,10 +13,11 @@ import (
 	"time"
 
 	"github.com/databricks/databricks-sql-go/driverctx"
+	dbsqlerr "github.com/databricks/databricks-sql-go/errors"
 	"github.com/databricks/databricks-sql-go/internal/cli_service"
 	"github.com/databricks/databricks-sql-go/internal/client"
-	dbsqlerr "github.com/databricks/databricks-sql-go/internal/err"
 	"github.com/databricks/databricks-sql-go/logger"
+	"github.com/pkg/errors"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -281,11 +282,19 @@ func TestContextTimeoutExample(t *testing.T) {
 	ctx1, cancel := context.WithTimeout(ogCtx, 5*time.Second)
 	defer cancel()
 	rows, err := db.QueryContext(ctx1, `SELECT id FROM RANGE(100000000) ORDER BY RANDOM() + 2 asc`)
-	require.ErrorContains(t, err, context.DeadlineExceeded.Error())
+	if err, ok := err.(interface{ StackTrace() errors.StackTrace }); ok {
+		fmt.Printf("Stack trace: %v", err.StackTrace())
+	}
+	require.True(t, errors.Is(err, context.DeadlineExceeded))
+	require.True(t, errors.Is(err, dbsqlerr.ExecutionError))
+	var ee dbsqlerr.DBExecutionError
+	require.True(t, errors.As(err, &ee))
+	require.Equal(t, "context-timeout-example", ee.CorrelationId())
 	require.Nil(t, rows)
-	_, ok := err.(dbsqlerr.Causer)
+
+	_, ok := err.(interface{ Cause() error })
 	assert.True(t, ok)
-	_, ok = err.(dbsqlerr.StackTracer)
+	_, ok = err.(interface{ StackTrace() errors.StackTrace })
 	assert.True(t, ok)
 	assert.Equal(t, 1, state.executeStatementCalls)
 	assert.GreaterOrEqual(t, state.getOperationStatusCalls, 1)

driverctx/ctx.go

@@ -11,15 +11,24 @@ type contextKey int
 const (
 	CorrelationIdContextKey contextKey = iota
 	ConnIdContextKey
+	QueryIdContextKey
+	QueryIdCallbackKey
+	ConnIdCallbackKey
 )
 
+type IdCallbackFunc func(string)
+
 // NewContextWithCorrelationId creates a new context with correlationId value. Used by Logger to populate field corrId.
 func NewContextWithCorrelationId(ctx context.Context, correlationId string) context.Context {
 	return context.WithValue(ctx, CorrelationIdContextKey, correlationId)
 }
 
 // CorrelationIdFromContext retrieves the correlationId stored in context.
 func CorrelationIdFromContext(ctx context.Context) string {
+	if ctx == nil {
+		return ""
+	}
+
 	corrId, ok := ctx.Value(CorrelationIdContextKey).(string)
 	if !ok {
 		return ""
@@ -29,14 +38,51 @@ func CorrelationIdFromContext(ctx context.Context) string {
 
 // NewContextWithConnId creates a new context with connectionId value.
 func NewContextWithConnId(ctx context.Context, connId string) context.Context {
+	if callback, ok := ctx.Value(ConnIdCallbackKey).(IdCallbackFunc); ok {
+		callback(connId)
+	}
 	return context.WithValue(ctx, ConnIdContextKey, connId)
 }
 
 // ConnIdFromContext retrieves the connectionId stored in context.
 func ConnIdFromContext(ctx context.Context) string {
+	if ctx == nil {
+		return ""
+	}
+
 	connId, ok := ctx.Value(ConnIdContextKey).(string)
 	if !ok {
 		return ""
 	}
 	return connId
 }
+
+// NewContextWithQueryId creates a new context with queryId value.
+func NewContextWithQueryId(ctx context.Context, queryId string) context.Context {
+	if callback, ok := ctx.Value(QueryIdCallbackKey).(IdCallbackFunc); ok {
+		callback(queryId)
+	}
+
+	return context.WithValue(ctx, QueryIdContextKey, queryId)
+}
+
+// QueryIdFromContext retrieves the queryId stored in context.
+func QueryIdFromContext(ctx context.Context) string {
+	if ctx == nil {
+		return ""
+	}
+
+	queryId, ok := ctx.Value(QueryIdContextKey).(string)
+	if !ok {
+		return ""
+	}
+	return queryId
+}
+
+func NewContextWithQueryIdCallback(ctx context.Context, callback IdCallbackFunc) context.Context {
+	return context.WithValue(ctx, QueryIdCallbackKey, callback)
+}
+
+func NewContextWithConnIdCallback(ctx context.Context, callback IdCallbackFunc) context.Context {
+	return context.WithValue(ctx, ConnIdCallbackKey, callback)
+}