Fix context loss in polling and connection close operations

Author: dgiagio

connection.go

@@ -55,7 +55,7 @@ func (c *conn) Close() error {
 
 	if err != nil {
 		log.Err(err).Msg("databricks: failed to close connection")
-		return dbsqlerrint.NewRequestError(ctx, dbsqlerr.ErrCloseConnection, err)
+		return dbsqlerrint.NewBadConnectionError(err)
 	}
 	return nil
 }
@@ -168,9 +168,7 @@ func (c *conn) QueryContext(ctx context.Context, query string, args []driver.Nam
 		return nil, dbsqlerrint.NewExecutionError(ctx, dbsqlerr.ErrQueryExecution, err, opStatusResp)
 	}
 
-	corrId := driverctx.CorrelationIdFromContext(ctx)
-	rows, err := rows.NewRows(c.id, corrId, exStmtResp.OperationHandle, c.client, c.cfg, exStmtResp.DirectResults)
-
+	rows, err := rows.NewRows(ctx, exStmtResp.OperationHandle, c.client, c.cfg, exStmtResp.DirectResults)
 	return rows, err
 
 }
@@ -367,7 +365,7 @@ func (c *conn) pollOperation(ctx context.Context, opHandle *cli_service.TOperati
 	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)
+	newCtx := driverctx.NewContextWithCorrelationId(ctx, corrId)
 	pollSentinel := sentinel.Sentinel{
 		OnDoneFn: func(statusResp any) (any, error) {
 			return statusResp, nil
@@ -566,7 +564,6 @@ func (c *conn) execStagingOperation(
 		return nil
 	}
 
-	corrId := driverctx.CorrelationIdFromContext(ctx)
 	var row driver.Rows
 	var err error
 
@@ -589,7 +586,7 @@ func (c *conn) execStagingOperation(
 	}
 
 	if len(driverctx.StagingPathsFromContext(ctx)) != 0 {
-		row, err = rows.NewRows(c.id, corrId, exStmtResp.OperationHandle, c.client, c.cfg, exStmtResp.DirectResults)
+		row, err = rows.NewRows(ctx, exStmtResp.OperationHandle, c.client, c.cfg, exStmtResp.DirectResults)
 		if err != nil {
 			return dbsqlerrint.NewDriverError(ctx, "error reading row.", err)
 		}

internal/rows/rows.go

@@ -67,22 +67,22 @@ var _ driver.RowsColumnTypeLength = (*rows)(nil)
 var _ dbsqlrows.Rows = (*rows)(nil)
 
 func NewRows(
-	connId string,
-	correlationId string,
+	ctx context.Context,
 	opHandle *cli_service.TOperationHandle,
 	client cli_service.TCLIService,
 	config *config.Config,
 	directResults *cli_service.TSparkDirectResults,
 ) (driver.Rows, dbsqlerr.DBError) {
 
+	connId := driverctx.ConnIdFromContext(ctx)
+	correlationId := driverctx.CorrelationIdFromContext(ctx)
+
 	var logger *dbsqllog.DBSQLLogger
-	var ctx context.Context
 	if opHandle != nil {
 		logger = dbsqllog.WithContext(connId, correlationId, dbsqlclient.SprintGuid(opHandle.OperationId.GUID))
-		ctx = driverctx.NewContextWithQueryId(driverctx.NewContextWithCorrelationId(driverctx.NewContextWithConnId(context.Background(), connId), correlationId), dbsqlclient.SprintGuid(opHandle.OperationId.GUID))
+		ctx = driverctx.NewContextWithQueryId(ctx, dbsqlclient.SprintGuid(opHandle.OperationId.GUID))
 	} else {
 		logger = dbsqllog.WithContext(connId, correlationId, "")
-		ctx = driverctx.NewContextWithCorrelationId(driverctx.NewContextWithConnId(context.Background(), connId), correlationId)
 	}
 
 	if client == nil {
@@ -140,13 +140,12 @@ func NewRows(
 	// the operations.
 	closedOnServer := directResults != nil && directResults.CloseOperation != nil
 	r.ResultPageIterator = rowscanner.NewResultPageIterator(
+		ctx,
 		d,
 		pageSize,
 		opHandle,
 		closedOnServer,
 		client,
-		connId,
-		correlationId,
 		r.logger(),
 	)
 
@@ -417,9 +416,8 @@ func (r *rows) getResultSetSchema() (*cli_service.TTableSchema, dbsqlerr.DBError
 		req := cli_service.TGetResultSetMetadataReq{
 			OperationHandle: r.opHandle,
 		}
-		ctx := driverctx.NewContextWithCorrelationId(driverctx.NewContextWithConnId(context.Background(), r.connId), r.correlationId)
 
-		resp, err2 := r.client.GetResultSetMetadata(ctx, &req)
+		resp, err2 := r.client.GetResultSetMetadata(r.ctx, &req)
 		if err2 != nil {
 			r.logger().Err(err2).Msg(err2.Error())
 			return nil, dbsqlerr_int.NewRequestError(r.ctx, errRowsMetadataFetchFailed, err)

internal/rows/rowscanner/resultPageIterator.go

@@ -45,32 +45,34 @@ func (d Direction) String() string {
 
 // Create a new result page iterator.
 func NewResultPageIterator(
+	ctx context.Context,
 	delimiter Delimiter,
 	maxPageSize int64,
 	opHandle *cli_service.TOperationHandle,
 	closedOnServer bool,
 	client cli_service.TCLIService,
-	connectionId string,
-	correlationId string,
 	logger *dbsqllog.DBSQLLogger,
 ) ResultPageIterator {
 
 	// delimiter and hasMoreRows are used to set up the point in the paginated
 	// result set that this iterator starts from.
 	return &resultPageIterator{
+		ctx:            ctx,
 		Delimiter:      delimiter,
 		isFinished:     closedOnServer,
 		maxPageSize:    maxPageSize,
 		opHandle:       opHandle,
 		closedOnServer: closedOnServer,
 		client:         client,
-		connectionId:   connectionId,
-		correlationId:  correlationId,
+		connectionId:   driverctx.ConnIdFromContext(ctx),
+		correlationId:  driverctx.CorrelationIdFromContext(ctx),
 		logger:         logger,
 	}
 }
 
 type resultPageIterator struct {
+	ctx context.Context
+
 	// Gives the parameters of the current result page
 	Delimiter
 
@@ -167,15 +169,14 @@ func (rpf *resultPageIterator) getNextPage() (*cli_service.TFetchResultsResp, er
 	nextPageStartRow := rpf.Start() + rpf.Count()
 
 	rpf.logger.Debug().Msgf("databricks: fetching result page for row %d", nextPageStartRow)
-	ctx := driverctx.NewContextWithCorrelationId(driverctx.NewContextWithConnId(context.Background(), rpf.connectionId), rpf.correlationId)
 
 	// Keep fetching in the appropriate direction until we have the expected page.
 	var fetchResult *cli_service.TFetchResultsResp
 	var b bool
 	for b = rpf.Contains(nextPageStartRow); !b; b = rpf.Contains(nextPageStartRow) {
 
 		direction := rpf.Direction(nextPageStartRow)
-		err := rpf.checkDirectionValid(ctx, direction)
+		err := rpf.checkDirectionValid(direction)
 		if err != nil {
 			return nil, err
 		}
@@ -190,10 +191,10 @@ func (rpf *resultPageIterator) getNextPage() (*cli_service.TFetchResultsResp, er
 			IncludeResultSetMetadata: &includeResultSetMetadata,
 		}
 
-		fetchResult, err = rpf.client.FetchResults(ctx, &req)
+		fetchResult, err = rpf.client.FetchResults(rpf.ctx, &req)
 		if err != nil {
 			rpf.logger.Err(err).Msg("databricks: Rows instance failed to retrieve results")
-			return nil, dbsqlerrint.NewRequestError(ctx, errRowsResultFetchFailed, err)
+			return nil, dbsqlerrint.NewRequestError(rpf.ctx, errRowsResultFetchFailed, err)
 		}
 
 		rpf.Delimiter = NewDelimiter(fetchResult.Results.StartRowOffset, CountRows(fetchResult.Results))
@@ -218,7 +219,7 @@ func (rpf *resultPageIterator) Close() (err error) {
 				OperationHandle: rpf.opHandle,
 			}
 
-			_, err = rpf.client.CloseOperation(context.Background(), &req)
+			_, err = rpf.client.CloseOperation(rpf.ctx, &req)
 			return err
 		}
 	}
@@ -283,11 +284,11 @@ func CountRows(rowSet *cli_service.TRowSet) int64 {
 }
 
 // Check if trying to fetch in the specified direction creates an error condition.
-func (rpf *resultPageIterator) checkDirectionValid(ctx context.Context, direction Direction) error {
+func (rpf *resultPageIterator) checkDirectionValid(direction Direction) error {
 	if direction == DirBack {
 		// can't fetch rows previous to the start
 		if rpf.Start() == 0 {
-			return dbsqlerrint.NewDriverError(ctx, ErrRowsFetchPriorToStart, nil)
+			return dbsqlerrint.NewDriverError(rpf.ctx, ErrRowsFetchPriorToStart, nil)
 		}
 	} else if direction == DirForward {
 		// can't fetch past the end of the query results
@@ -296,7 +297,7 @@ func (rpf *resultPageIterator) checkDirectionValid(ctx context.Context, directio
 		}
 	} else {
 		rpf.logger.Error().Msgf(errRowsUnandledFetchDirection(direction.String()))
-		return dbsqlerrint.NewDriverError(ctx, errRowsUnandledFetchDirection(direction.String()), nil)
+		return dbsqlerrint.NewDriverError(rpf.ctx, errRowsUnandledFetchDirection(direction.String()), nil)
 	}
 	return nil
 }