Logging, debugging, and tracing improvements (#52)

## Error and tracing improvements
Now every error we produce is created with `errors.New()` or it is
wrapped with `wrapErr()`. This make sure that a stack tracing is added
to the error as soon as possible.

Errors now are also properly labeled and descriptive. E.g.:
`log.Err(err).Msg("databricks: failed to ping")`

In debug mode, we also get logging of time spent in function.

## Logging improvements

```go
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
ctx = driverctx.NewContextWithCorrelationId(ctx, "abc")
defer cancel1()

rows, err1 := db.QueryContext(ctx, `select * from quickstart_table1`)
```

the log looks like this:
```
{"level":"debug","connId":"01ed6545-5669-1ec7-8c7e-6d8a1ea0ab16","corrId":"abc","queryId":"01ed6545-57cc-188a-bfc5-d9c0eaf8e189","time":1668558402,"message":"QueryContext elapsed time: 1.298712292s"}
```
This format makes it very easy to parse and create metrics in logging
tools.

**Connection ID**
Pretty much everything we do with the driver is based on a connection
id. This is our internal id to track what happens under that connection.
Remember that connections can be reused so this would track across
queries.

**Correlation ID**
If our users have an external ID, such as a request ID, we provide a way
to add it to the context so we can log with it. This helps in cases
where multiple connections are used in the same request.

**Query ID**
If we have a query ID, we also add to the logs so we can have another
layer of tracking, as sometimes the same query can be used with multiple
connections.

A new package `driverctx` was added to help adding specific driver
values to the context.

```go
ctx := driverctx.NewContextWithConnId(context.Background(), r.connId)
ctx = driverctx.NewContextWithCorrelationId(ctx, r.correlationId)
```

The `logger` package has 3 new public methods:

```go
func WithContext(connectionId string, correlationId string, queryId string) *DBSQLLogger {
   return &DBSQLLogger{Logger.With().Str("connId", connectionId).Str("corrId", correlationId).Str("queryId", queryId).Logger()}
}
 
func Track(msg string) (string, time.Time) {
   return msg, time.Now()
}
 
func Duration(msg string, start time.Time) {
   Logger.Debug().Msgf("%v elapsed time: %v", msg, time.Since(start))
}
``` 

Client does something like this:
```go
   log := logger.WithContext(queryctx.ConnIdFromContext(ctx), queryctx.CorrelationIdFromContext(ctx), "")
   defer log.Duration(logger.Track("OpenSession"))
```
Signed-off-by: Andre Furlan <[email protected]>

Author: andrefurlan-db

.github/workflows/go.yml

@@ -20,7 +20,7 @@ jobs:
         with:
           version: v1.48
   build-and-test:
-    name: Build
+    name: Test and Build
     strategy:
       matrix:
         go-version: [1.19.x]

connection.go

@@ -3,16 +3,19 @@ package dbsql
 import (
 	"context"
 	"database/sql/driver"
-	"fmt"
 	"time"
 
+	"github.com/databricks/databricks-sql-go/driverctx"
 	"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"
 	"github.com/databricks/databricks-sql-go/internal/sentinel"
 	"github.com/databricks/databricks-sql-go/logger"
+	"github.com/pkg/errors"
 )
 
 type conn struct {
+	id      string
 	cfg     *config.Config
 	client  cli_service.TCLIService
 	session *cli_service.TOpenSessionResp
@@ -29,32 +32,41 @@ func (c *conn) PrepareContext(ctx context.Context, query string) (driver.Stmt, e
 }
 
 func (c *conn) Close() error {
+	log := logger.WithContext(c.id, "", "")
+	ctx := driverctx.NewContextWithConnId(context.Background(), c.id)
 	sentinel := sentinel.Sentinel{
 		OnDoneFn: func(statusResp any) (any, error) {
-			return c.client.CloseSession(context.Background(), &cli_service.TCloseSessionReq{
+			return c.client.CloseSession(ctx, &cli_service.TCloseSessionReq{
 				SessionHandle: c.session.SessionHandle,
 			})
 		},
 	}
-	_, _, err := sentinel.Watch(context.Background(), c.cfg.PollInterval, 15*time.Second)
-
-	return err
+	_, _, err := sentinel.Watch(ctx, c.cfg.PollInterval, 15*time.Second)
+	if err != nil {
+		log.Err(err).Msg("databricks: failed to close connection")
+		return wrapErr(err, "failed to close connection")
+	}
+	return nil
 }
 
 // Not supported in Databricks
 func (c *conn) Begin() (driver.Tx, error) {
-	return nil, fmt.Errorf("databricks: transactions are not supported")
+	return nil, errors.New(ErrTransactionsNotSupported)
 }
 
 // Not supported in Databricks
 func (c *conn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
-	return nil, fmt.Errorf("databricks: transactions are not supported")
+	return nil, errors.New(ErrTransactionsNotSupported)
 }
 
 func (c *conn) Ping(ctx context.Context) error {
-	_, err := c.QueryContext(ctx, "select 1", nil)
+	log := logger.WithContext(c.id, driverctx.CorrelationIdFromContext(ctx), "")
+	ctx = driverctx.NewContextWithConnId(ctx, c.id)
+	ctx1, cancel := context.WithTimeout(ctx, 15*time.Second)
+	defer cancel()
+	_, err := c.QueryContext(ctx1, "select 1", nil)
 	if err != nil {
-		logger.Err(err).Msg("ping error")
+		log.Err(err).Msg("databricks: failed to ping")
 		return driver.ErrBadConn
 	}
 	return nil
@@ -76,13 +88,21 @@ func (c *conn) IsValid() bool {
 // ExecContext honors the context timeout and return when it is canceled.
 // 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"))
+	ctx = driverctx.NewContextWithConnId(ctx, c.id)
 	if len(args) > 0 {
-		return nil, fmt.Errorf("databricks: query parameters are not supported")
+		return nil, errors.New(ErrParametersNotSupported)
+	}
+	exStmtResp, opStatusResp, err := c.runQuery(ctx, query, args)
+
+	if exStmtResp != nil && exStmtResp.OperationHandle != nil {
+		log = logger.WithContext(c.id, driverctx.CorrelationIdFromContext(ctx), client.SprintGuid(exStmtResp.OperationHandle.OperationId.GUID))
 	}
-	_, opStatusResp, err := c.runQuery(ctx, query, args)
 
 	if err != nil {
-		return nil, err
+		log.Err(err).Msgf("databricks: failed to execute query: query %s", query)
+		return nil, wrapErrf(err, "failed to execute query")
 	}
 	res := result{AffectedRows: opStatusResp.GetNumModifiedRows()}
 
@@ -95,24 +115,37 @@ func (c *conn) ExecContext(ctx context.Context, query string, args []driver.Name
 // QueryContext honors the context timeout and return when it is canceled.
 // Statement QueryContext is the same as connection QueryContext
 func (c *conn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {
+	corrId := driverctx.CorrelationIdFromContext(ctx)
+	log := logger.WithContext(c.id, corrId, "")
+	msg, start := log.Track("QueryContext")
+
+	ctx = driverctx.NewContextWithConnId(ctx, c.id)
 	if len(args) > 0 {
-		return nil, fmt.Errorf("databricks: query parameters are not supported")
+		return nil, errors.New(ErrParametersNotSupported)
 	}
 	// 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)
 
+	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 {
-		return nil, err
+		log.Err(err).Msgf("databricks: failed to run query: query %s", query)
+		return nil, wrapErrf(err, "failed to run query")
 	}
 	// hold on to the operation handle
 	opHandle := exStmtResp.OperationHandle
 
 	rows := rows{
-		client:   c.client,
-		opHandle: opHandle,
-		pageSize: int64(c.cfg.MaxRows),
-		location: c.cfg.Location,
+		connId:        c.id,
+		correlationId: corrId,
+		client:        c.client,
+		opHandle:      opHandle,
+		pageSize:      int64(c.cfg.MaxRows),
+		location:      c.cfg.Location,
 	}
 
 	if exStmtResp.DirectResults != nil {
@@ -126,15 +159,19 @@ func (c *conn) QueryContext(ctx context.Context, query string, args []driver.Nam
 }
 
 func (c *conn) runQuery(ctx context.Context, query string, args []driver.NamedValue) (*cli_service.TExecuteStatementResp, *cli_service.TGetOperationStatusResp, error) {
+	log := logger.WithContext(c.id, driverctx.CorrelationIdFromContext(ctx), "")
 	// 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.executeStatement(ctx, query, args)
 
 	if err != nil {
-		return nil, nil, err
+		return exStmtResp, nil, err
 	}
 	// hold on to the operation handle
 	opHandle := exStmtResp.OperationHandle
+	if opHandle != nil && opHandle.OperationId != nil {
+		log = logger.WithContext(c.id, driverctx.CorrelationIdFromContext(ctx), client.SprintGuid(opHandle.OperationId.GUID))
+	}
 
 	if exStmtResp.DirectResults != nil {
 		opStatus := exStmtResp.DirectResults.GetOperationStatus()
@@ -148,15 +185,13 @@ func (c *conn) runQuery(ctx context.Context, query string, args []driver.NamedVa
 		// bad
 		case cli_service.TOperationState_CANCELED_STATE, cli_service.TOperationState_CLOSED_STATE, cli_service.TOperationState_ERROR_STATE, cli_service.TOperationState_TIMEDOUT_STATE:
 			// do we need to close the operation in these cases?
-			logger.Error().Msg(opStatus.GetErrorMessage())
-			logger.Debug().Msgf("bad state: %s", opStatus.GetOperationState())
-
-			return exStmtResp, opStatus, fmt.Errorf(opStatus.GetDisplayMessage())
+			logBadQueryState(log, opStatus)
+			return exStmtResp, opStatus, errors.New(opStatus.GetDisplayMessage())
 		// live states
 		case cli_service.TOperationState_INITIALIZED_STATE, cli_service.TOperationState_PENDING_STATE, cli_service.TOperationState_RUNNING_STATE:
 			statusResp, err := c.pollOperation(ctx, opHandle)
 			if err != nil {
-				return nil, statusResp, err
+				return exStmtResp, statusResp, err
 			}
 			switch statusResp.GetOperationState() {
 			// terminal states
@@ -166,20 +201,17 @@ func (c *conn) runQuery(ctx context.Context, query string, args []driver.NamedVa
 				return exStmtResp, opStatus, nil
 			// bad
 			case cli_service.TOperationState_CANCELED_STATE, cli_service.TOperationState_CLOSED_STATE, cli_service.TOperationState_ERROR_STATE, cli_service.TOperationState_TIMEDOUT_STATE:
-				logger.Debug().Msgf("bad state: %s", statusResp.GetOperationState())
-				logger.Error().Msg(statusResp.GetErrorMessage())
-				return exStmtResp, opStatus, fmt.Errorf(statusResp.GetDisplayMessage())
+				logBadQueryState(log, statusResp)
+				return exStmtResp, opStatus, errors.New(statusResp.GetDisplayMessage())
 				// live states
 			default:
-				logger.Debug().Msgf("bad state: %s", statusResp.GetOperationState())
-				logger.Error().Msg(statusResp.GetErrorMessage())
-				return exStmtResp, opStatus, fmt.Errorf("invalid operation state. This should not have happened")
+				logBadQueryState(log, statusResp)
+				return exStmtResp, opStatus, errors.New("invalid operation state. This should not have happened")
 			}
 		// weird states
 		default:
-			logger.Debug().Msgf("bad state: %s", opStatus.GetOperationState())
-			logger.Error().Msg(opStatus.GetErrorMessage())
-			return exStmtResp, opStatus, fmt.Errorf("invalid operation state. This should not have happened")
+			logBadQueryState(log, opStatus)
+			return exStmtResp, opStatus, errors.New("invalid operation state. This should not have happened")
 		}
 
 	} else {
@@ -195,18 +227,21 @@ func (c *conn) runQuery(ctx context.Context, query string, args []driver.NamedVa
 			return exStmtResp, statusResp, nil
 		// bad
 		case cli_service.TOperationState_CANCELED_STATE, cli_service.TOperationState_CLOSED_STATE, cli_service.TOperationState_ERROR_STATE, cli_service.TOperationState_TIMEDOUT_STATE:
-			logger.Debug().Msgf("bad state: %s", statusResp.GetOperationState())
-			logger.Error().Msg(statusResp.GetErrorMessage())
-			return exStmtResp, statusResp, fmt.Errorf(statusResp.GetDisplayMessage())
+			logBadQueryState(log, statusResp)
+			return exStmtResp, statusResp, errors.New(statusResp.GetDisplayMessage())
 			// live states
 		default:
-			logger.Debug().Msgf("bad state: %s", statusResp.GetOperationState())
-			logger.Error().Msg(statusResp.GetErrorMessage())
-			return exStmtResp, statusResp, fmt.Errorf("invalid operation state. This should not have happened")
+			logBadQueryState(log, statusResp)
+			return exStmtResp, statusResp, errors.New("invalid operation state. This should not have happened")
 		}
 	}
 }
 
+func logBadQueryState(log *logger.DBSQLLogger, opStatus *cli_service.TGetOperationStatusResp) {
+	log.Error().Msgf("databricks: query state: %s", opStatus.GetOperationState())
+	log.Error().Msg(opStatus.GetErrorMessage())
+}
+
 func (c *conn) executeStatement(ctx context.Context, query string, args []driver.NamedValue) (*cli_service.TExecuteStatementResp, error) {
 	sentinel := sentinel.Sentinel{
 		OnDoneFn: func(statusResp any) (any, error) {
@@ -223,8 +258,9 @@ func (c *conn) executeStatement(ctx context.Context, query string, args []driver
 				// CanDecompressLZ4Result_: &f,
 				// CanDownloadResult_: &t,
 			}
+			ctx = driverctx.NewContextWithConnId(ctx, c.id)
 			resp, err := c.client.ExecuteStatement(ctx, &req)
-			return resp, err
+			return resp, wrapErr(err, "failed to execute statement")
 		},
 	}
 	_, res, err := sentinel.Watch(ctx, c.cfg.PollInterval, 0)
@@ -233,51 +269,61 @@ func (c *conn) executeStatement(ctx context.Context, query string, args []driver
 	}
 	exStmtResp, ok := res.(*cli_service.TExecuteStatementResp)
 	if !ok {
-		return nil, fmt.Errorf("databricks: invalid execute statement response")
+		return exStmtResp, errors.New("databricks: invalid execute statement response")
 	}
 	return exStmtResp, err
 }
 
 func (c *conn) pollOperation(ctx context.Context, opHandle *cli_service.TOperationHandle) (*cli_service.TGetOperationStatusResp, error) {
+	corrId := driverctx.CorrelationIdFromContext(ctx)
+	log := logger.WithContext(c.id, corrId, client.SprintGuid(opHandle.OperationId.GUID))
 	var statusResp *cli_service.TGetOperationStatusResp
+	ctx = driverctx.NewContextWithConnId(ctx, c.id)
+	newCtx := driverctx.NewContextWithCorrelationId(driverctx.NewContextWithConnId(context.Background(), c.id), corrId)
 	pollSentinel := sentinel.Sentinel{
 		OnDoneFn: func(statusResp any) (any, error) {
 			return statusResp, nil
 		},
 		StatusFn: func() (sentinel.Done, any, error) {
 			var err error
-			logger.Debug().Msg("databricks: polling status")
-			statusResp, err = c.client.GetOperationStatus(context.Background(), &cli_service.TGetOperationStatusReq{
+			log.Debug().Msg("databricks: polling status")
+			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())
+			}
 			return func() bool {
 				// which other states?
+				if err != nil {
+					return true
+				}
 				switch statusResp.GetOperationState() {
 				case cli_service.TOperationState_INITIALIZED_STATE, cli_service.TOperationState_PENDING_STATE, cli_service.TOperationState_RUNNING_STATE:
 					return false
 				default:
-					logger.Debug().Msg("databricks: polling done")
+					log.Debug().Msg("databricks: polling done")
 					return true
 				}
 			}, statusResp, err
 		},
 		OnCancelFn: func() (any, error) {
-			logger.Debug().Msgf("databricks: canceling operation %s", opHandle.OperationId)
-			ret, err := c.client.CancelOperation(context.Background(), &cli_service.TCancelOperationReq{
+			log.Debug().Msg("databricks: canceling query")
+			ret, err := c.client.CancelOperation(newCtx, &cli_service.TCancelOperationReq{
 				OperationHandle: opHandle,
 			})
 			return ret, err
 		},
 	}
 	_, resp, err := pollSentinel.Watch(ctx, c.cfg.PollInterval, 0)
 	if err != nil {
-		return nil, err
+		return nil, wrapErr(err, "failed to poll query state")
 	}
 	statusResp, ok := resp.(*cli_service.TGetOperationStatusResp)
 	if !ok {
-		return nil, fmt.Errorf("could not read operation status")
+		return nil, errors.New("could not read query status")
 	}
-	return statusResp, err
+	return statusResp, nil
 }
 
 var _ driver.Conn = (*conn)(nil)