From 96bc53b8156b4a359ed457fc05f72a417dbf5aac Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Sun, 5 Jul 2026 15:39:24 +0000 Subject: [PATCH 1/6] refactor: extract Backend/Operation seam for SEA integration MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduce a pluggable execution-backend abstraction (internal/backend) and fold today's Thrift/HiveServer2 code into it as the sole implementation (internal/backend/thrift). conn no longer talks to the generated Thrift client directly; it holds one backend.Backend and drives session lifecycle and statement execution through the neutral Backend/Operation interfaces. No behavior change, no cgo, no new user-facing API. ## Why The driver is migrating to the SEA (Statement Execution API) protocol, which we plan to add via the Rust kernel behind a connect-time flag. Wiring a second protocol into conn as it stands would mean threading Thrift and SEA types through every database/sql method. This change draws the seam first — with Thrift as the only implementation — so the kernel backend can slot in later without touching conn, stmt, the rows layer, or any public API. It is the regression firewall the kernel PRs build on. ## What changes - internal/backend: neutral Backend / Operation interfaces + ExecRequest / Param request types. Only a backend implementation imports its protocol's client (the Thrift backend imports internal/cli_service); nothing else in the driver does for the execute/poll/close/session path. - internal/backend/thrift: the Thrift execute/poll/close/session logic moved here behind the interfaces; thriftOperation wraps the execute + status responses and exposes only neutral accessors. - conn / connector: drop the client + session fields for a single backend field; the 9 database/sql methods keep their signatures and delegate. The connector factory builds the Thrift backend (a future flag selects the kernel backend here). - parameters.go: split the parameter conversion at a neutral backend.Param boundary (inference stays in dbsql, wire-mapping moves to the backend) to avoid a dbsql<->backend import cycle. internal/querytags extracted for the same reason; dbsql.SerializeQueryTags forwards to it, unchanged. - internal/debuglog: a step-level debug tracer for the upcoming FFI work, emitted through the existing logger as a trace-level sibling and gated by the DBSQL_DEBUG_LOG env var (off by default). Every backend step is instrumented so a failing or slow step is visible in one ordered stream. ## Test plan - [x] go test ./... -count=1 passes (unit + moved Thrift-mechanism tests) - [x] Real-warehouse e2e (cancellation + CloudFetch) passes against a staging warehouse - [x] CGO_ENABLED=0 go build ./... still works (pure-Go default preserved) - [x] go vet, gofmt -s clean; internal/debuglog and internal/backend/thrift race-clean under -race - [x] Reviewed with isaac review and the code-review-squad; findings fixed Deferred to a follow-up (noted in the PR): the queryId enrichment on the execute path is now overwrite-always rather than fill-if-empty (affects only app-set query ids), and the parameter conversion makes an extra intermediate pass (inherent to the neutral seam). This pull request and its description were written with assistance from Claude Code. Signed-off-by: Mani Kaustubh Mathur --- connection.go | 386 +---- connection_test.go | 1437 +---------------- connector.go | 56 +- internal/backend/backend.go | 109 ++ internal/backend/thrift/backend.go | 441 +++++ internal/backend/thrift/backend_test.go | 1237 ++++++++++++++ .../thrift/debuglog_integration_test.go | 130 ++ internal/backend/thrift/export_test_helper.go | 38 + internal/backend/thrift/operation.go | 137 ++ internal/backend/thrift/operation_test.go | 175 ++ internal/backend/thrift/params.go | 41 + internal/backend/thrift/params_test.go | 69 + internal/debuglog/debuglog.go | 162 ++ internal/debuglog/debuglog_test.go | 322 ++++ internal/querytags/querytags.go | 40 + logger/logger.go | 25 +- parameter_test.go | 70 +- parameters.go | 45 +- query_tags.go | 23 +- statement_test.go | 7 +- 20 files changed, 3140 insertions(+), 1810 deletions(-) create mode 100644 internal/backend/backend.go create mode 100644 internal/backend/thrift/backend.go create mode 100644 internal/backend/thrift/backend_test.go create mode 100644 internal/backend/thrift/debuglog_integration_test.go create mode 100644 internal/backend/thrift/export_test_helper.go create mode 100644 internal/backend/thrift/operation.go create mode 100644 internal/backend/thrift/operation_test.go create mode 100644 internal/backend/thrift/params.go create mode 100644 internal/backend/thrift/params_test.go create mode 100644 internal/debuglog/debuglog.go create mode 100644 internal/debuglog/debuglog_test.go create mode 100644 internal/querytags/querytags.go diff --git a/connection.go b/connection.go index 3e6ac452..27c206e2 100644 --- a/connection.go +++ b/connection.go @@ -14,25 +14,22 @@ import ( "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/backend" "github.com/databricks/databricks-sql-go/internal/client" context2 "github.com/databricks/databricks-sql-go/internal/compat/context" "github.com/databricks/databricks-sql-go/internal/config" + "github.com/databricks/databricks-sql-go/internal/debuglog" dbsqlerrint "github.com/databricks/databricks-sql-go/internal/errors" "github.com/databricks/databricks-sql-go/internal/retry" "github.com/databricks/databricks-sql-go/internal/rows" - "github.com/databricks/databricks-sql-go/internal/sentinel" - "github.com/databricks/databricks-sql-go/internal/thrift_protocol" "github.com/databricks/databricks-sql-go/logger" "github.com/databricks/databricks-sql-go/telemetry" - "github.com/pkg/errors" ) type conn struct { id string cfg *config.Config - client cli_service.TCLIService - session *cli_service.TOpenSessionResp + backend backend.Backend telemetry *telemetry.Interceptor // Optional telemetry interceptor } @@ -55,9 +52,7 @@ func (c *conn) Close() error { // Time CloseSession so we can record DELETE_SESSION before flushing telemetry closeStart := time.Now() - _, err := c.client.CloseSession(ctx, &cli_service.TCloseSessionReq{ - SessionHandle: c.session.SessionHandle, - }) + err := c.backend.CloseSession(ctx) // Record DELETE_SESSION regardless of error (matches JDBC), then flush and release if c.telemetry != nil { @@ -111,7 +106,7 @@ func (c *conn) ResetSession(ctx context.Context) error { // 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 + return c.backend.SessionValid() } // ExecContext executes a query that doesn't return rows, such @@ -124,25 +119,37 @@ func (c *conn) ExecContext(ctx context.Context, query string, args []driver.Name log, _ := client.LoggerAndContext(ctx, nil) msg, start := logger.Track("ExecContext") defer log.Duration(msg, start) + defer debuglog.Track(ctx, "conn.ExecContext", "sql.len=%d args=%d", len(query), len(args))() corrId := driverctx.CorrelationIdFromContext(ctx) // Capture execution start time for telemetry before running the query executeStart := time.Now() - exStmtResp, opStatusResp, err := c.runQuery(ctx, query, args) - log, ctx = client.LoggerAndContext(ctx, exStmtResp) + op, err := c.runQuery(ctx, query, args) + // A nil op means execution never reached the backend (parameter conversion + // failed): no server statement to close or measure. Return the already-wrapped + // error directly. + if op == nil { + log.Err(err).Msgf("databricks: failed to execute query: query %s", query) + return nil, err + } + // Enrich the logger/context with the statement id. + if sid := op.StatementID(); sid != "" { + ctx = driverctx.NewContextWithQueryId(ctx, sid) + } + log = logger.WithContext(driverctx.ConnIdFromContext(ctx), corrId, driverctx.QueryIdFromContext(ctx)) // Telemetry: set up metric context BEFORE staging operation so that the // staging op's telemetryUpdate callback can attach tags to the metric context. var statementID string var closeOpErr error // Track CloseOperation errors for telemetry - if c.telemetry != nil && exStmtResp != nil && exStmtResp.OperationHandle != nil && exStmtResp.OperationHandle.OperationId != nil { - statementID = client.SprintGuid(exStmtResp.OperationHandle.OperationId.GUID) + if c.telemetry != nil && op.StatementID() != "" { + statementID = op.StatementID() ctx = c.telemetry.BeforeExecuteWithTime(ctx, c.id, statementID, executeStart) c.telemetry.AddTag(ctx, telemetry.TagOperationType, telemetry.OperationTypeExecuteStatement) } - stagingErr := c.execStagingOperation(exStmtResp, ctx) + stagingErr := c.execStagingOperation(op, ctx) if c.telemetry != nil && statementID != "" { defer func() { @@ -159,36 +166,33 @@ func (c *conn) ExecContext(ctx context.Context, query string, args []driver.Name }() } - if exStmtResp != nil && exStmtResp.OperationHandle != nil { - // since we have an operation handle we can close the operation if necessary - alreadyClosed := exStmtResp.DirectResults != nil && exStmtResp.DirectResults.CloseOperation != nil - newCtx := driverctx.NewContextWithCorrelationId(driverctx.NewContextWithConnId(context.Background(), c.id), corrId) - if !alreadyClosed && (opStatusResp == nil || opStatusResp.GetOperationState() != cli_service.TOperationState_CLOSED_STATE) { - closeOpStart := time.Now() - _, err1 := c.client.CloseOperation(newCtx, &cli_service.TCloseOperationReq{ - OperationHandle: exStmtResp.OperationHandle, - }) - if c.telemetry != nil { - c.telemetry.RecordOperation(ctx, c.id, statementID, telemetry.OperationTypeCloseStatement, time.Since(closeOpStart).Milliseconds(), err1) - } - if err1 != nil { - log.Err(err1).Msg("databricks: failed to close operation after executing statement") - closeOpErr = err1 // Capture for telemetry - } + // Close the server operation if one is still open. The backend decides + // whether a close RPC is actually needed and reports closed=true only when it + // issued one, so CLOSE_STATEMENT telemetry is recorded only for a real close. + newCtx := driverctx.NewContextWithCorrelationId(driverctx.NewContextWithConnId(context.Background(), c.id), corrId) + closeOpStart := time.Now() + closed, err1 := op.Close(newCtx) + if closed { + if c.telemetry != nil { + c.telemetry.RecordOperation(ctx, c.id, statementID, telemetry.OperationTypeCloseStatement, time.Since(closeOpStart).Milliseconds(), err1) + } + if err1 != nil { + log.Err(err1).Msg("databricks: failed to close operation after executing statement") + closeOpErr = err1 // Capture for telemetry } } if err != nil { log.Err(err).Msgf("databricks: failed to execute query: query %s", query) - return nil, dbsqlerrint.NewExecutionError(ctx, dbsqlerr.ErrQueryExecution, err, opStatusResp) + return nil, op.ExecutionError(ctx, err) } if stagingErr != nil { log.Err(stagingErr).Msgf("databricks: failed to execute query: query %s", query) - return nil, dbsqlerrint.NewExecutionError(ctx, dbsqlerr.ErrQueryExecution, stagingErr, opStatusResp) + return nil, op.ExecutionError(ctx, stagingErr) } - res := result{AffectedRows: opStatusResp.GetNumModifiedRows()} + res := result{AffectedRows: op.AffectedRows()} return &res, nil } @@ -243,21 +247,36 @@ func (c *conn) QueryContext(ctx context.Context, query string, args []driver.Nam ctx = driverctx.NewContextWithConnId(ctx, c.id) log, _ := client.LoggerAndContext(ctx, nil) msg, start := log.Track("QueryContext") + defer debuglog.Track(ctx, "conn.QueryContext", "sql.len=%d args=%d", len(query), len(args))() // first we try to get the results synchronously. // at any point in time that the context is done we must cancel and return + corrId := driverctx.CorrelationIdFromContext(ctx) + // Capture execution start time for telemetry before running the query executeStart := time.Now() - exStmtResp, opStatusResp, err := c.runQuery(ctx, query, args) - log, ctx = client.LoggerAndContext(ctx, exStmtResp) + op, err := c.runQuery(ctx, query, args) + // A nil op means execution never reached the backend (parameter conversion + // failed): no statement id, no telemetry to emit. Log and return the + // already-wrapped error directly. + if op == nil { + log.Err(err).Msg("databricks: failed to run query") // To log query we need to redact credentials + log.Duration(msg, start) + return nil, err + } + // Enrich the logger/context with the statement id. + if sid := op.StatementID(); sid != "" { + ctx = driverctx.NewContextWithQueryId(ctx, sid) + } + log = logger.WithContext(driverctx.ConnIdFromContext(ctx), corrId, driverctx.QueryIdFromContext(ctx)) defer log.Duration(msg, start) // Telemetry: set up metric context for the statement. // BeforeExecuteWithTime anchors startTime to before runQuery() ran. var statementID string - if c.telemetry != nil && exStmtResp != nil && exStmtResp.OperationHandle != nil && exStmtResp.OperationHandle.OperationId != nil { - statementID = client.SprintGuid(exStmtResp.OperationHandle.OperationId.GUID) + if c.telemetry != nil && op.StatementID() != "" { + statementID = op.StatementID() ctx = c.telemetry.BeforeExecuteWithTime(ctx, c.id, statementID, executeStart) c.telemetry.AddTag(ctx, telemetry.TagOperationType, telemetry.OperationTypeExecuteStatement) } @@ -270,7 +289,7 @@ func (c *conn) QueryContext(ctx context.Context, query string, args []driver.Nam c.telemetry.CompleteStatement(ctx, statementID, true) } log.Err(err).Msg("databricks: failed to run query") // To log query we need to redact credentials - return nil, dbsqlerrint.NewExecutionError(ctx, dbsqlerr.ErrQueryExecution, err, opStatusResp) + return nil, op.ExecutionError(ctx, err) } // Success path: freeze execute latency NOW (before row iteration inflates time.Since). @@ -358,7 +377,7 @@ func (c *conn) QueryContext(ctx context.Context, query string, args []driver.Nam } } - rows, err := rows.NewRows(ctx, exStmtResp.OperationHandle, c.client, c.cfg, exStmtResp.DirectResults, &rows.TelemetryCallbacks{ + rows, err := op.Results(ctx, &rows.TelemetryCallbacks{ OnChunkFetched: telemetryUpdate, OnClose: closeCallback, OnCloudFetchFile: cloudFetchCallback, @@ -367,262 +386,18 @@ 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) { - // 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) - var log *logger.DBSQLLogger - log, ctx = client.LoggerAndContext(ctx, exStmtResp) - - if err != nil { - return exStmtResp, nil, err - } - - opHandle := exStmtResp.OperationHandle - - if exStmtResp.DirectResults != nil { - opStatus := exStmtResp.DirectResults.GetOperationStatus() - - switch opStatus.GetOperationState() { - // terminal states - // good - case cli_service.TOperationState_FINISHED_STATE: - 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: - logBadQueryState(log, opStatus) - return exStmtResp, opStatus, unexpectedOperationState(opStatus) - // 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 exStmtResp, statusResp, err - } - switch statusResp.GetOperationState() { - // terminal states - // good - case cli_service.TOperationState_FINISHED_STATE: - 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: - logBadQueryState(log, statusResp) - return exStmtResp, statusResp, unexpectedOperationState(statusResp) - // live states - default: - logBadQueryState(log, statusResp) - return exStmtResp, statusResp, invalidOperationState(ctx, statusResp) - } - // weird states - default: - logBadQueryState(log, opStatus) - return exStmtResp, opStatus, invalidOperationState(ctx, opStatus) - } - - } else { - statusResp, err := c.pollOperation(ctx, opHandle) - if err != nil { - return exStmtResp, statusResp, err - } - switch statusResp.GetOperationState() { - // terminal states - // good - case cli_service.TOperationState_FINISHED_STATE: - 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: - logBadQueryState(log, statusResp) - return exStmtResp, statusResp, unexpectedOperationState(statusResp) - // live states - default: - logBadQueryState(log, statusResp) - return exStmtResp, statusResp, invalidOperationState(ctx, statusResp) - } - } -} - -func logBadQueryState(log *logger.DBSQLLogger, opStatus *cli_service.TGetOperationStatusResp) { - log.Error().Msgf("databricks: query state: %s", opStatus.GetOperationState()) - log.Error().Msg(opStatus.GetDisplayMessage()) - log.Debug().Msg(opStatus.GetDiagnosticInfo()) -} - -func unexpectedOperationState(opStatus *cli_service.TGetOperationStatusResp) error { - return errors.WithMessage(errors.New(opStatus.GetDisplayMessage()), dbsqlerr.ErrUnexpectedOperationState(opStatus.GetOperationState().String())) -} - -func invalidOperationState(ctx context.Context, opStatus *cli_service.TGetOperationStatusResp) error { - return dbsqlerrint.NewDriverError(ctx, dbsqlerr.ErrInvalidOperationState(opStatus.GetOperationState().String()), nil) -} - -func (c *conn) executeStatement(ctx context.Context, query string, args []driver.NamedValue) (*cli_service.TExecuteStatementResp, error) { - ctx = driverctx.NewContextWithConnId(ctx, c.id) - - parameters, err := convertNamedValuesToSparkParams(args) - if err != nil { - return nil, err - } - - req := cli_service.TExecuteStatementReq{ - SessionHandle: c.session.SessionHandle, - Statement: query, - RunAsync: true, - QueryTimeout: int64(c.cfg.QueryTimeout / time.Second), - } - - // Check protocol version for feature support - serverProtocolVersion := c.session.ServerProtocolVersion - - // Add direct results if supported - if thrift_protocol.SupportsDirectResults(serverProtocolVersion) { - req.GetDirectResults = &cli_service.TSparkGetDirectResults{ - MaxRows: int64(c.cfg.MaxRows), - } - } - - // Add LZ4 compression if supported and enabled - if thrift_protocol.SupportsLz4Compression(serverProtocolVersion) && c.cfg.UseLz4Compression { - req.CanDecompressLZ4Result_ = &c.cfg.UseLz4Compression - } - - // Add cloud fetch if supported and enabled - if thrift_protocol.SupportsCloudFetch(serverProtocolVersion) && c.cfg.UseCloudFetch { - req.CanDownloadResult_ = &c.cfg.UseCloudFetch - } - - // Add Arrow support if supported and enabled - if thrift_protocol.SupportsArrow(serverProtocolVersion) && c.cfg.UseArrowBatches { - req.CanReadArrowResult_ = &c.cfg.UseArrowBatches - req.UseArrowNativeTypes = &cli_service.TSparkArrowTypes{ - DecimalAsArrow: &c.cfg.UseArrowNativeDecimal, - TimestampAsArrow: &c.cfg.UseArrowNativeTimestamp, - ComplexTypesAsArrow: &c.cfg.UseArrowNativeComplexTypes, - IntervalTypesAsArrow: &c.cfg.UseArrowNativeIntervalTypes, - } - } - - // Add parameters if supported and provided - if thrift_protocol.SupportsParameterizedQueries(serverProtocolVersion) && len(parameters) > 0 { - req.Parameters = parameters - } - - // Add per-statement query tags if provided via context - if queryTags := driverctx.QueryTagsFromContext(ctx); len(queryTags) > 0 { - serialized := SerializeQueryTags(queryTags) - if serialized != "" { - if req.ConfOverlay == nil { - req.ConfOverlay = make(map[string]string) - } - req.ConfOverlay["query_tags"] = serialized - } - } - - resp, err := c.client.ExecuteStatement(ctx, &req) - var log *logger.DBSQLLogger - log, ctx = client.LoggerAndContext(ctx, resp) - - var shouldCancel = func(resp *cli_service.TExecuteStatementResp) bool { - if resp == nil { - return false - } - hasHandle := resp.OperationHandle != nil - isOpen := resp.DirectResults == nil || resp.DirectResults.CloseOperation == nil - return hasHandle && isOpen - } - - select { - default: - // Non-blocking check: continue if context not done - case <-ctx.Done(): - newCtx := driverctx.NewContextFromBackground(ctx) - // in case context is done, we need to cancel the operation if necessary - if err == nil && shouldCancel(resp) { - log.Debug().Msg("databricks: canceling query") - _, err1 := c.client.CancelOperation(newCtx, &cli_service.TCancelOperationReq{ - OperationHandle: resp.GetOperationHandle(), - }) - - if err1 != nil { - log.Err(err1).Msgf("databricks: cancel failed") - } else { - log.Debug().Msgf("databricks: cancel success") - } - } else { - log.Debug().Msg("databricks: query did not need cancellation") - } - return nil, ctx.Err() - } - - return resp, 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 := context2.WithoutCancel(ctx) - pollSentinel := sentinel.Sentinel{ - OnDoneFn: func(statusResp any) (any, error) { - return statusResp, nil - }, - StatusFn: func() (sentinel.Done, any, error) { - var err error - 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 { - 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: - log.Debug().Msg("databricks: polling done") - return true - } - }, statusResp, err - }, - OnCancelFn: func() (any, error) { - log.Debug().Msg("databricks: sentinel canceling query") - ret, err := c.client.CancelOperation(newCtx, &cli_service.TCancelOperationReq{ - OperationHandle: opHandle, - }) - return ret, err - }, - } - status, resp, err := pollSentinel.Watch(ctx, c.cfg.PollInterval, 0) +// runQuery converts the caller's parameters and executes the statement through +// the backend. It returns a nil Operation only when execution never reached the +// backend — parameter conversion failed before any server statement existed — in +// which case the wrapped error is returned directly. Once the backend is called, +// backend.Execute guarantees a non-nil Operation, so callers guard only for the +// pre-backend nil. +func (c *conn) runQuery(ctx context.Context, query string, args []driver.NamedValue) (backend.Operation, error) { + params, err := convertNamedValuesToParams(args) if err != nil { - log.Err(err).Msg("error polling operation status") - if status == sentinel.WatchTimeout { - err = dbsqlerrint.NewRequestError(ctx, dbsqlerr.ErrSentinelTimeout, err) - } - return nil, err - } - - statusResp, ok := resp.(*cli_service.TGetOperationStatusResp) - if !ok { - return nil, dbsqlerrint.NewDriverError(ctx, dbsqlerr.ErrReadQueryStatus, nil) + return nil, dbsqlerrint.NewExecutionError(ctx, dbsqlerr.ErrQueryExecution, err, nil) } - return statusResp, nil + return c.backend.Execute(ctx, backend.ExecRequest{Query: query, Params: params}) } func (c *conn) CheckNamedValue(nv *driver.NamedValue) error { @@ -916,28 +691,17 @@ func localPathIsAllowed(stagingAllowedLocalPaths []string, localFile string) boo } func (c *conn) execStagingOperation( - exStmtResp *cli_service.TExecuteStatementResp, + op backend.Operation, ctx context.Context) dbsqlerr.DBError { - if exStmtResp == nil || exStmtResp.OperationHandle == nil { - return nil - } + defer debuglog.Track(ctx, "conn.execStagingOperation", "stmt=%s", op.StatementID())() var row driver.Rows var err error - var isStagingOperation bool - if exStmtResp.DirectResults != nil && exStmtResp.DirectResults.ResultSetMetadata != nil { - isStagingOperation = exStmtResp.DirectResults.ResultSetMetadata.IsStagingOperation != nil && *exStmtResp.DirectResults.ResultSetMetadata.IsStagingOperation - } else { - req := cli_service.TGetResultSetMetadataReq{ - OperationHandle: exStmtResp.OperationHandle, - } - resp, err := c.client.GetResultSetMetadata(ctx, &req) - if err != nil { - return dbsqlerrint.NewDriverError(ctx, "error performing staging operation", err) - } - isStagingOperation = resp.IsStagingOperation != nil && *resp.IsStagingOperation + isStagingOperation, stagingErr := op.IsStaging(ctx) + if stagingErr != nil { + return dbsqlerrint.NewDriverError(ctx, "error performing staging operation", stagingErr) } if !isStagingOperation { @@ -952,7 +716,7 @@ func (c *conn) execStagingOperation( c.telemetry.AddTag(ctx, telemetry.TagBytesDownloaded, bytesDownloaded) } } - row, err = rows.NewRows(ctx, exStmtResp.OperationHandle, c.client, c.cfg, exStmtResp.DirectResults, &rows.TelemetryCallbacks{ + row, err = op.Results(ctx, &rows.TelemetryCallbacks{ OnChunkFetched: telemetryUpdate, }) if err != nil { diff --git a/connection_test.go b/connection_test.go index 7fd48ce9..77d1d5bb 100644 --- a/connection_test.go +++ b/connection_test.go @@ -17,1389 +17,14 @@ import ( "github.com/apache/thrift/lib/go/thrift" "github.com/pkg/errors" - "github.com/databricks/databricks-sql-go/driverctx" dbsqlerr "github.com/databricks/databricks-sql-go/errors" + thriftbackend "github.com/databricks/databricks-sql-go/internal/backend/thrift" "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/thrift_protocol" "github.com/stretchr/testify/assert" ) -func TestConn_executeStatement(t *testing.T) { - t.Parallel() - t.Run("executeStatement should err when client.ExecuteStatement fails", func(t *testing.T) { - var executeStatementCount int - executeStatement := func(ctx context.Context, req *cli_service.TExecuteStatementReq) (r *cli_service.TExecuteStatementResp, err error) { - executeStatementCount++ - return nil, fmt.Errorf("error") - } - testClient := &client.TestClient{ - FnExecuteStatement: executeStatement, - } - testConn := &conn{ - session: getTestSession(), - client: testClient, - cfg: config.WithDefaults(), - } - _, err := testConn.executeStatement(context.Background(), "select 1", []driver.NamedValue{}) - assert.Error(t, err) - assert.Equal(t, 1, executeStatementCount) - }) - - t.Run("executeStatement should return TExecuteStatementResp on success", func(t *testing.T) { - var executeStatementCount int - executeStatement := func(ctx context.Context, req *cli_service.TExecuteStatementReq) (r *cli_service.TExecuteStatementResp, err error) { - executeStatementCount++ - executeStatementResp := &cli_service.TExecuteStatementResp{ - Status: &cli_service.TStatus{ - StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, - }, - OperationHandle: &cli_service.TOperationHandle{ - OperationId: &cli_service.THandleIdentifier{ - GUID: []byte{1, 2, 3, 4, 2, 23, 4, 2, 3, 1, 2, 3, 4, 4, 223, 34, 54}, - Secret: []byte("b"), - }, - }, - DirectResults: &cli_service.TSparkDirectResults{ - OperationStatus: &cli_service.TGetOperationStatusResp{ - Status: &cli_service.TStatus{ - StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, - }, - OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_ERROR_STATE), - ErrorMessage: strPtr("error message"), - DisplayMessage: strPtr("display message"), - }, - ResultSetMetadata: &cli_service.TGetResultSetMetadataResp{ - Status: &cli_service.TStatus{ - StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, - }, - }, - ResultSet: &cli_service.TFetchResultsResp{ - Status: &cli_service.TStatus{ - StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, - }, - }, - }, - } - return executeStatementResp, nil - } - testClient := &client.TestClient{ - FnExecuteStatement: executeStatement, - } - testConn := &conn{ - session: getTestSession(), - client: testClient, - cfg: config.WithDefaults(), - } - _, err := testConn.executeStatement(context.Background(), "select 1", []driver.NamedValue{}) - - assert.NoError(t, err) - assert.Equal(t, 1, executeStatementCount) - }) - - t.Run("ExecStatement should close operation on success", func(t *testing.T) { - var executeStatementCount, closeOperationCount int - executeStatementResp := &cli_service.TExecuteStatementResp{ - Status: &cli_service.TStatus{ - StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, - }, - OperationHandle: &cli_service.TOperationHandle{ - OperationId: &cli_service.THandleIdentifier{ - GUID: []byte{1, 2, 3, 4, 2, 23, 4, 2, 3, 1, 2, 3, 4, 223, 34, 54}, - Secret: []byte("b"), - }, - }, - DirectResults: &cli_service.TSparkDirectResults{ - OperationStatus: &cli_service.TGetOperationStatusResp{ - Status: &cli_service.TStatus{ - StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, - }, - OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_ERROR_STATE), - ErrorMessage: strPtr("error message"), - DisplayMessage: strPtr("display message"), - }, - ResultSetMetadata: &cli_service.TGetResultSetMetadataResp{ - Status: &cli_service.TStatus{ - StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, - }, - }, - ResultSet: &cli_service.TFetchResultsResp{ - Status: &cli_service.TStatus{ - StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, - }, - }, - }, - } - - getResultSetMetadata := func(ctx context.Context, req *cli_service.TGetResultSetMetadataReq) (_r *cli_service.TGetResultSetMetadataResp, _err error) { - var b = false - return &cli_service.TGetResultSetMetadataResp{IsStagingOperation: &b}, nil - } - - testClient := &client.TestClient{ - FnExecuteStatement: func(ctx context.Context, req *cli_service.TExecuteStatementReq) (r *cli_service.TExecuteStatementResp, err error) { - executeStatementCount++ - return executeStatementResp, nil - }, - FnCloseOperation: func(ctx context.Context, req *cli_service.TCloseOperationReq) (_r *cli_service.TCloseOperationResp, _err error) { - closeOperationCount++ - return &cli_service.TCloseOperationResp{}, nil - }, - FnGetResultSetMetadata: getResultSetMetadata, - } - testConn := &conn{ - session: getTestSession(), - client: testClient, - cfg: config.WithDefaults(), - } - - type opStateTest struct { - state cli_service.TOperationState - err string - closeOperationCount int - } - - // test behaviour with all terminal operation states - operationStateTests := []opStateTest{ - {state: cli_service.TOperationState_ERROR_STATE, err: "error state", closeOperationCount: 1}, - {state: cli_service.TOperationState_FINISHED_STATE, err: "", closeOperationCount: 1}, - {state: cli_service.TOperationState_CANCELED_STATE, err: "cancelled state", closeOperationCount: 1}, - {state: cli_service.TOperationState_CLOSED_STATE, err: "closed state", closeOperationCount: 0}, - {state: cli_service.TOperationState_TIMEDOUT_STATE, err: "timeout state", closeOperationCount: 1}, - } - - for _, opTest := range operationStateTests { - closeOperationCount = 0 - executeStatementCount = 0 - executeStatementResp.DirectResults.OperationStatus.OperationState = &opTest.state //nolint:gosec // G601: pointer is used only within this loop iteration - executeStatementResp.DirectResults.OperationStatus.DisplayMessage = &opTest.err //nolint:gosec // G601: pointer is used only within this loop iteration - _, err := testConn.ExecContext(context.Background(), "select 1", []driver.NamedValue{}) - if opTest.err == "" { - assert.NoError(t, err) - } else { - assert.EqualError(t, err, "databricks: execution error: failed to execute query: unexpected operation state "+opTest.state.String()+": "+opTest.err) - } - assert.Equal(t, 1, executeStatementCount) - assert.Equal(t, opTest.closeOperationCount, closeOperationCount) - } - - // if the execute statement response contains direct results with a non-nil CloseOperation member - // we shouldn't call close - closeOperationCount = 0 - executeStatementCount = 0 - executeStatementResp.DirectResults.CloseOperation = &cli_service.TCloseOperationResp{} - finished := cli_service.TOperationState_FINISHED_STATE - executeStatementResp.DirectResults.OperationStatus.OperationState = &finished - _, err := testConn.ExecContext(context.Background(), "select 1", []driver.NamedValue{}) - assert.NoError(t, err) - assert.Equal(t, 1, executeStatementCount) - assert.Equal(t, 0, closeOperationCount) - }) - - t.Run("executeStatement should not call cancel if not needed", func(t *testing.T) { - var executeStatementCount int - var cancelOperationCount int - var cancel context.CancelFunc - executeStatement := func(ctx context.Context, req *cli_service.TExecuteStatementReq) (r *cli_service.TExecuteStatementResp, err error) { - executeStatementCount++ - cancel() - executeStatementResp := &cli_service.TExecuteStatementResp{ - Status: &cli_service.TStatus{ - StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, - }, - OperationHandle: &cli_service.TOperationHandle{ - OperationId: &cli_service.THandleIdentifier{ - GUID: []byte{1, 2, 3, 4, 2, 23, 4, 2, 3, 1, 2, 3, 4, 4, 223, 34, 54}, - Secret: []byte("b"), - }, - }, - DirectResults: &cli_service.TSparkDirectResults{ - OperationStatus: &cli_service.TGetOperationStatusResp{ - Status: &cli_service.TStatus{ - StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, - }, - OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_FINISHED_STATE), - ErrorMessage: strPtr("error message"), - DisplayMessage: strPtr("display message"), - }, - ResultSetMetadata: &cli_service.TGetResultSetMetadataResp{ - Status: &cli_service.TStatus{ - StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, - }, - }, - ResultSet: &cli_service.TFetchResultsResp{ - Status: &cli_service.TStatus{ - StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, - }, - }, - CloseOperation: &cli_service.TCloseOperationResp{ - Status: &cli_service.TStatus{ - StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, - }, - }, - }, - } - return executeStatementResp, nil - } - cancelOperation := func(ctx context.Context, req *cli_service.TCancelOperationReq) (r *cli_service.TCancelOperationResp, err error) { - cancelOperationCount++ - cancelOperationResp := &cli_service.TCancelOperationResp{ - Status: &cli_service.TStatus{ - StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, - }, - } - return cancelOperationResp, nil - } - testClient := &client.TestClient{ - FnExecuteStatement: executeStatement, - FnCancelOperation: cancelOperation, - } - testConn := &conn{ - session: getTestSession(), - client: testClient, - cfg: config.WithDefaults(), - } - - ctx := context.Background() - ctx, cancel = context.WithCancel(ctx) - defer cancel() - _, err := testConn.executeStatement(ctx, "select 1", []driver.NamedValue{}) - - assert.Error(t, err) - assert.Equal(t, 1, executeStatementCount) - assert.Equal(t, 0, cancelOperationCount) - }) - t.Run("executeStatement should call cancel if needed", func(t *testing.T) { - var executeStatementCount int - var cancelOperationCount int - var cancel context.CancelFunc - executeStatement := func(ctx context.Context, req *cli_service.TExecuteStatementReq) (r *cli_service.TExecuteStatementResp, err error) { - executeStatementCount++ - cancel() - executeStatementResp := &cli_service.TExecuteStatementResp{ - Status: &cli_service.TStatus{ - StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, - }, - OperationHandle: &cli_service.TOperationHandle{ - OperationId: &cli_service.THandleIdentifier{ - GUID: []byte{1, 2, 3, 4, 2, 23, 4, 2, 3, 1, 2, 3, 4, 4, 223, 34, 54}, - Secret: []byte("b"), - }, - }, - DirectResults: &cli_service.TSparkDirectResults{ - OperationStatus: &cli_service.TGetOperationStatusResp{ - Status: &cli_service.TStatus{ - StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, - }, - OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_FINISHED_STATE), - ErrorMessage: strPtr("error message"), - DisplayMessage: strPtr("display message"), - }, - ResultSetMetadata: &cli_service.TGetResultSetMetadataResp{ - Status: &cli_service.TStatus{ - StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, - }, - }, - ResultSet: &cli_service.TFetchResultsResp{ - Status: &cli_service.TStatus{ - StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, - }, - }, - }, - } - return executeStatementResp, nil - } - cancelOperation := func(ctx context.Context, req *cli_service.TCancelOperationReq) (r *cli_service.TCancelOperationResp, err error) { - cancelOperationCount++ - cancelOperationResp := &cli_service.TCancelOperationResp{ - Status: &cli_service.TStatus{ - StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, - }, - } - return cancelOperationResp, nil - } - testClient := &client.TestClient{ - FnExecuteStatement: executeStatement, - FnCancelOperation: cancelOperation, - } - testConn := &conn{ - session: getTestSession(), - client: testClient, - cfg: config.WithDefaults(), - } - ctx := context.Background() - ctx, cancel = context.WithCancel(ctx) - defer cancel() - _, err := testConn.executeStatement(ctx, "select 1", []driver.NamedValue{}) - - assert.Error(t, err) - assert.Equal(t, 1, executeStatementCount) - assert.Equal(t, 1, cancelOperationCount) - }) - -} - -func TestConn_executeStatement_ProtocolFeatures(t *testing.T) { - t.Parallel() - - protocols := []cli_service.TProtocolVersion{ - cli_service.TProtocolVersion_SPARK_CLI_SERVICE_PROTOCOL_V1, - cli_service.TProtocolVersion_SPARK_CLI_SERVICE_PROTOCOL_V2, - cli_service.TProtocolVersion_SPARK_CLI_SERVICE_PROTOCOL_V3, - cli_service.TProtocolVersion_SPARK_CLI_SERVICE_PROTOCOL_V4, - cli_service.TProtocolVersion_SPARK_CLI_SERVICE_PROTOCOL_V5, - cli_service.TProtocolVersion_SPARK_CLI_SERVICE_PROTOCOL_V6, - cli_service.TProtocolVersion_SPARK_CLI_SERVICE_PROTOCOL_V7, - cli_service.TProtocolVersion_SPARK_CLI_SERVICE_PROTOCOL_V8, - } - - testCases := []struct { - cfg *config.Config - supportsDirectResults func(version cli_service.TProtocolVersion) bool - supportsLz4Compression func(version cli_service.TProtocolVersion) bool - supportsCloudFetch func(version cli_service.TProtocolVersion) bool - supportsArrow func(version cli_service.TProtocolVersion) bool - supportsParameterizedQueries func(version cli_service.TProtocolVersion) bool - hasParameters bool - }{ - { - cfg: func() *config.Config { - cfg := config.WithDefaults() - cfg.UseLz4Compression = true - cfg.UseCloudFetch = true - cfg.UseArrowBatches = true - cfg.UseArrowNativeDecimal = true - cfg.UseArrowNativeTimestamp = true - cfg.UseArrowNativeComplexTypes = true - cfg.UseArrowNativeIntervalTypes = true - return cfg - }(), - supportsDirectResults: thrift_protocol.SupportsDirectResults, - supportsLz4Compression: thrift_protocol.SupportsLz4Compression, - supportsCloudFetch: thrift_protocol.SupportsCloudFetch, - supportsArrow: thrift_protocol.SupportsArrow, - supportsParameterizedQueries: thrift_protocol.SupportsParameterizedQueries, - hasParameters: true, - }, - { - cfg: func() *config.Config { - cfg := config.WithDefaults() - cfg.UseLz4Compression = false - cfg.UseCloudFetch = false - cfg.UseArrowBatches = false - return cfg - }(), - supportsDirectResults: thrift_protocol.SupportsDirectResults, - supportsLz4Compression: thrift_protocol.SupportsLz4Compression, - supportsCloudFetch: thrift_protocol.SupportsCloudFetch, - supportsArrow: thrift_protocol.SupportsArrow, - supportsParameterizedQueries: thrift_protocol.SupportsParameterizedQueries, - hasParameters: false, - }, - } - - for _, tc := range testCases { - for _, version := range protocols { - t.Run(fmt.Sprintf("protocol_v%d_withParams_%v", version, tc.hasParameters), func(t *testing.T) { - var capturedReq *cli_service.TExecuteStatementReq - executeStatement := func(ctx context.Context, req *cli_service.TExecuteStatementReq) (r *cli_service.TExecuteStatementResp, err error) { - capturedReq = req - executeStatementResp := &cli_service.TExecuteStatementResp{ - Status: &cli_service.TStatus{ - StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, - }, - OperationHandle: &cli_service.TOperationHandle{ - OperationId: &cli_service.THandleIdentifier{ - GUID: []byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}, - Secret: []byte("secret"), - }, - }, - DirectResults: &cli_service.TSparkDirectResults{ - OperationStatus: &cli_service.TGetOperationStatusResp{ - Status: &cli_service.TStatus{ - StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, - }, - OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_FINISHED_STATE), - }, - }, - } - return executeStatementResp, nil - } - - session := getTestSession() - session.ServerProtocolVersion = version - - testClient := &client.TestClient{ - FnExecuteStatement: executeStatement, - } - - testConn := &conn{ - session: session, - client: testClient, - cfg: tc.cfg, - } - - var args []driver.NamedValue - if tc.hasParameters { - args = []driver.NamedValue{ - {Name: "param1", Value: "value1"}, - } - } - - _, err := testConn.executeStatement(context.Background(), "SELECT 1", args) - assert.NoError(t, err) - - // Verify direct results - hasDirectResults := tc.supportsDirectResults(version) - assert.Equal(t, hasDirectResults, capturedReq.GetDirectResults != nil, "Direct results should be enabled if protocol supports it") - - // Verify LZ4 compression - shouldHaveLz4 := tc.supportsLz4Compression(version) && tc.cfg.UseLz4Compression - if shouldHaveLz4 { - assert.NotNil(t, capturedReq.CanDecompressLZ4Result_) - assert.True(t, *capturedReq.CanDecompressLZ4Result_) - } else { - assert.Nil(t, capturedReq.CanDecompressLZ4Result_) - } - - // Verify cloud fetch - shouldHaveCloudFetch := tc.supportsCloudFetch(version) && tc.cfg.UseCloudFetch - if shouldHaveCloudFetch { - assert.NotNil(t, capturedReq.CanDownloadResult_) - assert.True(t, *capturedReq.CanDownloadResult_) - } else { - assert.Nil(t, capturedReq.CanDownloadResult_) - } - - // Verify Arrow support - shouldHaveArrow := tc.supportsArrow(version) && tc.cfg.UseArrowBatches - if shouldHaveArrow { - assert.NotNil(t, capturedReq.CanReadArrowResult_) - assert.True(t, *capturedReq.CanReadArrowResult_) - assert.NotNil(t, capturedReq.UseArrowNativeTypes) - assert.Equal(t, tc.cfg.UseArrowNativeDecimal, *capturedReq.UseArrowNativeTypes.DecimalAsArrow) - assert.Equal(t, tc.cfg.UseArrowNativeTimestamp, *capturedReq.UseArrowNativeTypes.TimestampAsArrow) - assert.Equal(t, tc.cfg.UseArrowNativeComplexTypes, *capturedReq.UseArrowNativeTypes.ComplexTypesAsArrow) - assert.Equal(t, tc.cfg.UseArrowNativeIntervalTypes, *capturedReq.UseArrowNativeTypes.IntervalTypesAsArrow) - } else { - assert.Nil(t, capturedReq.CanReadArrowResult_) - assert.Nil(t, capturedReq.UseArrowNativeTypes) - } - - // Verify parameters - shouldHaveParams := tc.supportsParameterizedQueries(version) && tc.hasParameters - if shouldHaveParams { - assert.NotNil(t, capturedReq.Parameters) - assert.Len(t, capturedReq.Parameters, 1) - } else if tc.hasParameters { - // Even if we have parameters but protocol doesn't support it, we shouldn't set them - assert.Nil(t, capturedReq.Parameters) - } - }) - } - } -} - -func TestConn_executeStatement_QueryTags(t *testing.T) { - t.Parallel() - - makeTestConn := func(captureReq *(*cli_service.TExecuteStatementReq)) *conn { - executeStatement := func(ctx context.Context, req *cli_service.TExecuteStatementReq) (r *cli_service.TExecuteStatementResp, err error) { - *captureReq = req - return &cli_service.TExecuteStatementResp{ - Status: &cli_service.TStatus{ - StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, - }, - OperationHandle: &cli_service.TOperationHandle{ - OperationId: &cli_service.THandleIdentifier{ - GUID: []byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}, - Secret: []byte("secret"), - }, - }, - DirectResults: &cli_service.TSparkDirectResults{ - OperationStatus: &cli_service.TGetOperationStatusResp{ - Status: &cli_service.TStatus{ - StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, - }, - OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_FINISHED_STATE), - }, - }, - }, nil - } - - return &conn{ - session: getTestSession(), - client: &client.TestClient{ - FnExecuteStatement: executeStatement, - }, - cfg: config.WithDefaults(), - } - } - - t.Run("query tags from context are set in ConfOverlay", func(t *testing.T) { - var capturedReq *cli_service.TExecuteStatementReq - testConn := makeTestConn(&capturedReq) - - ctx := driverctx.NewContextWithQueryTags(context.Background(), map[string]string{ - "team": "engineering", - "app": "etl", - }) - - _, err := testConn.executeStatement(ctx, "SELECT 1", nil) - assert.NoError(t, err) - assert.NotNil(t, capturedReq.ConfOverlay) - // Map iteration is non-deterministic, so check both possible orderings - queryTags := capturedReq.ConfOverlay["query_tags"] - assert.True(t, - queryTags == "team:engineering,app:etl" || queryTags == "app:etl,team:engineering", - "unexpected query_tags value: %s", queryTags) - }) - - t.Run("no query tags in context means no ConfOverlay", func(t *testing.T) { - var capturedReq *cli_service.TExecuteStatementReq - testConn := makeTestConn(&capturedReq) - - _, err := testConn.executeStatement(context.Background(), "SELECT 1", nil) - assert.NoError(t, err) - assert.Nil(t, capturedReq.ConfOverlay) - }) - - t.Run("empty query tags map means no ConfOverlay", func(t *testing.T) { - var capturedReq *cli_service.TExecuteStatementReq - testConn := makeTestConn(&capturedReq) - - ctx := driverctx.NewContextWithQueryTags(context.Background(), map[string]string{}) - - _, err := testConn.executeStatement(ctx, "SELECT 1", nil) - assert.NoError(t, err) - assert.Nil(t, capturedReq.ConfOverlay) - }) - - t.Run("single query tag", func(t *testing.T) { - var capturedReq *cli_service.TExecuteStatementReq - testConn := makeTestConn(&capturedReq) - - ctx := driverctx.NewContextWithQueryTags(context.Background(), map[string]string{ - "team": "data-eng", - }) - - _, err := testConn.executeStatement(ctx, "SELECT 1", nil) - assert.NoError(t, err) - assert.Equal(t, "team:data-eng", capturedReq.ConfOverlay["query_tags"]) - }) - - t.Run("query tags with special characters in values", func(t *testing.T) { - var capturedReq *cli_service.TExecuteStatementReq - testConn := makeTestConn(&capturedReq) - - ctx := driverctx.NewContextWithQueryTags(context.Background(), map[string]string{ - "url": "http://host:8080", - }) - - _, err := testConn.executeStatement(ctx, "SELECT 1", nil) - assert.NoError(t, err) - assert.Equal(t, `url:http\://host\:8080`, capturedReq.ConfOverlay["query_tags"]) - }) - - t.Run("query tags with empty value", func(t *testing.T) { - var capturedReq *cli_service.TExecuteStatementReq - testConn := makeTestConn(&capturedReq) - - ctx := driverctx.NewContextWithQueryTags(context.Background(), map[string]string{ - "flag": "", - }) - - _, err := testConn.executeStatement(ctx, "SELECT 1", nil) - assert.NoError(t, err) - assert.Equal(t, "flag", capturedReq.ConfOverlay["query_tags"]) - }) - - t.Run("session-level and statement-level query tags coexist", func(t *testing.T) { - // Session-level tags are sent via TOpenSessionReq.Configuration at connect time. - // Statement-level tags are sent via TExecuteStatementReq.ConfOverlay at query time. - // They are independent fields on different requests, so both should work together. - - var capturedOpenReq *cli_service.TOpenSessionReq - var capturedExecReq *cli_service.TExecuteStatementReq - - testClient := &client.TestClient{ - FnOpenSession: func(ctx context.Context, req *cli_service.TOpenSessionReq) (*cli_service.TOpenSessionResp, error) { - capturedOpenReq = req - return &cli_service.TOpenSessionResp{ - Status: &cli_service.TStatus{ - StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, - }, - SessionHandle: &cli_service.TSessionHandle{ - SessionId: &cli_service.THandleIdentifier{ - GUID: []byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}, - }, - }, - }, nil - }, - FnExecuteStatement: func(ctx context.Context, req *cli_service.TExecuteStatementReq) (*cli_service.TExecuteStatementResp, error) { - capturedExecReq = req - return &cli_service.TExecuteStatementResp{ - Status: &cli_service.TStatus{ - StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, - }, - OperationHandle: &cli_service.TOperationHandle{ - OperationId: &cli_service.THandleIdentifier{ - GUID: []byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}, - Secret: []byte("secret"), - }, - }, - DirectResults: &cli_service.TSparkDirectResults{ - OperationStatus: &cli_service.TGetOperationStatusResp{ - Status: &cli_service.TStatus{ - StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, - }, - OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_FINISHED_STATE), - }, - }, - }, nil - }, - } - - // Simulate what connector.Connect() does: pass session params to OpenSession - sessionParams := map[string]string{ - "QUERY_TAGS": "team:platform,env:prod", - "ansi_mode": "false", - } - protocolVersion := int64(cli_service.TProtocolVersion_SPARK_CLI_SERVICE_PROTOCOL_V8) - session, err := testClient.OpenSession(context.Background(), &cli_service.TOpenSessionReq{ - ClientProtocolI64: &protocolVersion, - Configuration: sessionParams, - }) - assert.NoError(t, err) - - // Verify session-level tags were sent in OpenSession - assert.Equal(t, "team:platform,env:prod", capturedOpenReq.Configuration["QUERY_TAGS"]) - assert.Equal(t, "false", capturedOpenReq.Configuration["ansi_mode"]) - - // Create conn with session that has session-level tags - cfg := config.WithDefaults() - cfg.SessionParams = sessionParams - testConn := &conn{ - session: session, - client: testClient, - cfg: cfg, - } - - // Execute with statement-level tags - ctx := driverctx.NewContextWithQueryTags(context.Background(), map[string]string{ - "job": "nightly-etl", - }) - _, err = testConn.executeStatement(ctx, "SELECT 1", nil) - assert.NoError(t, err) - - // Statement-level tags should be in ConfOverlay - assert.Equal(t, "job:nightly-etl", capturedExecReq.ConfOverlay["query_tags"]) - - // ConfOverlay should ONLY have query_tags, not session params - _, hasAnsiMode := capturedExecReq.ConfOverlay["ansi_mode"] - assert.False(t, hasAnsiMode, "session params should not leak into ConfOverlay") - _, hasSessionQueryTags := capturedExecReq.ConfOverlay["QUERY_TAGS"] - assert.False(t, hasSessionQueryTags, "session-level QUERY_TAGS should not be in ConfOverlay") - }) -} - -func TestConn_pollOperation(t *testing.T) { - t.Parallel() - t.Run("pollOperation returns finished state response when query finishes", func(t *testing.T) { - var getOperationStatusCount int - getOperationStatus := func(ctx context.Context, req *cli_service.TGetOperationStatusReq) (r *cli_service.TGetOperationStatusResp, err error) { - getOperationStatusCount++ - getOperationStatusResp := &cli_service.TGetOperationStatusResp{ - OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_FINISHED_STATE), - } - return getOperationStatusResp, nil - } - testClient := &client.TestClient{ - FnGetOperationStatus: getOperationStatus, - } - testConn := &conn{ - session: getTestSession(), - client: testClient, - cfg: config.WithDefaults(), - } - res, err := testConn.pollOperation(context.Background(), &cli_service.TOperationHandle{ - OperationId: &cli_service.THandleIdentifier{ - GUID: []byte{1, 2, 3, 4, 2, 23, 4, 2, 3, 2, 4, 7, 8, 223, 34, 54}, - Secret: []byte("b"), - }, - }) - assert.NoError(t, err) - assert.Equal(t, 1, getOperationStatusCount) - assert.Equal(t, cli_service.TGetOperationStatusResp{ - OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_FINISHED_STATE), - }, *res) - }) - - t.Run("pollOperation returns closed state response when query has been closed", func(t *testing.T) { - var getOperationStatusCount int - getOperationStatus := func(ctx context.Context, req *cli_service.TGetOperationStatusReq) (r *cli_service.TGetOperationStatusResp, err error) { - getOperationStatusCount++ - getOperationStatusResp := &cli_service.TGetOperationStatusResp{ - OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_CLOSED_STATE), - } - return getOperationStatusResp, nil - } - testClient := &client.TestClient{ - FnGetOperationStatus: getOperationStatus, - } - testConn := &conn{ - session: getTestSession(), - client: testClient, - cfg: config.WithDefaults(), - } - res, err := testConn.pollOperation(context.Background(), &cli_service.TOperationHandle{ - OperationId: &cli_service.THandleIdentifier{ - GUID: []byte{1, 2, 3, 4, 2, 23, 4, 2, 3, 2, 3, 4, 4, 223, 34, 54}, - Secret: []byte("b"), - }, - }) - assert.NoError(t, err) - assert.Equal(t, 1, getOperationStatusCount) - assert.Equal(t, cli_service.TGetOperationStatusResp{ - OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_CLOSED_STATE), - }, *res) - }) - - t.Run("pollOperation returns closed state response when query has been closed", func(t *testing.T) { - var getOperationStatusCount int - getOperationStatus := func(ctx context.Context, req *cli_service.TGetOperationStatusReq) (r *cli_service.TGetOperationStatusResp, err error) { - getOperationStatusCount++ - getOperationStatusResp := &cli_service.TGetOperationStatusResp{ - OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_CLOSED_STATE), - } - return getOperationStatusResp, nil - } - testClient := &client.TestClient{ - FnGetOperationStatus: getOperationStatus, - } - testConn := &conn{ - session: getTestSession(), - client: testClient, - cfg: config.WithDefaults(), - } - res, err := testConn.pollOperation(context.Background(), &cli_service.TOperationHandle{ - OperationId: &cli_service.THandleIdentifier{ - GUID: []byte{1, 2, 3, 4, 2, 23, 4, 2, 3, 2, 3, 4, 4, 223, 34, 54}, - Secret: []byte("b"), - }, - }) - assert.NoError(t, err) - assert.Equal(t, 1, getOperationStatusCount) - assert.Equal(t, cli_service.TGetOperationStatusResp{ - OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_CLOSED_STATE), - }, *res) - }) - - t.Run("pollOperation returns unknown state response when query state is unknown", func(t *testing.T) { - var getOperationStatusCount int - getOperationStatus := func(ctx context.Context, req *cli_service.TGetOperationStatusReq) (r *cli_service.TGetOperationStatusResp, err error) { - getOperationStatusCount++ - getOperationStatusResp := &cli_service.TGetOperationStatusResp{ - OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_UKNOWN_STATE), - } - return getOperationStatusResp, nil - } - testClient := &client.TestClient{ - FnGetOperationStatus: getOperationStatus, - } - testConn := &conn{ - session: getTestSession(), - client: testClient, - cfg: config.WithDefaults(), - } - res, err := testConn.pollOperation(context.Background(), &cli_service.TOperationHandle{ - OperationId: &cli_service.THandleIdentifier{ - GUID: []byte{1, 2, 3, 4, 2, 23, 4, 2, 3, 1, 2, 3, 4, 223, 34, 54}, - Secret: []byte("b"), - }, - }) - assert.NoError(t, err) - assert.Equal(t, 1, getOperationStatusCount) - assert.Equal(t, cli_service.TGetOperationStatusResp{ - OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_UKNOWN_STATE), - }, *res) - }) - - t.Run("pollOperation returns error state response when query errors", func(t *testing.T) { - var getOperationStatusCount int - getOperationStatus := func(ctx context.Context, req *cli_service.TGetOperationStatusReq) (r *cli_service.TGetOperationStatusResp, err error) { - getOperationStatusCount++ - getOperationStatusResp := &cli_service.TGetOperationStatusResp{ - OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_ERROR_STATE), - } - return getOperationStatusResp, nil - } - testClient := &client.TestClient{ - FnGetOperationStatus: getOperationStatus, - } - testConn := &conn{ - session: getTestSession(), - client: testClient, - cfg: config.WithDefaults(), - } - res, err := testConn.pollOperation(context.Background(), &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"), - }, - }) - assert.NoError(t, err) - assert.Equal(t, 1, getOperationStatusCount) - assert.Equal(t, cli_service.TGetOperationStatusResp{ - OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_ERROR_STATE), - }, *res) - }) - - t.Run("pollOperation returns finished state response after query cycles through various states", func(t *testing.T) { - var getOperationStatusCount int - getOperationStatus := func(ctx context.Context, req *cli_service.TGetOperationStatusReq) (r *cli_service.TGetOperationStatusResp, err error) { - operationStates := [4]cli_service.TOperationState{cli_service.TOperationState_INITIALIZED_STATE, cli_service.TOperationState_PENDING_STATE, - cli_service.TOperationState_RUNNING_STATE, cli_service.TOperationState_FINISHED_STATE} - getOperationStatusCount++ - getOperationStatusResp := &cli_service.TGetOperationStatusResp{ - OperationState: cli_service.TOperationStatePtr(operationStates[getOperationStatusCount-1]), - } - return getOperationStatusResp, nil - } - testClient := &client.TestClient{ - FnGetOperationStatus: getOperationStatus, - } - testConn := &conn{ - session: getTestSession(), - client: testClient, - cfg: config.WithDefaults(), - } - res, err := testConn.pollOperation(context.Background(), &cli_service.TOperationHandle{ - OperationId: &cli_service.THandleIdentifier{ - GUID: []byte{1, 2, 3, 4, 2, 23, 4, 2, 3, 1, 3, 4, 4, 223, 34, 54}, - Secret: []byte("b"), - }, - }) - assert.NoError(t, err) - assert.Equal(t, 4, getOperationStatusCount) - assert.Equal(t, cli_service.TGetOperationStatusResp{ - OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_FINISHED_STATE), - }, *res) - }) - - t.Run("pollOperation returns cancel err when context times out before get operation", func(t *testing.T) { - var getOperationStatusCount, cancelOperationCount int - getOperationStatus := func(ctx context.Context, req *cli_service.TGetOperationStatusReq) (r *cli_service.TGetOperationStatusResp, err error) { - operationStates := [4]cli_service.TOperationState{cli_service.TOperationState_INITIALIZED_STATE, cli_service.TOperationState_PENDING_STATE, - cli_service.TOperationState_RUNNING_STATE, cli_service.TOperationState_FINISHED_STATE} - getOperationStatusCount++ - getOperationStatusResp := &cli_service.TGetOperationStatusResp{ - OperationState: cli_service.TOperationStatePtr(operationStates[getOperationStatusCount-1]), - } - return getOperationStatusResp, nil - } - cancelOperation := func(ctx context.Context, req *cli_service.TCancelOperationReq) (r *cli_service.TCancelOperationResp, err error) { - cancelOperationCount++ - cancelOperationResp := &cli_service.TCancelOperationResp{ - Status: &cli_service.TStatus{ - StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, - }, - } - return cancelOperationResp, nil - } - testClient := &client.TestClient{ - FnGetOperationStatus: getOperationStatus, - FnCancelOperation: cancelOperation, - } - testConn := &conn{ - session: getTestSession(), - client: testClient, - cfg: config.WithDefaults(), - } - 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"), - }, - }) - assert.Error(t, err) - assert.Equal(t, 0, getOperationStatusCount) - assert.Equal(t, 1, cancelOperationCount) - assert.Nil(t, res) - }) - - t.Run("pollOperation returns cancel err when context times out before get operation", func(t *testing.T) { - var getOperationStatusCount, cancelOperationCount int - getOperationStatus := func(ctx context.Context, req *cli_service.TGetOperationStatusReq) (r *cli_service.TGetOperationStatusResp, err error) { - operationStates := [4]cli_service.TOperationState{cli_service.TOperationState_INITIALIZED_STATE, cli_service.TOperationState_PENDING_STATE, - cli_service.TOperationState_RUNNING_STATE, cli_service.TOperationState_FINISHED_STATE} - getOperationStatusCount++ - getOperationStatusResp := &cli_service.TGetOperationStatusResp{ - OperationState: cli_service.TOperationStatePtr(operationStates[getOperationStatusCount-1]), - } - return getOperationStatusResp, nil - } - cancelOperation := func(ctx context.Context, req *cli_service.TCancelOperationReq) (r *cli_service.TCancelOperationResp, err error) { - cancelOperationCount++ - cancelOperationResp := &cli_service.TCancelOperationResp{ - Status: &cli_service.TStatus{ - StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, - }, - } - return cancelOperationResp, nil - } - testClient := &client.TestClient{ - FnGetOperationStatus: getOperationStatus, - FnCancelOperation: cancelOperation, - } - cfg := config.WithDefaults() - cfg.PollInterval = 100 * time.Millisecond - testConn := &conn{ - session: getTestSession(), - client: testClient, - cfg: cfg, - } - ctx, cancel := context.WithTimeout(context.Background(), 250*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, 3, 4, 223, 34, 54}, - Secret: []byte("b"), - }, - }) - - assert.Error(t, err) - assert.GreaterOrEqual(t, getOperationStatusCount, 1) - assert.Equal(t, 1, cancelOperationCount) - assert.Nil(t, res) - }) - - t.Run("pollOperation returns cancel err when context is cancelled", func(t *testing.T) { - var getOperationStatusCount, cancelOperationCount int - getOperationStatus := func(ctx context.Context, req *cli_service.TGetOperationStatusReq) (r *cli_service.TGetOperationStatusResp, err error) { - operationStates := [4]cli_service.TOperationState{cli_service.TOperationState_INITIALIZED_STATE, cli_service.TOperationState_PENDING_STATE, - cli_service.TOperationState_RUNNING_STATE, cli_service.TOperationState_FINISHED_STATE} - getOperationStatusCount++ - getOperationStatusResp := &cli_service.TGetOperationStatusResp{ - OperationState: cli_service.TOperationStatePtr(operationStates[getOperationStatusCount-1]), - } - return getOperationStatusResp, nil - } - cancelOperation := func(ctx context.Context, req *cli_service.TCancelOperationReq) (r *cli_service.TCancelOperationResp, err error) { - cancelOperationCount++ - cancelOperationResp := &cli_service.TCancelOperationResp{ - Status: &cli_service.TStatus{ - StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, - }, - } - return cancelOperationResp, nil - } - testClient := &client.TestClient{ - FnGetOperationStatus: getOperationStatus, - FnCancelOperation: cancelOperation, - } - cfg := config.WithDefaults() - cfg.PollInterval = 100 * time.Millisecond - testConn := &conn{ - session: getTestSession(), - client: testClient, - cfg: cfg, - } - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - go func() { - time.Sleep(150 * time.Millisecond) - 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"), - }, - }) - assert.Error(t, err) - assert.GreaterOrEqual(t, getOperationStatusCount, 1) - assert.GreaterOrEqual(t, 1, cancelOperationCount) - assert.Nil(t, res) - }) -} - -func TestConn_runQuery(t *testing.T) { - t.Parallel() - t.Run("runQuery should err when client.ExecuteStatement fails", func(t *testing.T) { - var executeStatementCount int - executeStatement := func(ctx context.Context, req *cli_service.TExecuteStatementReq) (r *cli_service.TExecuteStatementResp, err error) { - executeStatementCount++ - return nil, fmt.Errorf("error") - } - testClient := &client.TestClient{ - FnExecuteStatement: executeStatement, - } - testConn := &conn{ - session: getTestSession(), - client: testClient, - cfg: config.WithDefaults(), - } - exStmtResp, opStatusResp, err := testConn.runQuery(context.Background(), "select 1", []driver.NamedValue{}) - assert.Error(t, err) - assert.Nil(t, exStmtResp) - assert.Nil(t, opStatusResp) - assert.Equal(t, 1, executeStatementCount) - }) - - t.Run("runQuery should err when pollOperation fails", func(t *testing.T) { - var executeStatementCount, getOperationStatusCount int - executeStatement := func(ctx context.Context, req *cli_service.TExecuteStatementReq) (r *cli_service.TExecuteStatementResp, err error) { - executeStatementCount++ - executeStatementResp := &cli_service.TExecuteStatementResp{ - Status: &cli_service.TStatus{ - StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, - }, - OperationHandle: &cli_service.TOperationHandle{ - OperationId: &cli_service.THandleIdentifier{ - GUID: []byte{1, 2, 3, 4, 2, 23, 4, 2, 3, 1, 3, 4, 4, 223, 34, 54}, - Secret: []byte("b"), - }, - }, - } - return executeStatementResp, nil - } - - getOperationStatus := func(ctx context.Context, req *cli_service.TGetOperationStatusReq) (r *cli_service.TGetOperationStatusResp, err error) { - getOperationStatusCount++ - getOperationStatusResp := &cli_service.TGetOperationStatusResp{ - OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_ERROR_STATE), - } - return getOperationStatusResp, fmt.Errorf("error on get operation status") - } - - testClient := &client.TestClient{ - FnExecuteStatement: executeStatement, - FnGetOperationStatus: getOperationStatus, - } - testConn := &conn{ - session: getTestSession(), - client: testClient, - cfg: config.WithDefaults(), - } - exStmtResp, opStatusResp, err := testConn.runQuery(context.Background(), "select 1", []driver.NamedValue{}) - - assert.Error(t, err) - assert.Equal(t, 1, executeStatementCount) - assert.Equal(t, 1, getOperationStatusCount) - assert.NotNil(t, exStmtResp) - assert.Nil(t, opStatusResp) - }) - - t.Run("runQuery should return resp when query is finished", func(t *testing.T) { - var executeStatementCount, getOperationStatusCount int - executeStatement := func(ctx context.Context, req *cli_service.TExecuteStatementReq) (r *cli_service.TExecuteStatementResp, err error) { - executeStatementCount++ - executeStatementResp := &cli_service.TExecuteStatementResp{ - Status: &cli_service.TStatus{ - StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, - }, - OperationHandle: &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"), - }, - }, - } - return executeStatementResp, nil - } - var numModRows int64 = 2 - - getOperationStatus := func(ctx context.Context, req *cli_service.TGetOperationStatusReq) (r *cli_service.TGetOperationStatusResp, err error) { - getOperationStatusCount++ - getOperationStatusResp := &cli_service.TGetOperationStatusResp{ - OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_FINISHED_STATE), - NumModifiedRows: &numModRows, - } - return getOperationStatusResp, nil - } - - testClient := &client.TestClient{ - FnExecuteStatement: executeStatement, - FnGetOperationStatus: getOperationStatus, - } - testConn := &conn{ - session: getTestSession(), - client: testClient, - cfg: config.WithDefaults(), - } - exStmtResp, opStatusResp, err := testConn.runQuery(context.Background(), "select 1", []driver.NamedValue{}) - - assert.NoError(t, err) - assert.Equal(t, 1, executeStatementCount) - assert.Equal(t, 1, getOperationStatusCount) - assert.NotNil(t, exStmtResp) - assert.NotNil(t, opStatusResp) - assert.Equal(t, &numModRows, opStatusResp.NumModifiedRows) - }) - - t.Run("runQuery should return resp and error when query is canceled", func(t *testing.T) { - var executeStatementCount, getOperationStatusCount int - executeStatement := func(ctx context.Context, req *cli_service.TExecuteStatementReq) (r *cli_service.TExecuteStatementResp, err error) { - executeStatementCount++ - executeStatementResp := &cli_service.TExecuteStatementResp{ - Status: &cli_service.TStatus{ - StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, - }, - OperationHandle: &cli_service.TOperationHandle{ - OperationId: &cli_service.THandleIdentifier{ - GUID: []byte{1, 2, 3, 4, 2, 23, 4, 2, 3, 2, 4, 4, 223, 34, 23, 54}, - Secret: []byte("b"), - }, - }, - } - return executeStatementResp, nil - } - var numModRows int64 = 3 - - getOperationStatus := func(ctx context.Context, req *cli_service.TGetOperationStatusReq) (r *cli_service.TGetOperationStatusResp, err error) { - getOperationStatusCount++ - getOperationStatusResp := &cli_service.TGetOperationStatusResp{ - OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_CANCELED_STATE), - NumModifiedRows: &numModRows, - } - return getOperationStatusResp, nil - } - - testClient := &client.TestClient{ - FnExecuteStatement: executeStatement, - FnGetOperationStatus: getOperationStatus, - } - testConn := &conn{ - session: getTestSession(), - client: testClient, - cfg: config.WithDefaults(), - } - exStmtResp, opStatusResp, err := testConn.runQuery(context.Background(), "select 1", []driver.NamedValue{}) - - assert.Error(t, err) - assert.Equal(t, 1, executeStatementCount) - assert.Equal(t, 1, getOperationStatusCount) - assert.NotNil(t, exStmtResp) - assert.NotNil(t, opStatusResp) - assert.Equal(t, &numModRows, opStatusResp.NumModifiedRows) - }) - - t.Run("runQuery should return resp when query is finished with DirectResults", func(t *testing.T) { - var executeStatementCount, getOperationStatusCount int - executeStatement := func(ctx context.Context, req *cli_service.TExecuteStatementReq) (r *cli_service.TExecuteStatementResp, err error) { - executeStatementCount++ - executeStatementResp := &cli_service.TExecuteStatementResp{ - Status: &cli_service.TStatus{ - StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, - }, - OperationHandle: &cli_service.TOperationHandle{ - OperationId: &cli_service.THandleIdentifier{ - GUID: []byte{1, 2, 3, 4, 2, 23, 4, 2, 3, 2, 4, 4, 223, 34, 54, 87}, - Secret: []byte("b"), - }, - }, - DirectResults: &cli_service.TSparkDirectResults{ - OperationStatus: &cli_service.TGetOperationStatusResp{ - Status: &cli_service.TStatus{ - StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, - }, - OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_FINISHED_STATE), - }, - }, - } - return executeStatementResp, nil - } - - getOperationStatus := func(ctx context.Context, req *cli_service.TGetOperationStatusReq) (r *cli_service.TGetOperationStatusResp, err error) { - getOperationStatusCount++ - getOperationStatusResp := &cli_service.TGetOperationStatusResp{ - OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_FINISHED_STATE), - } - return getOperationStatusResp, nil - } - - testClient := &client.TestClient{ - FnExecuteStatement: executeStatement, - FnGetOperationStatus: getOperationStatus, - } - testConn := &conn{ - session: getTestSession(), - client: testClient, - cfg: config.WithDefaults(), - } - exStmtResp, opStatusResp, err := testConn.runQuery(context.Background(), "select 1", []driver.NamedValue{}) - - assert.NoError(t, err) - assert.Equal(t, 1, executeStatementCount) - assert.Equal(t, 0, getOperationStatusCount) // GetOperationStatus should not be called, already provided in DirectResults - assert.NotNil(t, exStmtResp) - assert.NotNil(t, opStatusResp) - }) - - t.Run("runQuery should return resp and err when query is cancelled with DirectResults", func(t *testing.T) { - var executeStatementCount, getOperationStatusCount int - executeStatement := func(ctx context.Context, req *cli_service.TExecuteStatementReq) (r *cli_service.TExecuteStatementResp, err error) { - executeStatementCount++ - executeStatementResp := &cli_service.TExecuteStatementResp{ - Status: &cli_service.TStatus{ - StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, - }, - OperationHandle: &cli_service.TOperationHandle{ - OperationId: &cli_service.THandleIdentifier{ - GUID: []byte{1, 2, 3, 4, 2, 23, 4, 2, 3, 2, 3, 4, 4, 223, 34, 54}, - Secret: []byte("b"), - }, - }, - DirectResults: &cli_service.TSparkDirectResults{ - OperationStatus: &cli_service.TGetOperationStatusResp{ - Status: &cli_service.TStatus{ - StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, - }, - OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_CANCELED_STATE), - }, - }, - } - return executeStatementResp, nil - } - - getOperationStatus := func(ctx context.Context, req *cli_service.TGetOperationStatusReq) (r *cli_service.TGetOperationStatusResp, err error) { - getOperationStatusCount++ - getOperationStatusResp := &cli_service.TGetOperationStatusResp{ - OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_FINISHED_STATE), - } - return getOperationStatusResp, nil - } - - testClient := &client.TestClient{ - FnExecuteStatement: executeStatement, - FnGetOperationStatus: getOperationStatus, - } - testConn := &conn{ - session: getTestSession(), - client: testClient, - cfg: config.WithDefaults(), - } - exStmtResp, opStatusResp, err := testConn.runQuery(context.Background(), "select 1", []driver.NamedValue{}) - - assert.Error(t, err) - assert.Equal(t, 1, executeStatementCount) - assert.Equal(t, 0, getOperationStatusCount) // GetOperationStatus should not be called, already provided in DirectResults - assert.NotNil(t, exStmtResp) - assert.NotNil(t, opStatusResp) - }) - - t.Run("runQuery should return resp when query is finished but DirectResults still live", func(t *testing.T) { - var executeStatementCount, getOperationStatusCount int - executeStatement := func(ctx context.Context, req *cli_service.TExecuteStatementReq) (r *cli_service.TExecuteStatementResp, err error) { - executeStatementCount++ - executeStatementResp := &cli_service.TExecuteStatementResp{ - Status: &cli_service.TStatus{ - StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, - }, - OperationHandle: &cli_service.TOperationHandle{ - OperationId: &cli_service.THandleIdentifier{ - GUID: []byte{1, 2, 3, 4, 2, 23, 4, 2, 3, 2, 3, 4, 4, 223, 34, 54}, - Secret: []byte("b"), - }, - }, - DirectResults: &cli_service.TSparkDirectResults{ - OperationStatus: &cli_service.TGetOperationStatusResp{ - Status: &cli_service.TStatus{ - StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, - }, - OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_INITIALIZED_STATE), - }, - }, - } - return executeStatementResp, nil - } - var numModRows int64 = 3 - getOperationStatus := func(ctx context.Context, req *cli_service.TGetOperationStatusReq) (r *cli_service.TGetOperationStatusResp, err error) { - getOperationStatusCount++ - getOperationStatusResp := &cli_service.TGetOperationStatusResp{ - OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_FINISHED_STATE), - NumModifiedRows: &numModRows, - } - return getOperationStatusResp, nil - } - - testClient := &client.TestClient{ - FnExecuteStatement: executeStatement, - FnGetOperationStatus: getOperationStatus, - } - testConn := &conn{ - session: getTestSession(), - client: testClient, - cfg: config.WithDefaults(), - } - exStmtResp, opStatusResp, err := testConn.runQuery(context.Background(), "select 1", []driver.NamedValue{}) - - assert.NoError(t, err) - assert.Equal(t, 1, executeStatementCount) - assert.Equal(t, &numModRows, opStatusResp.NumModifiedRows) - assert.Equal(t, 1, getOperationStatusCount) - assert.NotNil(t, exStmtResp) - assert.NotNil(t, opStatusResp) - }) - - t.Run("runQuery should return resp and err when query is cancelled after DirectResults still live", func(t *testing.T) { - var executeStatementCount, getOperationStatusCount int - executeStatement := func(ctx context.Context, req *cli_service.TExecuteStatementReq) (r *cli_service.TExecuteStatementResp, err error) { - executeStatementCount++ - executeStatementResp := &cli_service.TExecuteStatementResp{ - Status: &cli_service.TStatus{ - StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, - }, - OperationHandle: &cli_service.TOperationHandle{ - OperationId: &cli_service.THandleIdentifier{ - GUID: []byte{1, 2, 3, 4, 2, 23, 4, 2, 3, 2, 3, 4, 4, 223, 34, 54}, - Secret: []byte("b"), - }, - }, - DirectResults: &cli_service.TSparkDirectResults{ - OperationStatus: &cli_service.TGetOperationStatusResp{ - Status: &cli_service.TStatus{ - StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, - }, - OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_INITIALIZED_STATE), - }, - }, - } - return executeStatementResp, nil - } - - getOperationStatus := func(ctx context.Context, req *cli_service.TGetOperationStatusReq) (r *cli_service.TGetOperationStatusResp, err error) { - getOperationStatusCount++ - getOperationStatusResp := &cli_service.TGetOperationStatusResp{ - OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_CANCELED_STATE), - } - return getOperationStatusResp, nil - } - - testClient := &client.TestClient{ - FnExecuteStatement: executeStatement, - FnGetOperationStatus: getOperationStatus, - } - testConn := &conn{ - session: getTestSession(), - client: testClient, - cfg: config.WithDefaults(), - } - exStmtResp, opStatusResp, err := testConn.runQuery(context.Background(), "select 1", []driver.NamedValue{}) - - assert.Error(t, err) - assert.Equal(t, 1, executeStatementCount) - assert.Equal(t, 1, getOperationStatusCount) - assert.NotNil(t, exStmtResp) - assert.NotNil(t, opStatusResp) - }) -} - func TestConn_ExecContext(t *testing.T) { t.Parallel() t.Run("ExecContext currently does not support query parameters", func(t *testing.T) { @@ -1407,9 +32,8 @@ func TestConn_ExecContext(t *testing.T) { testClient := &client.TestClient{} testConn := &conn{ - session: getTestSession(), - client: testClient, cfg: config.WithDefaults(), + backend: thriftbackend.NewForTest(testClient, getTestSession(), config.WithDefaults()), } res, err := testConn.ExecContext(context.Background(), "select 1", []driver.NamedValue{ {Value: 1, Name: "name"}, @@ -1447,9 +71,8 @@ func TestConn_ExecContext(t *testing.T) { }, } testConn := &conn{ - session: getTestSession(), - client: testClient, cfg: config.WithDefaults(), + backend: thriftbackend.NewForTest(testClient, getTestSession(), config.WithDefaults()), } res, err := testConn.ExecContext(context.Background(), "select 1", []driver.NamedValue{}) @@ -1500,9 +123,8 @@ func TestConn_ExecContext(t *testing.T) { FnGetResultSetMetadata: getResultSetMetadata, } testConn := &conn{ - session: getTestSession(), - client: testClient, cfg: config.WithDefaults(), + backend: thriftbackend.NewForTest(testClient, getTestSession(), config.WithDefaults()), } res, err := testConn.ExecContext(context.Background(), "insert 10", []driver.NamedValue{}) @@ -1568,9 +190,8 @@ func TestConn_ExecContext(t *testing.T) { } testConn := &conn{ - session: getTestSession(), - client: testClient, cfg: config.WithDefaults(), + backend: thriftbackend.NewForTest(testClient, getTestSession(), config.WithDefaults()), } ctx := context.Background() ctx, cancel = context.WithCancel(ctx) @@ -1601,9 +222,8 @@ func TestConn_QueryContext(t *testing.T) { testClient := &client.TestClient{} testConn := &conn{ - session: getTestSession(), - client: testClient, cfg: config.WithDefaults(), + backend: thriftbackend.NewForTest(testClient, getTestSession(), config.WithDefaults()), } res, err := testConn.QueryContext(context.Background(), "select 1", []driver.NamedValue{ {Value: 1, Name: "name"}, @@ -1636,9 +256,8 @@ func TestConn_QueryContext(t *testing.T) { FnExecuteStatement: executeStatement, } testConn := &conn{ - session: getTestSession(), - client: testClient, cfg: config.WithDefaults(), + backend: thriftbackend.NewForTest(testClient, getTestSession(), config.WithDefaults()), } res, err := testConn.QueryContext(context.Background(), "select 1", []driver.NamedValue{}) @@ -1679,9 +298,8 @@ func TestConn_QueryContext(t *testing.T) { FnGetOperationStatus: getOperationStatus, } testConn := &conn{ - session: getTestSession(), - client: testClient, cfg: config.WithDefaults(), + backend: thriftbackend.NewForTest(testClient, getTestSession(), config.WithDefaults()), } rows, err := testConn.QueryContext(context.Background(), "select 1", []driver.NamedValue{}) @@ -1714,9 +332,8 @@ func TestConn_Ping(t *testing.T) { FnExecuteStatement: executeStatement, } testConn := &conn{ - session: getTestSession(), - client: testClient, cfg: config.WithDefaults(), + backend: thriftbackend.NewForTest(testClient, getTestSession(), config.WithDefaults()), } err := testConn.Ping(context.Background()) @@ -1763,9 +380,8 @@ func TestConn_Ping(t *testing.T) { } testConn := &conn{ - session: getTestSession(), - client: testClient, cfg: config.WithDefaults(), + backend: thriftbackend.NewForTest(testClient, getTestSession(), config.WithDefaults()), } err := testConn.Ping(context.Background()) @@ -1778,9 +394,8 @@ func TestConn_Ping(t *testing.T) { func TestConn_Begin(t *testing.T) { t.Run("Begin not supported", func(t *testing.T) { testConn := &conn{ - session: getTestSession(), - client: &client.TestClient{}, cfg: config.WithDefaults(), + backend: thriftbackend.NewForTest(&client.TestClient{}, getTestSession(), config.WithDefaults()), } res, err := testConn.Begin() assert.Nil(t, res) @@ -1791,9 +406,8 @@ func TestConn_Begin(t *testing.T) { func TestConn_BeginTx(t *testing.T) { t.Run("BeginTx not supported", func(t *testing.T) { testConn := &conn{ - session: getTestSession(), - client: &client.TestClient{}, cfg: config.WithDefaults(), + backend: thriftbackend.NewForTest(&client.TestClient{}, getTestSession(), config.WithDefaults()), } res, err := testConn.BeginTx(context.Background(), driver.TxOptions{}) assert.Nil(t, res) @@ -1804,9 +418,8 @@ func TestConn_BeginTx(t *testing.T) { func TestConn_ResetSession(t *testing.T) { t.Run("ResetSession not currently supported", func(t *testing.T) { testConn := &conn{ - session: getTestSession(), - client: &client.TestClient{}, cfg: config.WithDefaults(), + backend: thriftbackend.NewForTest(&client.TestClient{}, getTestSession(), config.WithDefaults()), } res := testConn.ResetSession(context.Background()) assert.Nil(t, res) @@ -1831,9 +444,8 @@ func TestConn_Close(t *testing.T) { FnCloseSession: closeSession, } testConn := &conn{ - session: getTestSession(), - client: testClient, cfg: config.WithDefaults(), + backend: thriftbackend.NewForTest(testClient, getTestSession(), config.WithDefaults()), } err := testConn.Close() @@ -1858,9 +470,8 @@ func TestConn_Close(t *testing.T) { FnCloseSession: closeSession, } testConn := &conn{ - session: getTestSession(), - client: testClient, cfg: config.WithDefaults(), + backend: thriftbackend.NewForTest(testClient, getTestSession(), config.WithDefaults()), } err := testConn.Close() @@ -1873,9 +484,8 @@ func TestConn_Prepare(t *testing.T) { t.Run("Prepare returns stmt struct", func(t *testing.T) { testClient := &client.TestClient{} testConn := &conn{ - session: getTestSession(), - client: testClient, cfg: config.WithDefaults(), + backend: thriftbackend.NewForTest(testClient, getTestSession(), config.WithDefaults()), } stmt, err := testConn.Prepare("query string") assert.NoError(t, err) @@ -1887,9 +497,8 @@ func TestConn_PrepareContext(t *testing.T) { t.Run("PrepareContext returns stmt struct", func(t *testing.T) { testClient := &client.TestClient{} testConn := &conn{ - session: getTestSession(), - client: testClient, cfg: config.WithDefaults(), + backend: thriftbackend.NewForTest(testClient, getTestSession(), config.WithDefaults()), } stmt, err := testConn.PrepareContext(context.Background(), "query string") assert.NoError(t, err) @@ -1901,9 +510,8 @@ func TestConn_execStagingOperation(t *testing.T) { t.Run("handles nil IsStagingOperation from DirectResults", func(t *testing.T) { testClient := &client.TestClient{} testConn := &conn{ - session: getTestSession(), - client: testClient, cfg: config.WithDefaults(), + backend: thriftbackend.NewForTest(testClient, getTestSession(), config.WithDefaults()), } // Create response with nil IsStagingOperation in DirectResults @@ -1942,7 +550,8 @@ func TestConn_execStagingOperation(t *testing.T) { testClient.FnGetResultSetMetadata = getResultSetMetadata ctx := context.Background() - err := testConn.execStagingOperation(exStmtResp, ctx) + op := thriftbackend.NewForTest(testClient, getTestSession(), config.WithDefaults()).OperationForTest(exStmtResp, nil) + err := testConn.execStagingOperation(op, ctx) assert.Nil(t, err) assert.Equal(t, 0, getResultSetMetadataCount) // should not be called since DirectResults.ResultSetMetadata exists @@ -1951,9 +560,8 @@ func TestConn_execStagingOperation(t *testing.T) { t.Run("handles nil IsStagingOperation from GetResultSetMetadata", func(t *testing.T) { testClient := &client.TestClient{} testConn := &conn{ - session: getTestSession(), - client: testClient, cfg: config.WithDefaults(), + backend: thriftbackend.NewForTest(testClient, getTestSession(), config.WithDefaults()), } // Create response with nil DirectResults.ResultSetMetadata @@ -1984,7 +592,8 @@ func TestConn_execStagingOperation(t *testing.T) { testClient.FnGetResultSetMetadata = getResultSetMetadata ctx := context.Background() - err := testConn.execStagingOperation(exStmtResp, ctx) + op := thriftbackend.NewForTest(testClient, getTestSession(), config.WithDefaults()).OperationForTest(exStmtResp, nil) + err := testConn.execStagingOperation(op, ctx) assert.Nil(t, err) assert.Equal(t, 1, getResultSetMetadataCount) // should be called since DirectResults.ResultSetMetadata is nil diff --git a/connector.go b/connector.go index 306a6f89..cf5a9c2e 100644 --- a/connector.go +++ b/connector.go @@ -4,7 +4,6 @@ import ( "context" "crypto/tls" "database/sql/driver" - "fmt" "net/http" "net/url" "regexp" @@ -16,11 +15,10 @@ import ( "github.com/databricks/databricks-sql-go/auth/pat" "github.com/databricks/databricks-sql-go/auth/tokenprovider" "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/backend/thrift" "github.com/databricks/databricks-sql-go/internal/client" "github.com/databricks/databricks-sql-go/internal/config" - dbsqlerrint "github.com/databricks/databricks-sql-go/internal/errors" + "github.com/databricks/databricks-sql-go/internal/debuglog" "github.com/databricks/databricks-sql-go/logger" "github.com/databricks/databricks-sql-go/telemetry" ) @@ -32,53 +30,25 @@ type connector struct { // 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 - if c.cfg.Catalog != "" { - catalogName = cli_service.TIdentifierPtr(cli_service.TIdentifier(c.cfg.Catalog)) - } - if c.cfg.Schema != "" { - schemaName = cli_service.TIdentifierPtr(cli_service.TIdentifier(c.cfg.Schema)) - } + defer debuglog.Track(ctx, "connector.Connect", "host=%s", c.cfg.Host)() - tclient, err := client.InitThriftClient(c.cfg, c.client) + // Build the execution backend. Currently always the Thrift backend; a second + // SEA-via-kernel backend will be selected here by a connect-time flag. + be, err := thrift.New(ctx, c.cfg, c.client) if err != nil { - return nil, dbsqlerrint.NewDriverError(ctx, dbsqlerr.ErrThriftClient, err) - } - - // Prepare session configuration - sessionParams := make(map[string]string) - for k, v := range c.cfg.SessionParams { - sessionParams[k] = v + return nil, err } - if c.cfg.EnableMetricViewMetadata { - sessionParams["spark.sql.thriftserver.metadata.metricview.enabled"] = "true" - } - - protocolVersion := int64(c.cfg.ThriftProtocolVersion) - sessionStart := time.Now() - session, err := tclient.OpenSession(ctx, &cli_service.TOpenSessionReq{ - ClientProtocolI64: &protocolVersion, - Configuration: sessionParams, - InitialNamespace: &cli_service.TNamespace{ - CatalogName: catalogName, - SchemaName: schemaName, - }, - CanUseMultipleCatalogs: &c.cfg.CanUseMultipleCatalogs, - }) - sessionLatencyMs := time.Since(sessionStart).Milliseconds() - - if err != nil { - 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) + if err := be.OpenSession(ctx); err != nil { + return nil, err } + sessionLatencyMs := time.Since(sessionStart).Milliseconds() conn := &conn{ - id: client.SprintGuid(session.SessionHandle.GetSessionId().GUID), + id: be.SessionID(), cfg: c.cfg, - client: tclient, - session: session, + backend: be, } log := logger.WithContext(conn.id, driverctx.CorrelationIdFromContext(ctx), "") @@ -107,7 +77,7 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) { conn.telemetry.RecordOperation(ctx, conn.id, "", telemetry.OperationTypeCreateSession, sessionLatencyMs, nil) } - log.Info().Msgf("connect: host=%s port=%d httpPath=%s serverProtocolVersion=0x%X", c.cfg.Host, c.cfg.Port, c.cfg.HTTPPath, session.ServerProtocolVersion) + log.Info().Msgf("connect: host=%s port=%d httpPath=%s serverProtocolVersion=0x%X", c.cfg.Host, c.cfg.Port, c.cfg.HTTPPath, be.ServerProtocolVersion()) return conn, nil } diff --git a/internal/backend/backend.go b/internal/backend/backend.go new file mode 100644 index 00000000..79d83cab --- /dev/null +++ b/internal/backend/backend.go @@ -0,0 +1,109 @@ +// Package backend defines the execution-backend abstraction for the +// databricks-sql-go driver. A connection holds one Backend, created by the +// connector at connect time, and drives session lifecycle and statement +// execution through it without depending on a concrete protocol. +// +// The interfaces are expressed in neutral terms — plain Go request structs, a +// driver.Rows result, and small accessors for statement id, affected rows, and +// sqlstate-aware error wrapping — so that only a Backend implementation imports +// its protocol's client (e.g. the Thrift backend imports internal/cli_service). +package backend + +import ( + "context" + "database/sql/driver" + + dbsqlrows "github.com/databricks/databricks-sql-go/internal/rows" +) + +// Backend is the per-connection execution backend. Exactly one Backend backs one +// conn, which (via database/sql's pool) is used by one goroutine at a time — so +// implementations need not be safe for concurrent use by multiple goroutines. +type Backend interface { + // OpenSession establishes the server-side session. Called once by the + // connector factory at connect time, before the conn is handed to the pool. + OpenSession(ctx context.Context) error + + // CloseSession tears down the server-side session. Called from conn.Close. + CloseSession(ctx context.Context) error + + // SessionValid backs driver.Validator (conn.IsValid): it reports whether the + // session is still usable, so a broken conn is evicted from the pool rather + // than reused. It does no I/O — it inspects state captured at OpenSession. + SessionValid() bool + + // SessionID is the formatted server session id (conn.id). Valid only after a + // successful OpenSession; empty otherwise. + SessionID() string + + // Execute runs a statement to a terminal state and returns an Operation + // describing the outcome. + // + // An implementation MUST return a non-nil Operation, even when err is non-nil: + // the caller reads StatementID, wraps the error via ExecutionError, and closes + // the operation via Close on both the success and failure paths. When the + // server returned no handle, the accessors return zero values (StatementID "", + // Close closed=false) rather than panicking. + Execute(ctx context.Context, req ExecRequest) (Operation, error) +} + +// Operation is a submitted statement that has reached (or attempted to reach) a +// terminal state. It exposes only neutral accessors, keeping the concrete +// protocol types inside the backend implementation. +type Operation interface { + // StatementID is the server statement/operation id (empty if the server + // returned no handle). + StatementID() string + + // AffectedRows is the modified-row count for ExecContext results. Unlike the + // other accessors it is defined only on the success path — the caller reads it + // only after Execute returned a nil error — so an implementation need not make + // it safe to call on a failed/handle-less operation. + AffectedRows() int64 + + // Results builds the driver.Rows for the operation's result set. The callbacks + // carry the telemetry hooks (chunk timing, close, CloudFetch file); the caller + // supplies whichever subset it needs (a full query wires all three, the + // staging path wires only chunk timing). + Results(ctx context.Context, callbacks *dbsqlrows.TelemetryCallbacks) (driver.Rows, error) + + // IsStaging reports whether this is a staging (PUT/GET/REMOVE) operation. It + // may issue a metadata lookup if the answer is not already available from + // direct results. + IsStaging(ctx context.Context) (bool, error) + + // Close best-effort closes the server-side operation, if one is still open. + // It is idempotent and safe to call regardless of whether execution + // succeeded. closed reports whether a close RPC was actually issued (false + // when the operation had no handle or was already closed), so the caller can + // record CLOSE_STATEMENT telemetry only when a real close happened. + Close(ctx context.Context) (closed bool, err error) + + // ExecutionError wraps cause as the driver's execution error, attaching this + // operation's terminal sqlstate. Returns nil when cause is nil. + ExecutionError(ctx context.Context, cause error) error +} + +// ExecRequest is a backend-neutral statement-execution request. Per-backend wire +// options (Arrow flags, CloudFetch, direct results, query tags, timeouts) are +// derived from the config and context inside the backend implementation. +type ExecRequest struct { + // Query is the SQL text to execute. + Query string + // Params are the bound query parameters, already type-inferred and rendered + // to their string wire form (see Param); empty when the statement has none. + // Inference lives in the dbsql package because it depends on the public + // dbsql.Parameter/SqlType types; each backend maps these neutral params to its + // own wire type. + Params []Param +} + +// Param is one bound query parameter in backend-neutral form: the name (empty +// for a positional parameter), the SQL type name as the server expects it (e.g. +// "BIGINT", "DECIMAL(2,1)", "VOID"), and the value already rendered to its +// string wire representation. A nil Value denotes a SQL NULL (type "VOID"). +type Param struct { + Name string + Type string + Value *string +} diff --git a/internal/backend/thrift/backend.go b/internal/backend/thrift/backend.go new file mode 100644 index 00000000..444130eb --- /dev/null +++ b/internal/backend/thrift/backend.go @@ -0,0 +1,441 @@ +// Package thrift is the Thrift/HiveServer2 implementation of backend.Backend. +// It owns the generated TCLIService client, the open-session response, and the +// execute/poll/close/session RPCs, so that internal/cli_service is imported only +// here and not by the rest of the driver. +// +// Steps are instrumented via internal/debuglog (gated, ordered, timed, +// function-tagged) so a failing or slow step is visible across the +// execute/poll/fetch flow. +package thrift + +import ( + "context" + "fmt" + "net/http" + "time" + + "github.com/databricks/databricks-sql-go/driverctx" + dbsqlerr "github.com/databricks/databricks-sql-go/errors" + "github.com/databricks/databricks-sql-go/internal/backend" + "github.com/databricks/databricks-sql-go/internal/cli_service" + "github.com/databricks/databricks-sql-go/internal/client" + context2 "github.com/databricks/databricks-sql-go/internal/compat/context" + "github.com/databricks/databricks-sql-go/internal/config" + "github.com/databricks/databricks-sql-go/internal/debuglog" + dbsqlerrint "github.com/databricks/databricks-sql-go/internal/errors" + "github.com/databricks/databricks-sql-go/internal/querytags" + "github.com/databricks/databricks-sql-go/internal/sentinel" + "github.com/databricks/databricks-sql-go/internal/thrift_protocol" + "github.com/databricks/databricks-sql-go/logger" + "github.com/pkg/errors" +) + +// Backend is the Thrift implementation of backend.Backend. One Backend backs one +// conn and is used by a single goroutine at a time (database/sql pool +// discipline), so it holds no locks. +type Backend struct { + cfg *config.Config + client cli_service.TCLIService + session *cli_service.TOpenSessionResp + // sessionID is the formatted session GUID, computed once in OpenSession and + // reused. Formatting it (SprintGuid) allocates, and it is read on every query + // (connId enrichment, logging, telemetry), so it is cached rather than + // recomputed per call. + sessionID string +} + +var _ backend.Backend = (*Backend)(nil) + +// New builds a Thrift backend from the driver config and the shared HTTP client. +// The Thrift client is constructed here; the session is not opened until +// OpenSession. ctx carries the caller's correlation id onto a client-init error. +func New(ctx context.Context, cfg *config.Config, httpClient *http.Client) (*Backend, error) { + defer debuglog.Track(ctx, "thrift.New", "host=%s", cfg.Host)() + + tclient, err := client.InitThriftClient(cfg, httpClient) + if err != nil { + debuglog.Logf(ctx, "thrift.New", "InitThriftClient failed: %v", err) + return nil, dbsqlerrint.NewDriverError(ctx, dbsqlerr.ErrThriftClient, err) + } + return &Backend{cfg: cfg, client: tclient}, nil +} + +// OpenSession opens the server-side Thrift session, wiring the session params +// and initial namespace (catalog/schema) from the config. +func (b *Backend) OpenSession(ctx context.Context) error { + defer debuglog.Track(ctx, "thrift.Backend.OpenSession", "host=%s", b.cfg.Host)() + + var catalogName *cli_service.TIdentifier + var schemaName *cli_service.TIdentifier + if b.cfg.Catalog != "" { + catalogName = cli_service.TIdentifierPtr(cli_service.TIdentifier(b.cfg.Catalog)) + } + if b.cfg.Schema != "" { + schemaName = cli_service.TIdentifierPtr(cli_service.TIdentifier(b.cfg.Schema)) + } + + sessionParams := make(map[string]string) + for k, v := range b.cfg.SessionParams { + sessionParams[k] = v + } + if b.cfg.EnableMetricViewMetadata { + sessionParams["spark.sql.thriftserver.metadata.metricview.enabled"] = "true" + } + + protocolVersion := int64(b.cfg.ThriftProtocolVersion) + + debuglog.Logf(ctx, "thrift.Backend.OpenSession", "sending OpenSession protocolVersion=0x%X catalog=%q schema=%q", protocolVersion, b.cfg.Catalog, b.cfg.Schema) + session, err := b.client.OpenSession(ctx, &cli_service.TOpenSessionReq{ + ClientProtocolI64: &protocolVersion, + Configuration: sessionParams, + InitialNamespace: &cli_service.TNamespace{ + CatalogName: catalogName, + SchemaName: schemaName, + }, + CanUseMultipleCatalogs: &b.cfg.CanUseMultipleCatalogs, + }) + if err != nil { + debuglog.Logf(ctx, "thrift.Backend.OpenSession", "OpenSession failed: %v", err) + return dbsqlerrint.NewRequestError(ctx, fmt.Sprintf("error connecting: host=%s port=%d, httpPath=%s", b.cfg.Host, b.cfg.Port, b.cfg.HTTPPath), err) + } + b.session = session + if session.SessionHandle != nil { + b.sessionID = client.SprintGuid(session.SessionHandle.GetSessionId().GUID) + } + debuglog.Logf(ctx, "thrift.Backend.OpenSession", "session opened id=%s serverProtocol=0x%X", b.sessionID, session.ServerProtocolVersion) + return nil +} + +// SessionID returns the formatted server session id (conn.id's value), computed +// once in OpenSession. Empty until OpenSession succeeds. +func (b *Backend) SessionID() string { + return b.sessionID +} + +// ServerProtocolVersion returns the negotiated Thrift protocol version for the +// open session, for the connector's connect log. It is Thrift-specific and so is +// not part of the neutral backend.Backend interface; the connector reads it off +// the concrete *Backend it constructed. Returns 0 before OpenSession. +func (b *Backend) ServerProtocolVersion() cli_service.TProtocolVersion { + if b.session == nil { + return 0 + } + return b.session.ServerProtocolVersion +} + +// SessionValid backs conn.IsValid: the session is usable while its open-status +// code is SUCCESS. No I/O — matches today's behavior exactly. +func (b *Backend) SessionValid() bool { + if b.session == nil { + return false + } + return b.session.GetStatus().StatusCode == cli_service.TStatusCode_SUCCESS_STATUS +} + +// CloseSession closes the server-side session, returning the raw RPC error. The +// caller owns the DELETE_SESSION telemetry timing and the error classification. +func (b *Backend) CloseSession(ctx context.Context) error { + defer debuglog.Track(ctx, "thrift.Backend.CloseSession", "id=%s", b.SessionID())() + + if b.session == nil { + return nil + } + _, err := b.client.CloseSession(ctx, &cli_service.TCloseSessionReq{ + SessionHandle: b.session.SessionHandle, + }) + if err != nil { + debuglog.Logf(ctx, "thrift.Backend.CloseSession", "CloseSession failed: %v", err) + } + return err +} + +// Execute runs the statement to a terminal state and returns a thriftOperation +// carrying the execute and status responses. Per the backend.Backend contract +// the returned Operation is non-nil even on error. +func (b *Backend) Execute(ctx context.Context, req backend.ExecRequest) (backend.Operation, error) { + defer debuglog.Track(ctx, "thrift.Backend.Execute", "sql.len=%d params=%d", len(req.Query), len(req.Params))() + + exStmtResp, opStatusResp, err := b.runQuery(ctx, req) + op := &thriftOperation{ + backend: b, + exStmtResp: exStmtResp, + opStatusResp: opStatusResp, + } + if err != nil { + debuglog.Logf(ctx, "thrift.Backend.Execute", "runQuery error: %v", err) + } + return op, err +} + +// runQuery executes the statement, then drives the operation to a terminal state +// via direct results or by polling. +func (b *Backend) runQuery(ctx context.Context, req backend.ExecRequest) (*cli_service.TExecuteStatementResp, *cli_service.TGetOperationStatusResp, error) { + defer debuglog.Track(ctx, "thrift.Backend.runQuery", "")() + + exStmtResp, err := b.executeStatement(ctx, req) + var log *logger.DBSQLLogger + log, ctx = client.LoggerAndContext(ctx, exStmtResp) + + if err != nil { + return exStmtResp, nil, err + } + + opHandle := exStmtResp.OperationHandle + + if exStmtResp.DirectResults != nil { + opStatus := exStmtResp.DirectResults.GetOperationStatus() + debuglog.Logf(ctx, "thrift.Backend.runQuery", "direct results present, state=%s", opStatus.GetOperationState()) + + switch opStatus.GetOperationState() { + // terminal states + // good + case cli_service.TOperationState_FINISHED_STATE: + 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: + logBadQueryState(log, opStatus) + return exStmtResp, opStatus, unexpectedOperationState(opStatus) + // live states + case cli_service.TOperationState_INITIALIZED_STATE, + cli_service.TOperationState_PENDING_STATE, + cli_service.TOperationState_RUNNING_STATE: + statusResp, err := b.pollOperation(ctx, opHandle) + if err != nil { + return exStmtResp, statusResp, err + } + switch statusResp.GetOperationState() { + // terminal states + // good + case cli_service.TOperationState_FINISHED_STATE: + 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: + logBadQueryState(log, statusResp) + return exStmtResp, statusResp, unexpectedOperationState(statusResp) + // live states + default: + logBadQueryState(log, statusResp) + return exStmtResp, statusResp, invalidOperationState(ctx, statusResp) + } + // weird states + default: + logBadQueryState(log, opStatus) + return exStmtResp, opStatus, invalidOperationState(ctx, opStatus) + } + + } else { + statusResp, err := b.pollOperation(ctx, opHandle) + if err != nil { + return exStmtResp, statusResp, err + } + switch statusResp.GetOperationState() { + // terminal states + // good + case cli_service.TOperationState_FINISHED_STATE: + 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: + logBadQueryState(log, statusResp) + return exStmtResp, statusResp, unexpectedOperationState(statusResp) + // live states + default: + logBadQueryState(log, statusResp) + return exStmtResp, statusResp, invalidOperationState(ctx, statusResp) + } + } +} + +// executeStatement issues the ExecuteStatement RPC, gating each wire option +// (direct results, LZ4, CloudFetch, Arrow, parameters, query tags) on server +// protocol support, and cancels the operation if the context is done. +func (b *Backend) executeStatement(ctx context.Context, req backend.ExecRequest) (*cli_service.TExecuteStatementResp, error) { + ctx = driverctx.NewContextWithConnId(ctx, b.SessionID()) + defer debuglog.Track(ctx, "thrift.Backend.executeStatement", "")() + + parameters := toSparkParameters(req.Params) + + thriftReq := cli_service.TExecuteStatementReq{ + SessionHandle: b.session.SessionHandle, + Statement: req.Query, + RunAsync: true, + QueryTimeout: int64(b.cfg.QueryTimeout / time.Second), + } + + // Check protocol version for feature support + serverProtocolVersion := b.session.ServerProtocolVersion + + // Add direct results if supported + if thrift_protocol.SupportsDirectResults(serverProtocolVersion) { + thriftReq.GetDirectResults = &cli_service.TSparkGetDirectResults{ + MaxRows: int64(b.cfg.MaxRows), + } + } + + // Add LZ4 compression if supported and enabled + if thrift_protocol.SupportsLz4Compression(serverProtocolVersion) && b.cfg.UseLz4Compression { + thriftReq.CanDecompressLZ4Result_ = &b.cfg.UseLz4Compression + } + + // Add cloud fetch if supported and enabled + if thrift_protocol.SupportsCloudFetch(serverProtocolVersion) && b.cfg.UseCloudFetch { + thriftReq.CanDownloadResult_ = &b.cfg.UseCloudFetch + } + + // Add Arrow support if supported and enabled + if thrift_protocol.SupportsArrow(serverProtocolVersion) && b.cfg.UseArrowBatches { + thriftReq.CanReadArrowResult_ = &b.cfg.UseArrowBatches + thriftReq.UseArrowNativeTypes = &cli_service.TSparkArrowTypes{ + DecimalAsArrow: &b.cfg.UseArrowNativeDecimal, + TimestampAsArrow: &b.cfg.UseArrowNativeTimestamp, + ComplexTypesAsArrow: &b.cfg.UseArrowNativeComplexTypes, + IntervalTypesAsArrow: &b.cfg.UseArrowNativeIntervalTypes, + } + } + + // Add parameters if supported and provided + if thrift_protocol.SupportsParameterizedQueries(serverProtocolVersion) && len(parameters) > 0 { + thriftReq.Parameters = parameters + } + + // Add per-statement query tags if provided via context + if queryTags := driverctx.QueryTagsFromContext(ctx); len(queryTags) > 0 { + serialized := querytags.Serialize(queryTags) + if serialized != "" { + if thriftReq.ConfOverlay == nil { + thriftReq.ConfOverlay = make(map[string]string) + } + thriftReq.ConfOverlay["query_tags"] = serialized + } + } + + debuglog.Logf(ctx, "thrift.Backend.executeStatement", "sending ExecuteStatement runAsync=true directResults=%t params=%d", thriftReq.GetDirectResults != nil, len(parameters)) + resp, err := b.client.ExecuteStatement(ctx, &thriftReq) + var log *logger.DBSQLLogger + log, ctx = client.LoggerAndContext(ctx, resp) + + var shouldCancel = func(resp *cli_service.TExecuteStatementResp) bool { + if resp == nil { + return false + } + hasHandle := resp.OperationHandle != nil + isOpen := resp.DirectResults == nil || resp.DirectResults.CloseOperation == nil + return hasHandle && isOpen + } + + select { + default: + // Non-blocking check: continue if context not done + case <-ctx.Done(): + newCtx := driverctx.NewContextFromBackground(ctx) + // in case context is done, we need to cancel the operation if necessary + if err == nil && shouldCancel(resp) { + debuglog.Logf(newCtx, "thrift.Backend.executeStatement", "context done, canceling query") + log.Debug().Msg("databricks: canceling query") + _, err1 := b.client.CancelOperation(newCtx, &cli_service.TCancelOperationReq{ + OperationHandle: resp.GetOperationHandle(), + }) + + if err1 != nil { + log.Err(err1).Msgf("databricks: cancel failed") + } else { + log.Debug().Msgf("databricks: cancel success") + } + } else { + log.Debug().Msg("databricks: query did not need cancellation") + } + return nil, ctx.Err() + } + + return resp, err +} + +// pollOperation polls the operation status until it reaches a terminal state, +// cancelling the operation if the context is done (via the sentinel poll loop). +func (b *Backend) pollOperation(ctx context.Context, opHandle *cli_service.TOperationHandle) (*cli_service.TGetOperationStatusResp, error) { + corrId := driverctx.CorrelationIdFromContext(ctx) + opID := client.SprintGuid(opHandle.OperationId.GUID) + log := logger.WithContext(b.SessionID(), corrId, opID) + defer debuglog.Track(ctx, "thrift.Backend.pollOperation", "op=%s", opID)() + + var statusResp *cli_service.TGetOperationStatusResp + ctx = driverctx.NewContextWithConnId(ctx, b.SessionID()) + newCtx := context2.WithoutCancel(ctx) + pollSentinel := sentinel.Sentinel{ + OnDoneFn: func(statusResp any) (any, error) { + return statusResp, nil + }, + StatusFn: func() (sentinel.Done, any, error) { + var err error + log.Debug().Msg("databricks: polling status") + debuglog.Logf(ctx, "thrift.Backend.pollOperation", "GetOperationStatus") + statusResp, err = b.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 { + 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: + log.Debug().Msg("databricks: polling done") + return true + } + }, statusResp, err + }, + OnCancelFn: func() (any, error) { + log.Debug().Msg("databricks: sentinel canceling query") + debuglog.Logf(ctx, "thrift.Backend.pollOperation", "sentinel canceling query") + ret, err := b.client.CancelOperation(newCtx, &cli_service.TCancelOperationReq{ + OperationHandle: opHandle, + }) + return ret, err + }, + } + status, resp, err := pollSentinel.Watch(ctx, b.cfg.PollInterval, 0) + if err != nil { + log.Err(err).Msg("error polling operation status") + if status == sentinel.WatchTimeout { + err = dbsqlerrint.NewRequestError(ctx, dbsqlerr.ErrSentinelTimeout, err) + } + return nil, err + } + + statusResp, ok := resp.(*cli_service.TGetOperationStatusResp) + if !ok { + return nil, dbsqlerrint.NewDriverError(ctx, dbsqlerr.ErrReadQueryStatus, nil) + } + return statusResp, nil +} + +// --- operation-state helpers --- + +func unexpectedOperationState(opStatus *cli_service.TGetOperationStatusResp) error { + return errors.WithMessage(errors.New(opStatus.GetDisplayMessage()), dbsqlerr.ErrUnexpectedOperationState(opStatus.GetOperationState().String())) +} + +func invalidOperationState(ctx context.Context, opStatus *cli_service.TGetOperationStatusResp) error { + return dbsqlerrint.NewDriverError(ctx, dbsqlerr.ErrInvalidOperationState(opStatus.GetOperationState().String()), nil) +} + +func logBadQueryState(log *logger.DBSQLLogger, opStatus *cli_service.TGetOperationStatusResp) { + log.Error().Msgf("databricks: query state: %s", opStatus.GetOperationState()) + log.Error().Msg(opStatus.GetDisplayMessage()) + log.Debug().Msg(opStatus.GetDiagnosticInfo()) +} diff --git a/internal/backend/thrift/backend_test.go b/internal/backend/thrift/backend_test.go new file mode 100644 index 00000000..0ad13c3d --- /dev/null +++ b/internal/backend/thrift/backend_test.go @@ -0,0 +1,1237 @@ +package thrift + +// White-box tests for the Thrift execute/poll state machine. They drive the +// unexported executeStatement/pollOperation/runQuery against a mock TCLIService. + +import ( + "context" + "database/sql/driver" + "fmt" + "testing" + "time" + + "github.com/databricks/databricks-sql-go/driverctx" + "github.com/databricks/databricks-sql-go/internal/backend" + "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/thrift_protocol" + "github.com/stretchr/testify/assert" +) + +// newTestBackend builds a Backend around a mock client and canned session, +// bypassing New/OpenSession. +func newTestBackend(cli cli_service.TCLIService, session *cli_service.TOpenSessionResp, cfg *config.Config) *Backend { + b := &Backend{cfg: cfg, client: cli, session: session} + if session != nil && session.SessionHandle != nil { + b.sessionID = client.SprintGuid(session.SessionHandle.GetSessionId().GUID) + } + return b +} + +// namedToParams converts driver.NamedValues to backend.Params with a minimal +// string-typed mapping — enough for these tests, which only exercise param +// passthrough, not the full type inference (that lives in the dbsql package). +func namedToParams(values []driver.NamedValue) []backend.Param { + if len(values) == 0 { + return nil + } + params := make([]backend.Param, 0, len(values)) + for _, v := range values { + s := fmt.Sprintf("%v", v.Value) + val := s + params = append(params, backend.Param{Name: v.Name, Type: "STRING", Value: &val}) + } + return params +} + +func getTestSession() *cli_service.TOpenSessionResp { + return &cli_service.TOpenSessionResp{SessionHandle: &cli_service.TSessionHandle{ + SessionId: &cli_service.THandleIdentifier{ + GUID: []byte{1, 2, 3, 4, 2, 23, 4, 2, 3, 2, 3, 4, 4, 223, 34, 54}, + }, + }} +} + +func strPtr(s string) *string { return &s } + +func TestBackend_executeStatement(t *testing.T) { + t.Parallel() + t.Run("executeStatement should err when client.ExecuteStatement fails", func(t *testing.T) { + var executeStatementCount int + executeStatement := func(ctx context.Context, req *cli_service.TExecuteStatementReq) (r *cli_service.TExecuteStatementResp, err error) { + executeStatementCount++ + return nil, fmt.Errorf("error") + } + testClient := &client.TestClient{ + FnExecuteStatement: executeStatement, + } + be := newTestBackend(testClient, getTestSession(), config.WithDefaults()) + _, err := be.executeStatement(context.Background(), backend.ExecRequest{Query: "select 1", Params: namedToParams([]driver.NamedValue{})}) + assert.Error(t, err) + assert.Equal(t, 1, executeStatementCount) + }) + + t.Run("executeStatement should return TExecuteStatementResp on success", func(t *testing.T) { + var executeStatementCount int + executeStatement := func(ctx context.Context, req *cli_service.TExecuteStatementReq) (r *cli_service.TExecuteStatementResp, err error) { + executeStatementCount++ + executeStatementResp := &cli_service.TExecuteStatementResp{ + Status: &cli_service.TStatus{ + StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, + }, + OperationHandle: &cli_service.TOperationHandle{ + OperationId: &cli_service.THandleIdentifier{ + GUID: []byte{1, 2, 3, 4, 2, 23, 4, 2, 3, 1, 2, 3, 4, 4, 223, 34, 54}, + Secret: []byte("b"), + }, + }, + DirectResults: &cli_service.TSparkDirectResults{ + OperationStatus: &cli_service.TGetOperationStatusResp{ + Status: &cli_service.TStatus{ + StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, + }, + OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_ERROR_STATE), + ErrorMessage: strPtr("error message"), + DisplayMessage: strPtr("display message"), + }, + ResultSetMetadata: &cli_service.TGetResultSetMetadataResp{ + Status: &cli_service.TStatus{ + StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, + }, + }, + ResultSet: &cli_service.TFetchResultsResp{ + Status: &cli_service.TStatus{ + StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, + }, + }, + }, + } + return executeStatementResp, nil + } + testClient := &client.TestClient{ + FnExecuteStatement: executeStatement, + } + be := newTestBackend(testClient, getTestSession(), config.WithDefaults()) + _, err := be.executeStatement(context.Background(), backend.ExecRequest{Query: "select 1", Params: namedToParams([]driver.NamedValue{})}) + + assert.NoError(t, err) + assert.Equal(t, 1, executeStatementCount) + }) + + t.Run("executeStatement should not call cancel if not needed", func(t *testing.T) { + var executeStatementCount int + var cancelOperationCount int + var cancel context.CancelFunc + executeStatement := func(ctx context.Context, req *cli_service.TExecuteStatementReq) (r *cli_service.TExecuteStatementResp, err error) { + executeStatementCount++ + cancel() + executeStatementResp := &cli_service.TExecuteStatementResp{ + Status: &cli_service.TStatus{ + StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, + }, + OperationHandle: &cli_service.TOperationHandle{ + OperationId: &cli_service.THandleIdentifier{ + GUID: []byte{1, 2, 3, 4, 2, 23, 4, 2, 3, 1, 2, 3, 4, 4, 223, 34, 54}, + Secret: []byte("b"), + }, + }, + DirectResults: &cli_service.TSparkDirectResults{ + OperationStatus: &cli_service.TGetOperationStatusResp{ + Status: &cli_service.TStatus{ + StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, + }, + OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_FINISHED_STATE), + ErrorMessage: strPtr("error message"), + DisplayMessage: strPtr("display message"), + }, + ResultSetMetadata: &cli_service.TGetResultSetMetadataResp{ + Status: &cli_service.TStatus{ + StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, + }, + }, + ResultSet: &cli_service.TFetchResultsResp{ + Status: &cli_service.TStatus{ + StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, + }, + }, + CloseOperation: &cli_service.TCloseOperationResp{ + Status: &cli_service.TStatus{ + StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, + }, + }, + }, + } + return executeStatementResp, nil + } + cancelOperation := func(ctx context.Context, req *cli_service.TCancelOperationReq) (r *cli_service.TCancelOperationResp, err error) { + cancelOperationCount++ + cancelOperationResp := &cli_service.TCancelOperationResp{ + Status: &cli_service.TStatus{ + StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, + }, + } + return cancelOperationResp, nil + } + testClient := &client.TestClient{ + FnExecuteStatement: executeStatement, + FnCancelOperation: cancelOperation, + } + be := newTestBackend(testClient, getTestSession(), config.WithDefaults()) + + ctx := context.Background() + ctx, cancel = context.WithCancel(ctx) + defer cancel() + _, err := be.executeStatement(ctx, backend.ExecRequest{Query: "select 1", Params: namedToParams([]driver.NamedValue{})}) + + assert.Error(t, err) + assert.Equal(t, 1, executeStatementCount) + assert.Equal(t, 0, cancelOperationCount) + }) + t.Run("executeStatement should call cancel if needed", func(t *testing.T) { + var executeStatementCount int + var cancelOperationCount int + var cancel context.CancelFunc + executeStatement := func(ctx context.Context, req *cli_service.TExecuteStatementReq) (r *cli_service.TExecuteStatementResp, err error) { + executeStatementCount++ + cancel() + executeStatementResp := &cli_service.TExecuteStatementResp{ + Status: &cli_service.TStatus{ + StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, + }, + OperationHandle: &cli_service.TOperationHandle{ + OperationId: &cli_service.THandleIdentifier{ + GUID: []byte{1, 2, 3, 4, 2, 23, 4, 2, 3, 1, 2, 3, 4, 4, 223, 34, 54}, + Secret: []byte("b"), + }, + }, + DirectResults: &cli_service.TSparkDirectResults{ + OperationStatus: &cli_service.TGetOperationStatusResp{ + Status: &cli_service.TStatus{ + StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, + }, + OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_FINISHED_STATE), + ErrorMessage: strPtr("error message"), + DisplayMessage: strPtr("display message"), + }, + ResultSetMetadata: &cli_service.TGetResultSetMetadataResp{ + Status: &cli_service.TStatus{ + StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, + }, + }, + ResultSet: &cli_service.TFetchResultsResp{ + Status: &cli_service.TStatus{ + StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, + }, + }, + }, + } + return executeStatementResp, nil + } + cancelOperation := func(ctx context.Context, req *cli_service.TCancelOperationReq) (r *cli_service.TCancelOperationResp, err error) { + cancelOperationCount++ + cancelOperationResp := &cli_service.TCancelOperationResp{ + Status: &cli_service.TStatus{ + StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, + }, + } + return cancelOperationResp, nil + } + testClient := &client.TestClient{ + FnExecuteStatement: executeStatement, + FnCancelOperation: cancelOperation, + } + be := newTestBackend(testClient, getTestSession(), config.WithDefaults()) + ctx := context.Background() + ctx, cancel = context.WithCancel(ctx) + defer cancel() + _, err := be.executeStatement(ctx, backend.ExecRequest{Query: "select 1", Params: namedToParams([]driver.NamedValue{})}) + + assert.Error(t, err) + assert.Equal(t, 1, executeStatementCount) + assert.Equal(t, 1, cancelOperationCount) + }) + +} + +func TestBackend_executeStatement_ProtocolFeatures(t *testing.T) { + t.Parallel() + + protocols := []cli_service.TProtocolVersion{ + cli_service.TProtocolVersion_SPARK_CLI_SERVICE_PROTOCOL_V1, + cli_service.TProtocolVersion_SPARK_CLI_SERVICE_PROTOCOL_V2, + cli_service.TProtocolVersion_SPARK_CLI_SERVICE_PROTOCOL_V3, + cli_service.TProtocolVersion_SPARK_CLI_SERVICE_PROTOCOL_V4, + cli_service.TProtocolVersion_SPARK_CLI_SERVICE_PROTOCOL_V5, + cli_service.TProtocolVersion_SPARK_CLI_SERVICE_PROTOCOL_V6, + cli_service.TProtocolVersion_SPARK_CLI_SERVICE_PROTOCOL_V7, + cli_service.TProtocolVersion_SPARK_CLI_SERVICE_PROTOCOL_V8, + } + + testCases := []struct { + cfg *config.Config + supportsDirectResults func(version cli_service.TProtocolVersion) bool + supportsLz4Compression func(version cli_service.TProtocolVersion) bool + supportsCloudFetch func(version cli_service.TProtocolVersion) bool + supportsArrow func(version cli_service.TProtocolVersion) bool + supportsParameterizedQueries func(version cli_service.TProtocolVersion) bool + hasParameters bool + }{ + { + cfg: func() *config.Config { + cfg := config.WithDefaults() + cfg.UseLz4Compression = true + cfg.UseCloudFetch = true + cfg.UseArrowBatches = true + cfg.UseArrowNativeDecimal = true + cfg.UseArrowNativeTimestamp = true + cfg.UseArrowNativeComplexTypes = true + cfg.UseArrowNativeIntervalTypes = true + return cfg + }(), + supportsDirectResults: thrift_protocol.SupportsDirectResults, + supportsLz4Compression: thrift_protocol.SupportsLz4Compression, + supportsCloudFetch: thrift_protocol.SupportsCloudFetch, + supportsArrow: thrift_protocol.SupportsArrow, + supportsParameterizedQueries: thrift_protocol.SupportsParameterizedQueries, + hasParameters: true, + }, + { + cfg: func() *config.Config { + cfg := config.WithDefaults() + cfg.UseLz4Compression = false + cfg.UseCloudFetch = false + cfg.UseArrowBatches = false + return cfg + }(), + supportsDirectResults: thrift_protocol.SupportsDirectResults, + supportsLz4Compression: thrift_protocol.SupportsLz4Compression, + supportsCloudFetch: thrift_protocol.SupportsCloudFetch, + supportsArrow: thrift_protocol.SupportsArrow, + supportsParameterizedQueries: thrift_protocol.SupportsParameterizedQueries, + hasParameters: false, + }, + } + + for _, tc := range testCases { + for _, version := range protocols { + t.Run(fmt.Sprintf("protocol_v%d_withParams_%v", version, tc.hasParameters), func(t *testing.T) { + var capturedReq *cli_service.TExecuteStatementReq + executeStatement := func(ctx context.Context, req *cli_service.TExecuteStatementReq) (r *cli_service.TExecuteStatementResp, err error) { + capturedReq = req + executeStatementResp := &cli_service.TExecuteStatementResp{ + Status: &cli_service.TStatus{ + StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, + }, + OperationHandle: &cli_service.TOperationHandle{ + OperationId: &cli_service.THandleIdentifier{ + GUID: []byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}, + Secret: []byte("secret"), + }, + }, + DirectResults: &cli_service.TSparkDirectResults{ + OperationStatus: &cli_service.TGetOperationStatusResp{ + Status: &cli_service.TStatus{ + StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, + }, + OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_FINISHED_STATE), + }, + }, + } + return executeStatementResp, nil + } + + session := getTestSession() + session.ServerProtocolVersion = version + + testClient := &client.TestClient{ + FnExecuteStatement: executeStatement, + } + + be := newTestBackend(testClient, session, tc.cfg) + + var args []driver.NamedValue + if tc.hasParameters { + args = []driver.NamedValue{ + {Name: "param1", Value: "value1"}, + } + } + + _, err := be.executeStatement(context.Background(), backend.ExecRequest{Query: "SELECT 1", Params: namedToParams(args)}) + assert.NoError(t, err) + + // Verify direct results + hasDirectResults := tc.supportsDirectResults(version) + assert.Equal(t, hasDirectResults, capturedReq.GetDirectResults != nil, "Direct results should be enabled if protocol supports it") + + // Verify LZ4 compression + shouldHaveLz4 := tc.supportsLz4Compression(version) && tc.cfg.UseLz4Compression + if shouldHaveLz4 { + assert.NotNil(t, capturedReq.CanDecompressLZ4Result_) + assert.True(t, *capturedReq.CanDecompressLZ4Result_) + } else { + assert.Nil(t, capturedReq.CanDecompressLZ4Result_) + } + + // Verify cloud fetch + shouldHaveCloudFetch := tc.supportsCloudFetch(version) && tc.cfg.UseCloudFetch + if shouldHaveCloudFetch { + assert.NotNil(t, capturedReq.CanDownloadResult_) + assert.True(t, *capturedReq.CanDownloadResult_) + } else { + assert.Nil(t, capturedReq.CanDownloadResult_) + } + + // Verify Arrow support + shouldHaveArrow := tc.supportsArrow(version) && tc.cfg.UseArrowBatches + if shouldHaveArrow { + assert.NotNil(t, capturedReq.CanReadArrowResult_) + assert.True(t, *capturedReq.CanReadArrowResult_) + assert.NotNil(t, capturedReq.UseArrowNativeTypes) + assert.Equal(t, tc.cfg.UseArrowNativeDecimal, *capturedReq.UseArrowNativeTypes.DecimalAsArrow) + assert.Equal(t, tc.cfg.UseArrowNativeTimestamp, *capturedReq.UseArrowNativeTypes.TimestampAsArrow) + assert.Equal(t, tc.cfg.UseArrowNativeComplexTypes, *capturedReq.UseArrowNativeTypes.ComplexTypesAsArrow) + assert.Equal(t, tc.cfg.UseArrowNativeIntervalTypes, *capturedReq.UseArrowNativeTypes.IntervalTypesAsArrow) + } else { + assert.Nil(t, capturedReq.CanReadArrowResult_) + assert.Nil(t, capturedReq.UseArrowNativeTypes) + } + + // Verify parameters + shouldHaveParams := tc.supportsParameterizedQueries(version) && tc.hasParameters + if shouldHaveParams { + assert.NotNil(t, capturedReq.Parameters) + assert.Len(t, capturedReq.Parameters, 1) + } else if tc.hasParameters { + // Even if we have parameters but protocol doesn't support it, we shouldn't set them + assert.Nil(t, capturedReq.Parameters) + } + }) + } + } +} + +func TestBackend_executeStatement_QueryTags(t *testing.T) { + t.Parallel() + + makeTestBackend := func(captureReq *(*cli_service.TExecuteStatementReq)) *Backend { + executeStatement := func(ctx context.Context, req *cli_service.TExecuteStatementReq) (r *cli_service.TExecuteStatementResp, err error) { + *captureReq = req + return &cli_service.TExecuteStatementResp{ + Status: &cli_service.TStatus{ + StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, + }, + OperationHandle: &cli_service.TOperationHandle{ + OperationId: &cli_service.THandleIdentifier{ + GUID: []byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}, + Secret: []byte("secret"), + }, + }, + DirectResults: &cli_service.TSparkDirectResults{ + OperationStatus: &cli_service.TGetOperationStatusResp{ + Status: &cli_service.TStatus{ + StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, + }, + OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_FINISHED_STATE), + }, + }, + }, nil + } + + return newTestBackend( + &client.TestClient{FnExecuteStatement: executeStatement}, + getTestSession(), + config.WithDefaults(), + ) + } + + t.Run("query tags from context are set in ConfOverlay", func(t *testing.T) { + var capturedReq *cli_service.TExecuteStatementReq + be := makeTestBackend(&capturedReq) + + ctx := driverctx.NewContextWithQueryTags(context.Background(), map[string]string{ + "team": "engineering", + "app": "etl", + }) + + _, err := be.executeStatement(ctx, backend.ExecRequest{Query: "SELECT 1", Params: namedToParams(nil)}) + assert.NoError(t, err) + assert.NotNil(t, capturedReq.ConfOverlay) + // Map iteration is non-deterministic, so check both possible orderings + queryTags := capturedReq.ConfOverlay["query_tags"] + assert.True(t, + queryTags == "team:engineering,app:etl" || queryTags == "app:etl,team:engineering", + "unexpected query_tags value: %s", queryTags) + }) + + t.Run("no query tags in context means no ConfOverlay", func(t *testing.T) { + var capturedReq *cli_service.TExecuteStatementReq + be := makeTestBackend(&capturedReq) + + _, err := be.executeStatement(context.Background(), backend.ExecRequest{Query: "SELECT 1", Params: namedToParams(nil)}) + assert.NoError(t, err) + assert.Nil(t, capturedReq.ConfOverlay) + }) + + t.Run("empty query tags map means no ConfOverlay", func(t *testing.T) { + var capturedReq *cli_service.TExecuteStatementReq + be := makeTestBackend(&capturedReq) + + ctx := driverctx.NewContextWithQueryTags(context.Background(), map[string]string{}) + + _, err := be.executeStatement(ctx, backend.ExecRequest{Query: "SELECT 1", Params: namedToParams(nil)}) + assert.NoError(t, err) + assert.Nil(t, capturedReq.ConfOverlay) + }) + + t.Run("single query tag", func(t *testing.T) { + var capturedReq *cli_service.TExecuteStatementReq + be := makeTestBackend(&capturedReq) + + ctx := driverctx.NewContextWithQueryTags(context.Background(), map[string]string{ + "team": "data-eng", + }) + + _, err := be.executeStatement(ctx, backend.ExecRequest{Query: "SELECT 1", Params: namedToParams(nil)}) + assert.NoError(t, err) + assert.Equal(t, "team:data-eng", capturedReq.ConfOverlay["query_tags"]) + }) + + t.Run("query tags with special characters in values", func(t *testing.T) { + var capturedReq *cli_service.TExecuteStatementReq + be := makeTestBackend(&capturedReq) + + ctx := driverctx.NewContextWithQueryTags(context.Background(), map[string]string{ + "url": "http://host:8080", + }) + + _, err := be.executeStatement(ctx, backend.ExecRequest{Query: "SELECT 1", Params: namedToParams(nil)}) + assert.NoError(t, err) + assert.Equal(t, `url:http\://host\:8080`, capturedReq.ConfOverlay["query_tags"]) + }) + + t.Run("query tags with empty value", func(t *testing.T) { + var capturedReq *cli_service.TExecuteStatementReq + be := makeTestBackend(&capturedReq) + + ctx := driverctx.NewContextWithQueryTags(context.Background(), map[string]string{ + "flag": "", + }) + + _, err := be.executeStatement(ctx, backend.ExecRequest{Query: "SELECT 1", Params: namedToParams(nil)}) + assert.NoError(t, err) + assert.Equal(t, "flag", capturedReq.ConfOverlay["query_tags"]) + }) + + t.Run("session-level and statement-level query tags coexist", func(t *testing.T) { + // Session-level tags are sent via TOpenSessionReq.Configuration at connect time. + // Statement-level tags are sent via TExecuteStatementReq.ConfOverlay at query time. + // They are independent fields on different requests, so both should work together. + + var capturedOpenReq *cli_service.TOpenSessionReq + var capturedExecReq *cli_service.TExecuteStatementReq + + testClient := &client.TestClient{ + FnOpenSession: func(ctx context.Context, req *cli_service.TOpenSessionReq) (*cli_service.TOpenSessionResp, error) { + capturedOpenReq = req + return &cli_service.TOpenSessionResp{ + Status: &cli_service.TStatus{ + StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, + }, + SessionHandle: &cli_service.TSessionHandle{ + SessionId: &cli_service.THandleIdentifier{ + GUID: []byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}, + }, + }, + }, nil + }, + FnExecuteStatement: func(ctx context.Context, req *cli_service.TExecuteStatementReq) (*cli_service.TExecuteStatementResp, error) { + capturedExecReq = req + return &cli_service.TExecuteStatementResp{ + Status: &cli_service.TStatus{ + StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, + }, + OperationHandle: &cli_service.TOperationHandle{ + OperationId: &cli_service.THandleIdentifier{ + GUID: []byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}, + Secret: []byte("secret"), + }, + }, + DirectResults: &cli_service.TSparkDirectResults{ + OperationStatus: &cli_service.TGetOperationStatusResp{ + Status: &cli_service.TStatus{ + StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, + }, + OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_FINISHED_STATE), + }, + }, + }, nil + }, + } + + // Simulate what connector.Connect() does: pass session params to OpenSession + sessionParams := map[string]string{ + "QUERY_TAGS": "team:platform,env:prod", + "ansi_mode": "false", + } + protocolVersion := int64(cli_service.TProtocolVersion_SPARK_CLI_SERVICE_PROTOCOL_V8) + session, err := testClient.OpenSession(context.Background(), &cli_service.TOpenSessionReq{ + ClientProtocolI64: &protocolVersion, + Configuration: sessionParams, + }) + assert.NoError(t, err) + + // Verify session-level tags were sent in OpenSession + assert.Equal(t, "team:platform,env:prod", capturedOpenReq.Configuration["QUERY_TAGS"]) + assert.Equal(t, "false", capturedOpenReq.Configuration["ansi_mode"]) + + // Create conn with session that has session-level tags + cfg := config.WithDefaults() + cfg.SessionParams = sessionParams + be := newTestBackend(testClient, session, cfg) + + // Execute with statement-level tags + ctx := driverctx.NewContextWithQueryTags(context.Background(), map[string]string{ + "job": "nightly-etl", + }) + _, err = be.executeStatement(ctx, backend.ExecRequest{Query: "SELECT 1", Params: namedToParams(nil)}) + assert.NoError(t, err) + + // Statement-level tags should be in ConfOverlay + assert.Equal(t, "job:nightly-etl", capturedExecReq.ConfOverlay["query_tags"]) + + // ConfOverlay should ONLY have query_tags, not session params + _, hasAnsiMode := capturedExecReq.ConfOverlay["ansi_mode"] + assert.False(t, hasAnsiMode, "session params should not leak into ConfOverlay") + _, hasSessionQueryTags := capturedExecReq.ConfOverlay["QUERY_TAGS"] + assert.False(t, hasSessionQueryTags, "session-level QUERY_TAGS should not be in ConfOverlay") + }) +} + +func TestBackend_pollOperation(t *testing.T) { + t.Parallel() + t.Run("pollOperation returns finished state response when query finishes", func(t *testing.T) { + var getOperationStatusCount int + getOperationStatus := func(ctx context.Context, req *cli_service.TGetOperationStatusReq) (r *cli_service.TGetOperationStatusResp, err error) { + getOperationStatusCount++ + getOperationStatusResp := &cli_service.TGetOperationStatusResp{ + OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_FINISHED_STATE), + } + return getOperationStatusResp, nil + } + testClient := &client.TestClient{ + FnGetOperationStatus: getOperationStatus, + } + be := newTestBackend(testClient, getTestSession(), config.WithDefaults()) + res, err := be.pollOperation(context.Background(), &cli_service.TOperationHandle{ + OperationId: &cli_service.THandleIdentifier{ + GUID: []byte{1, 2, 3, 4, 2, 23, 4, 2, 3, 2, 4, 7, 8, 223, 34, 54}, + Secret: []byte("b"), + }, + }) + assert.NoError(t, err) + assert.Equal(t, 1, getOperationStatusCount) + assert.Equal(t, cli_service.TGetOperationStatusResp{ + OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_FINISHED_STATE), + }, *res) + }) + + t.Run("pollOperation returns closed state response when query has been closed", func(t *testing.T) { + var getOperationStatusCount int + getOperationStatus := func(ctx context.Context, req *cli_service.TGetOperationStatusReq) (r *cli_service.TGetOperationStatusResp, err error) { + getOperationStatusCount++ + getOperationStatusResp := &cli_service.TGetOperationStatusResp{ + OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_CLOSED_STATE), + } + return getOperationStatusResp, nil + } + testClient := &client.TestClient{ + FnGetOperationStatus: getOperationStatus, + } + be := newTestBackend(testClient, getTestSession(), config.WithDefaults()) + res, err := be.pollOperation(context.Background(), &cli_service.TOperationHandle{ + OperationId: &cli_service.THandleIdentifier{ + GUID: []byte{1, 2, 3, 4, 2, 23, 4, 2, 3, 2, 3, 4, 4, 223, 34, 54}, + Secret: []byte("b"), + }, + }) + assert.NoError(t, err) + assert.Equal(t, 1, getOperationStatusCount) + assert.Equal(t, cli_service.TGetOperationStatusResp{ + OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_CLOSED_STATE), + }, *res) + }) + + t.Run("pollOperation returns closed state response when query has been closed", func(t *testing.T) { + var getOperationStatusCount int + getOperationStatus := func(ctx context.Context, req *cli_service.TGetOperationStatusReq) (r *cli_service.TGetOperationStatusResp, err error) { + getOperationStatusCount++ + getOperationStatusResp := &cli_service.TGetOperationStatusResp{ + OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_CLOSED_STATE), + } + return getOperationStatusResp, nil + } + testClient := &client.TestClient{ + FnGetOperationStatus: getOperationStatus, + } + be := newTestBackend(testClient, getTestSession(), config.WithDefaults()) + res, err := be.pollOperation(context.Background(), &cli_service.TOperationHandle{ + OperationId: &cli_service.THandleIdentifier{ + GUID: []byte{1, 2, 3, 4, 2, 23, 4, 2, 3, 2, 3, 4, 4, 223, 34, 54}, + Secret: []byte("b"), + }, + }) + assert.NoError(t, err) + assert.Equal(t, 1, getOperationStatusCount) + assert.Equal(t, cli_service.TGetOperationStatusResp{ + OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_CLOSED_STATE), + }, *res) + }) + + t.Run("pollOperation returns unknown state response when query state is unknown", func(t *testing.T) { + var getOperationStatusCount int + getOperationStatus := func(ctx context.Context, req *cli_service.TGetOperationStatusReq) (r *cli_service.TGetOperationStatusResp, err error) { + getOperationStatusCount++ + getOperationStatusResp := &cli_service.TGetOperationStatusResp{ + OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_UKNOWN_STATE), + } + return getOperationStatusResp, nil + } + testClient := &client.TestClient{ + FnGetOperationStatus: getOperationStatus, + } + be := newTestBackend(testClient, getTestSession(), config.WithDefaults()) + res, err := be.pollOperation(context.Background(), &cli_service.TOperationHandle{ + OperationId: &cli_service.THandleIdentifier{ + GUID: []byte{1, 2, 3, 4, 2, 23, 4, 2, 3, 1, 2, 3, 4, 223, 34, 54}, + Secret: []byte("b"), + }, + }) + assert.NoError(t, err) + assert.Equal(t, 1, getOperationStatusCount) + assert.Equal(t, cli_service.TGetOperationStatusResp{ + OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_UKNOWN_STATE), + }, *res) + }) + + t.Run("pollOperation returns error state response when query errors", func(t *testing.T) { + var getOperationStatusCount int + getOperationStatus := func(ctx context.Context, req *cli_service.TGetOperationStatusReq) (r *cli_service.TGetOperationStatusResp, err error) { + getOperationStatusCount++ + getOperationStatusResp := &cli_service.TGetOperationStatusResp{ + OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_ERROR_STATE), + } + return getOperationStatusResp, nil + } + testClient := &client.TestClient{ + FnGetOperationStatus: getOperationStatus, + } + be := newTestBackend(testClient, getTestSession(), config.WithDefaults()) + res, err := be.pollOperation(context.Background(), &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"), + }, + }) + assert.NoError(t, err) + assert.Equal(t, 1, getOperationStatusCount) + assert.Equal(t, cli_service.TGetOperationStatusResp{ + OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_ERROR_STATE), + }, *res) + }) + + t.Run("pollOperation returns finished state response after query cycles through various states", func(t *testing.T) { + var getOperationStatusCount int + getOperationStatus := func(ctx context.Context, req *cli_service.TGetOperationStatusReq) (r *cli_service.TGetOperationStatusResp, err error) { + operationStates := [4]cli_service.TOperationState{cli_service.TOperationState_INITIALIZED_STATE, cli_service.TOperationState_PENDING_STATE, + cli_service.TOperationState_RUNNING_STATE, cli_service.TOperationState_FINISHED_STATE} + getOperationStatusCount++ + getOperationStatusResp := &cli_service.TGetOperationStatusResp{ + OperationState: cli_service.TOperationStatePtr(operationStates[getOperationStatusCount-1]), + } + return getOperationStatusResp, nil + } + testClient := &client.TestClient{ + FnGetOperationStatus: getOperationStatus, + } + be := newTestBackend(testClient, getTestSession(), config.WithDefaults()) + res, err := be.pollOperation(context.Background(), &cli_service.TOperationHandle{ + OperationId: &cli_service.THandleIdentifier{ + GUID: []byte{1, 2, 3, 4, 2, 23, 4, 2, 3, 1, 3, 4, 4, 223, 34, 54}, + Secret: []byte("b"), + }, + }) + assert.NoError(t, err) + assert.Equal(t, 4, getOperationStatusCount) + assert.Equal(t, cli_service.TGetOperationStatusResp{ + OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_FINISHED_STATE), + }, *res) + }) + + t.Run("pollOperation returns cancel err when context times out before get operation", func(t *testing.T) { + var getOperationStatusCount, cancelOperationCount int + getOperationStatus := func(ctx context.Context, req *cli_service.TGetOperationStatusReq) (r *cli_service.TGetOperationStatusResp, err error) { + operationStates := [4]cli_service.TOperationState{cli_service.TOperationState_INITIALIZED_STATE, cli_service.TOperationState_PENDING_STATE, + cli_service.TOperationState_RUNNING_STATE, cli_service.TOperationState_FINISHED_STATE} + getOperationStatusCount++ + getOperationStatusResp := &cli_service.TGetOperationStatusResp{ + OperationState: cli_service.TOperationStatePtr(operationStates[getOperationStatusCount-1]), + } + return getOperationStatusResp, nil + } + cancelOperation := func(ctx context.Context, req *cli_service.TCancelOperationReq) (r *cli_service.TCancelOperationResp, err error) { + cancelOperationCount++ + cancelOperationResp := &cli_service.TCancelOperationResp{ + Status: &cli_service.TStatus{ + StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, + }, + } + return cancelOperationResp, nil + } + testClient := &client.TestClient{ + FnGetOperationStatus: getOperationStatus, + FnCancelOperation: cancelOperation, + } + be := newTestBackend(testClient, getTestSession(), config.WithDefaults()) + ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + defer cancel() + res, err := be.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"), + }, + }) + assert.Error(t, err) + assert.Equal(t, 0, getOperationStatusCount) + assert.Equal(t, 1, cancelOperationCount) + assert.Nil(t, res) + }) + + t.Run("pollOperation returns cancel err when context times out before get operation", func(t *testing.T) { + var getOperationStatusCount, cancelOperationCount int + getOperationStatus := func(ctx context.Context, req *cli_service.TGetOperationStatusReq) (r *cli_service.TGetOperationStatusResp, err error) { + operationStates := [4]cli_service.TOperationState{cli_service.TOperationState_INITIALIZED_STATE, cli_service.TOperationState_PENDING_STATE, + cli_service.TOperationState_RUNNING_STATE, cli_service.TOperationState_FINISHED_STATE} + getOperationStatusCount++ + getOperationStatusResp := &cli_service.TGetOperationStatusResp{ + OperationState: cli_service.TOperationStatePtr(operationStates[getOperationStatusCount-1]), + } + return getOperationStatusResp, nil + } + cancelOperation := func(ctx context.Context, req *cli_service.TCancelOperationReq) (r *cli_service.TCancelOperationResp, err error) { + cancelOperationCount++ + cancelOperationResp := &cli_service.TCancelOperationResp{ + Status: &cli_service.TStatus{ + StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, + }, + } + return cancelOperationResp, nil + } + testClient := &client.TestClient{ + FnGetOperationStatus: getOperationStatus, + FnCancelOperation: cancelOperation, + } + cfg := config.WithDefaults() + cfg.PollInterval = 100 * time.Millisecond + be := newTestBackend(testClient, getTestSession(), cfg) + ctx, cancel := context.WithTimeout(context.Background(), 250*time.Millisecond) + defer cancel() + res, err := be.pollOperation(ctx, &cli_service.TOperationHandle{ + OperationId: &cli_service.THandleIdentifier{ + GUID: []byte{1, 2, 3, 4, 2, 23, 4, 2, 3, 1, 2, 3, 4, 223, 34, 54}, + Secret: []byte("b"), + }, + }) + + assert.Error(t, err) + assert.GreaterOrEqual(t, getOperationStatusCount, 1) + assert.Equal(t, 1, cancelOperationCount) + assert.Nil(t, res) + }) + + t.Run("pollOperation returns cancel err when context is cancelled", func(t *testing.T) { + var getOperationStatusCount, cancelOperationCount int + getOperationStatus := func(ctx context.Context, req *cli_service.TGetOperationStatusReq) (r *cli_service.TGetOperationStatusResp, err error) { + operationStates := [4]cli_service.TOperationState{cli_service.TOperationState_INITIALIZED_STATE, cli_service.TOperationState_PENDING_STATE, + cli_service.TOperationState_RUNNING_STATE, cli_service.TOperationState_FINISHED_STATE} + getOperationStatusCount++ + getOperationStatusResp := &cli_service.TGetOperationStatusResp{ + OperationState: cli_service.TOperationStatePtr(operationStates[getOperationStatusCount-1]), + } + return getOperationStatusResp, nil + } + cancelOperation := func(ctx context.Context, req *cli_service.TCancelOperationReq) (r *cli_service.TCancelOperationResp, err error) { + cancelOperationCount++ + cancelOperationResp := &cli_service.TCancelOperationResp{ + Status: &cli_service.TStatus{ + StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, + }, + } + return cancelOperationResp, nil + } + testClient := &client.TestClient{ + FnGetOperationStatus: getOperationStatus, + FnCancelOperation: cancelOperation, + } + cfg := config.WithDefaults() + cfg.PollInterval = 100 * time.Millisecond + be := newTestBackend(testClient, getTestSession(), cfg) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + go func() { + time.Sleep(150 * time.Millisecond) + cancel() + }() + res, err := be.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"), + }, + }) + assert.Error(t, err) + assert.GreaterOrEqual(t, getOperationStatusCount, 1) + assert.GreaterOrEqual(t, 1, cancelOperationCount) + assert.Nil(t, res) + }) +} + +func TestBackend_runQuery(t *testing.T) { + t.Parallel() + t.Run("runQuery should err when client.ExecuteStatement fails", func(t *testing.T) { + var executeStatementCount int + executeStatement := func(ctx context.Context, req *cli_service.TExecuteStatementReq) (r *cli_service.TExecuteStatementResp, err error) { + executeStatementCount++ + return nil, fmt.Errorf("error") + } + testClient := &client.TestClient{ + FnExecuteStatement: executeStatement, + } + be := newTestBackend(testClient, getTestSession(), config.WithDefaults()) + exStmtResp, opStatusResp, err := be.runQuery(context.Background(), backend.ExecRequest{Query: "select 1", Params: namedToParams([]driver.NamedValue{})}) + assert.Error(t, err) + assert.Nil(t, exStmtResp) + assert.Nil(t, opStatusResp) + assert.Equal(t, 1, executeStatementCount) + }) + + t.Run("runQuery should err when pollOperation fails", func(t *testing.T) { + var executeStatementCount, getOperationStatusCount int + executeStatement := func(ctx context.Context, req *cli_service.TExecuteStatementReq) (r *cli_service.TExecuteStatementResp, err error) { + executeStatementCount++ + executeStatementResp := &cli_service.TExecuteStatementResp{ + Status: &cli_service.TStatus{ + StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, + }, + OperationHandle: &cli_service.TOperationHandle{ + OperationId: &cli_service.THandleIdentifier{ + GUID: []byte{1, 2, 3, 4, 2, 23, 4, 2, 3, 1, 3, 4, 4, 223, 34, 54}, + Secret: []byte("b"), + }, + }, + } + return executeStatementResp, nil + } + + getOperationStatus := func(ctx context.Context, req *cli_service.TGetOperationStatusReq) (r *cli_service.TGetOperationStatusResp, err error) { + getOperationStatusCount++ + getOperationStatusResp := &cli_service.TGetOperationStatusResp{ + OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_ERROR_STATE), + } + return getOperationStatusResp, fmt.Errorf("error on get operation status") + } + + testClient := &client.TestClient{ + FnExecuteStatement: executeStatement, + FnGetOperationStatus: getOperationStatus, + } + be := newTestBackend(testClient, getTestSession(), config.WithDefaults()) + exStmtResp, opStatusResp, err := be.runQuery(context.Background(), backend.ExecRequest{Query: "select 1", Params: namedToParams([]driver.NamedValue{})}) + + assert.Error(t, err) + assert.Equal(t, 1, executeStatementCount) + assert.Equal(t, 1, getOperationStatusCount) + assert.NotNil(t, exStmtResp) + assert.Nil(t, opStatusResp) + }) + + t.Run("runQuery should return resp when query is finished", func(t *testing.T) { + var executeStatementCount, getOperationStatusCount int + executeStatement := func(ctx context.Context, req *cli_service.TExecuteStatementReq) (r *cli_service.TExecuteStatementResp, err error) { + executeStatementCount++ + executeStatementResp := &cli_service.TExecuteStatementResp{ + Status: &cli_service.TStatus{ + StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, + }, + OperationHandle: &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"), + }, + }, + } + return executeStatementResp, nil + } + var numModRows int64 = 2 + + getOperationStatus := func(ctx context.Context, req *cli_service.TGetOperationStatusReq) (r *cli_service.TGetOperationStatusResp, err error) { + getOperationStatusCount++ + getOperationStatusResp := &cli_service.TGetOperationStatusResp{ + OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_FINISHED_STATE), + NumModifiedRows: &numModRows, + } + return getOperationStatusResp, nil + } + + testClient := &client.TestClient{ + FnExecuteStatement: executeStatement, + FnGetOperationStatus: getOperationStatus, + } + be := newTestBackend(testClient, getTestSession(), config.WithDefaults()) + exStmtResp, opStatusResp, err := be.runQuery(context.Background(), backend.ExecRequest{Query: "select 1", Params: namedToParams([]driver.NamedValue{})}) + + assert.NoError(t, err) + assert.Equal(t, 1, executeStatementCount) + assert.Equal(t, 1, getOperationStatusCount) + assert.NotNil(t, exStmtResp) + assert.NotNil(t, opStatusResp) + assert.Equal(t, &numModRows, opStatusResp.NumModifiedRows) + }) + + t.Run("runQuery should return resp and error when query is canceled", func(t *testing.T) { + var executeStatementCount, getOperationStatusCount int + executeStatement := func(ctx context.Context, req *cli_service.TExecuteStatementReq) (r *cli_service.TExecuteStatementResp, err error) { + executeStatementCount++ + executeStatementResp := &cli_service.TExecuteStatementResp{ + Status: &cli_service.TStatus{ + StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, + }, + OperationHandle: &cli_service.TOperationHandle{ + OperationId: &cli_service.THandleIdentifier{ + GUID: []byte{1, 2, 3, 4, 2, 23, 4, 2, 3, 2, 4, 4, 223, 34, 23, 54}, + Secret: []byte("b"), + }, + }, + } + return executeStatementResp, nil + } + var numModRows int64 = 3 + + getOperationStatus := func(ctx context.Context, req *cli_service.TGetOperationStatusReq) (r *cli_service.TGetOperationStatusResp, err error) { + getOperationStatusCount++ + getOperationStatusResp := &cli_service.TGetOperationStatusResp{ + OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_CANCELED_STATE), + NumModifiedRows: &numModRows, + } + return getOperationStatusResp, nil + } + + testClient := &client.TestClient{ + FnExecuteStatement: executeStatement, + FnGetOperationStatus: getOperationStatus, + } + be := newTestBackend(testClient, getTestSession(), config.WithDefaults()) + exStmtResp, opStatusResp, err := be.runQuery(context.Background(), backend.ExecRequest{Query: "select 1", Params: namedToParams([]driver.NamedValue{})}) + + assert.Error(t, err) + assert.Equal(t, 1, executeStatementCount) + assert.Equal(t, 1, getOperationStatusCount) + assert.NotNil(t, exStmtResp) + assert.NotNil(t, opStatusResp) + assert.Equal(t, &numModRows, opStatusResp.NumModifiedRows) + }) + + t.Run("runQuery should return resp when query is finished with DirectResults", func(t *testing.T) { + var executeStatementCount, getOperationStatusCount int + executeStatement := func(ctx context.Context, req *cli_service.TExecuteStatementReq) (r *cli_service.TExecuteStatementResp, err error) { + executeStatementCount++ + executeStatementResp := &cli_service.TExecuteStatementResp{ + Status: &cli_service.TStatus{ + StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, + }, + OperationHandle: &cli_service.TOperationHandle{ + OperationId: &cli_service.THandleIdentifier{ + GUID: []byte{1, 2, 3, 4, 2, 23, 4, 2, 3, 2, 4, 4, 223, 34, 54, 87}, + Secret: []byte("b"), + }, + }, + DirectResults: &cli_service.TSparkDirectResults{ + OperationStatus: &cli_service.TGetOperationStatusResp{ + Status: &cli_service.TStatus{ + StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, + }, + OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_FINISHED_STATE), + }, + }, + } + return executeStatementResp, nil + } + + getOperationStatus := func(ctx context.Context, req *cli_service.TGetOperationStatusReq) (r *cli_service.TGetOperationStatusResp, err error) { + getOperationStatusCount++ + getOperationStatusResp := &cli_service.TGetOperationStatusResp{ + OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_FINISHED_STATE), + } + return getOperationStatusResp, nil + } + + testClient := &client.TestClient{ + FnExecuteStatement: executeStatement, + FnGetOperationStatus: getOperationStatus, + } + be := newTestBackend(testClient, getTestSession(), config.WithDefaults()) + exStmtResp, opStatusResp, err := be.runQuery(context.Background(), backend.ExecRequest{Query: "select 1", Params: namedToParams([]driver.NamedValue{})}) + + assert.NoError(t, err) + assert.Equal(t, 1, executeStatementCount) + assert.Equal(t, 0, getOperationStatusCount) // GetOperationStatus should not be called, already provided in DirectResults + assert.NotNil(t, exStmtResp) + assert.NotNil(t, opStatusResp) + }) + + t.Run("runQuery should return resp and err when query is cancelled with DirectResults", func(t *testing.T) { + var executeStatementCount, getOperationStatusCount int + executeStatement := func(ctx context.Context, req *cli_service.TExecuteStatementReq) (r *cli_service.TExecuteStatementResp, err error) { + executeStatementCount++ + executeStatementResp := &cli_service.TExecuteStatementResp{ + Status: &cli_service.TStatus{ + StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, + }, + OperationHandle: &cli_service.TOperationHandle{ + OperationId: &cli_service.THandleIdentifier{ + GUID: []byte{1, 2, 3, 4, 2, 23, 4, 2, 3, 2, 3, 4, 4, 223, 34, 54}, + Secret: []byte("b"), + }, + }, + DirectResults: &cli_service.TSparkDirectResults{ + OperationStatus: &cli_service.TGetOperationStatusResp{ + Status: &cli_service.TStatus{ + StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, + }, + OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_CANCELED_STATE), + }, + }, + } + return executeStatementResp, nil + } + + getOperationStatus := func(ctx context.Context, req *cli_service.TGetOperationStatusReq) (r *cli_service.TGetOperationStatusResp, err error) { + getOperationStatusCount++ + getOperationStatusResp := &cli_service.TGetOperationStatusResp{ + OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_FINISHED_STATE), + } + return getOperationStatusResp, nil + } + + testClient := &client.TestClient{ + FnExecuteStatement: executeStatement, + FnGetOperationStatus: getOperationStatus, + } + be := newTestBackend(testClient, getTestSession(), config.WithDefaults()) + exStmtResp, opStatusResp, err := be.runQuery(context.Background(), backend.ExecRequest{Query: "select 1", Params: namedToParams([]driver.NamedValue{})}) + + assert.Error(t, err) + assert.Equal(t, 1, executeStatementCount) + assert.Equal(t, 0, getOperationStatusCount) // GetOperationStatus should not be called, already provided in DirectResults + assert.NotNil(t, exStmtResp) + assert.NotNil(t, opStatusResp) + }) + + t.Run("runQuery should return resp when query is finished but DirectResults still live", func(t *testing.T) { + var executeStatementCount, getOperationStatusCount int + executeStatement := func(ctx context.Context, req *cli_service.TExecuteStatementReq) (r *cli_service.TExecuteStatementResp, err error) { + executeStatementCount++ + executeStatementResp := &cli_service.TExecuteStatementResp{ + Status: &cli_service.TStatus{ + StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, + }, + OperationHandle: &cli_service.TOperationHandle{ + OperationId: &cli_service.THandleIdentifier{ + GUID: []byte{1, 2, 3, 4, 2, 23, 4, 2, 3, 2, 3, 4, 4, 223, 34, 54}, + Secret: []byte("b"), + }, + }, + DirectResults: &cli_service.TSparkDirectResults{ + OperationStatus: &cli_service.TGetOperationStatusResp{ + Status: &cli_service.TStatus{ + StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, + }, + OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_INITIALIZED_STATE), + }, + }, + } + return executeStatementResp, nil + } + var numModRows int64 = 3 + getOperationStatus := func(ctx context.Context, req *cli_service.TGetOperationStatusReq) (r *cli_service.TGetOperationStatusResp, err error) { + getOperationStatusCount++ + getOperationStatusResp := &cli_service.TGetOperationStatusResp{ + OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_FINISHED_STATE), + NumModifiedRows: &numModRows, + } + return getOperationStatusResp, nil + } + + testClient := &client.TestClient{ + FnExecuteStatement: executeStatement, + FnGetOperationStatus: getOperationStatus, + } + be := newTestBackend(testClient, getTestSession(), config.WithDefaults()) + exStmtResp, opStatusResp, err := be.runQuery(context.Background(), backend.ExecRequest{Query: "select 1", Params: namedToParams([]driver.NamedValue{})}) + + assert.NoError(t, err) + assert.Equal(t, 1, executeStatementCount) + assert.Equal(t, &numModRows, opStatusResp.NumModifiedRows) + assert.Equal(t, 1, getOperationStatusCount) + assert.NotNil(t, exStmtResp) + assert.NotNil(t, opStatusResp) + }) + + t.Run("runQuery should return resp and err when query is cancelled after DirectResults still live", func(t *testing.T) { + var executeStatementCount, getOperationStatusCount int + executeStatement := func(ctx context.Context, req *cli_service.TExecuteStatementReq) (r *cli_service.TExecuteStatementResp, err error) { + executeStatementCount++ + executeStatementResp := &cli_service.TExecuteStatementResp{ + Status: &cli_service.TStatus{ + StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, + }, + OperationHandle: &cli_service.TOperationHandle{ + OperationId: &cli_service.THandleIdentifier{ + GUID: []byte{1, 2, 3, 4, 2, 23, 4, 2, 3, 2, 3, 4, 4, 223, 34, 54}, + Secret: []byte("b"), + }, + }, + DirectResults: &cli_service.TSparkDirectResults{ + OperationStatus: &cli_service.TGetOperationStatusResp{ + Status: &cli_service.TStatus{ + StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, + }, + OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_INITIALIZED_STATE), + }, + }, + } + return executeStatementResp, nil + } + + getOperationStatus := func(ctx context.Context, req *cli_service.TGetOperationStatusReq) (r *cli_service.TGetOperationStatusResp, err error) { + getOperationStatusCount++ + getOperationStatusResp := &cli_service.TGetOperationStatusResp{ + OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_CANCELED_STATE), + } + return getOperationStatusResp, nil + } + + testClient := &client.TestClient{ + FnExecuteStatement: executeStatement, + FnGetOperationStatus: getOperationStatus, + } + be := newTestBackend(testClient, getTestSession(), config.WithDefaults()) + exStmtResp, opStatusResp, err := be.runQuery(context.Background(), backend.ExecRequest{Query: "select 1", Params: namedToParams([]driver.NamedValue{})}) + + assert.Error(t, err) + assert.Equal(t, 1, executeStatementCount) + assert.Equal(t, 1, getOperationStatusCount) + assert.NotNil(t, exStmtResp) + assert.NotNil(t, opStatusResp) + }) +} diff --git a/internal/backend/thrift/debuglog_integration_test.go b/internal/backend/thrift/debuglog_integration_test.go new file mode 100644 index 00000000..91676b48 --- /dev/null +++ b/internal/backend/thrift/debuglog_integration_test.go @@ -0,0 +1,130 @@ +package thrift + +// Demonstrates that step debug logging (internal/debuglog) fires through the +// Thrift backend's execute path, producing ordered, function-tagged events, so a +// regression that silently drops the instrumentation is caught. + +import ( + "context" + "encoding/json" + "strings" + "sync" + "testing" + + "github.com/databricks/databricks-sql-go/internal/backend" + "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/debuglog" + "github.com/databricks/databricks-sql-go/logger" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +type logEntry struct { + Seq uint64 `json:"seq"` + Fn string `json:"fn"` + Phase string `json:"phase"` +} + +// captureBuf is a concurrency-safe sink for the debug logger. +type captureBuf struct { + mu sync.Mutex + buf strings.Builder +} + +func (c *captureBuf) Write(p []byte) (int, error) { + c.mu.Lock() + defer c.mu.Unlock() + return c.buf.Write(p) +} + +func (c *captureBuf) String() string { + c.mu.Lock() + defer c.mu.Unlock() + return c.buf.String() +} + +func TestBackend_DebugLoggingOrderedAndNested(t *testing.T) { + prev := debuglog.SetEnabled(true) + var buf captureBuf + logger.SetLogOutput(&buf) + t.Cleanup(func() { + debuglog.SetEnabled(prev) + logger.SetLogOutput(nil) + }) + + // A successful direct-results execute so we traverse Execute -> runQuery -> + // executeStatement without needing a poll loop. + executeStatement := func(ctx context.Context, req *cli_service.TExecuteStatementReq) (*cli_service.TExecuteStatementResp, error) { + return &cli_service.TExecuteStatementResp{ + Status: &cli_service.TStatus{StatusCode: cli_service.TStatusCode_SUCCESS_STATUS}, + OperationHandle: &cli_service.TOperationHandle{OperationId: &cli_service.THandleIdentifier{GUID: []byte("0123456789abcdef"), Secret: []byte("s")}}, + DirectResults: &cli_service.TSparkDirectResults{ + OperationStatus: &cli_service.TGetOperationStatusResp{ + Status: &cli_service.TStatus{StatusCode: cli_service.TStatusCode_SUCCESS_STATUS}, + OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_FINISHED_STATE), + }, + }, + }, nil + } + be := newTestBackend(&client.TestClient{FnExecuteStatement: executeStatement}, getTestSession(), config.WithDefaults()) + + _, err := be.Execute(context.Background(), backend.ExecRequest{Query: "select 1"}) + assert.NoError(t, err) + + entries := parseEntries(t, buf.String()) + + // All three nested steps must be present. + for _, fn := range []string{"thrift.Backend.Execute", "thrift.Backend.runQuery", "thrift.Backend.executeStatement"} { + assert.True(t, hasStep(entries, fn, "enter"), "expected an enter for %s", fn) + assert.True(t, hasStep(entries, fn, "done"), "expected a done for %s", fn) + } + + // Sequence numbers strictly increase across the whole stream — the total + // order that lets Go and kernel-Rust lines interleave into one ordered log. + require.GreaterOrEqual(t, len(entries), 6, "expected at least enter+done for 3 nested steps") + for i := 1; i < len(entries); i++ { + assert.Greater(t, entries[i].Seq, entries[i-1].Seq, "sequence must strictly increase") + } + + // The outermost Execute enters before the inner executeStatement enters and + // finishes after it — the nesting is reflected in the stream order. + execEnter := indexOfStep(entries, "thrift.Backend.Execute", "enter") + stmtEnter := indexOfStep(entries, "thrift.Backend.executeStatement", "enter") + execDone := indexOfStep(entries, "thrift.Backend.Execute", "done") + require.True(t, execEnter >= 0 && stmtEnter >= 0 && execDone >= 0, "all three markers present") + assert.Less(t, execEnter, stmtEnter, "Execute should enter before executeStatement") + assert.Less(t, stmtEnter, execDone, "executeStatement should run inside Execute") +} + +func parseEntries(t *testing.T, out string) []logEntry { + t.Helper() + var entries []logEntry + for _, line := range strings.Split(strings.TrimRight(out, "\n"), "\n") { + if line == "" { + continue + } + var e logEntry + if err := json.Unmarshal([]byte(line), &e); err != nil { + t.Fatalf("line is not JSON: %q (%v)", line, err) + } + if e.Fn != "" { // ignore any non-debug lines (e.g. the logger's init info line) + entries = append(entries, e) + } + } + return entries +} + +func hasStep(entries []logEntry, fn, phase string) bool { + return indexOfStep(entries, fn, phase) >= 0 +} + +func indexOfStep(entries []logEntry, fn, phase string) int { + for i, e := range entries { + if e.Fn == fn && e.Phase == phase { + return i + } + } + return -1 +} diff --git a/internal/backend/thrift/export_test_helper.go b/internal/backend/thrift/export_test_helper.go new file mode 100644 index 00000000..edb8781f --- /dev/null +++ b/internal/backend/thrift/export_test_helper.go @@ -0,0 +1,38 @@ +package thrift + +import ( + "github.com/databricks/databricks-sql-go/internal/backend" + "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" +) + +// Test helpers exported for the dbsql package's tests. They are defined in a +// non-_test file because Go's test-only visibility would hide them from a +// different package; they take internal-only types and are not part of the +// package's real API surface. + +// ParamsToSparkForTest exposes the neutral-param -> TSparkParameter mapping so +// the dbsql parameter tests can assert on the Thrift wire form end-to-end. +func ParamsToSparkForTest(params []backend.Param) []*cli_service.TSparkParameter { + return toSparkParameters(params) +} + +// NewForTest builds a Thrift Backend from a mock client, a canned session, and a +// config, bypassing New and OpenSession, so the dbsql connection tests can back a +// conn with an injected mock client. It populates the cached sessionID from the +// handle exactly as OpenSession would. +func NewForTest(cli cli_service.TCLIService, session *cli_service.TOpenSessionResp, cfg *config.Config) *Backend { + b := &Backend{cfg: cfg, client: cli, session: session} + if session != nil && session.SessionHandle != nil { + b.sessionID = client.SprintGuid(session.SessionHandle.GetSessionId().GUID) + } + return b +} + +// OperationForTest builds an Operation wrapping a canned execute response and +// status, using this backend for any follow-up RPCs (e.g. the +// GetResultSetMetadata that IsStaging may issue). +func (b *Backend) OperationForTest(exStmtResp *cli_service.TExecuteStatementResp, opStatusResp *cli_service.TGetOperationStatusResp) *thriftOperation { + return &thriftOperation{backend: b, exStmtResp: exStmtResp, opStatusResp: opStatusResp} +} diff --git a/internal/backend/thrift/operation.go b/internal/backend/thrift/operation.go new file mode 100644 index 00000000..e2adff83 --- /dev/null +++ b/internal/backend/thrift/operation.go @@ -0,0 +1,137 @@ +package thrift + +import ( + "context" + "database/sql/driver" + + dbsqlerr "github.com/databricks/databricks-sql-go/errors" + "github.com/databricks/databricks-sql-go/internal/backend" + "github.com/databricks/databricks-sql-go/internal/cli_service" + "github.com/databricks/databricks-sql-go/internal/client" + "github.com/databricks/databricks-sql-go/internal/debuglog" + dbsqlerrint "github.com/databricks/databricks-sql-go/internal/errors" + dbsqlrows "github.com/databricks/databricks-sql-go/internal/rows" +) + +// thriftOperation is the Thrift implementation of backend.Operation. It carries +// the execute response (holding the operation handle and optional direct +// results) and the terminal status response. Its constructor does no I/O; the +// RPCs (result fetch, staging-metadata lookup, close) happen in the methods. +type thriftOperation struct { + backend *Backend + exStmtResp *cli_service.TExecuteStatementResp + opStatusResp *cli_service.TGetOperationStatusResp + // statementID caches the formatted operation GUID (SprintGuid allocates and + // StatementID is read several times per query). An Operation is used by one + // goroutine at a time (pool discipline), so a plain memo needs no lock. + statementID string + statementIDSet bool +} + +var _ backend.Operation = (*thriftOperation)(nil) + +// hasHandle reports whether the server returned an operation handle. Every +// accessor tolerates its absence so the non-nil-Operation contract holds even on +// the pre-handle error path. +func (o *thriftOperation) hasHandle() bool { + return o.exStmtResp != nil && o.exStmtResp.OperationHandle != nil +} + +// StatementID is the formatted operation id, or "" when no handle exists. +// Computed once and cached. +func (o *thriftOperation) StatementID() string { + if o.statementIDSet { + return o.statementID + } + if o.hasHandle() && o.exStmtResp.OperationHandle.OperationId != nil { + o.statementID = client.SprintGuid(o.exStmtResp.OperationHandle.OperationId.GUID) + } + o.statementIDSet = true + return o.statementID +} + +// AffectedRows is the modified-row count for ExecContext, from the terminal +// status response. +func (o *thriftOperation) AffectedRows() int64 { + return o.opStatusResp.GetNumModifiedRows() +} + +// Results builds the driver.Rows for the operation's result set by delegating to +// the rows layer with the Thrift handle, client, and direct results. The caller +// supplies the telemetry callbacks (full for a query, chunk-timing-only for a +// staging read). +func (o *thriftOperation) Results(ctx context.Context, callbacks *dbsqlrows.TelemetryCallbacks) (driver.Rows, error) { + defer debuglog.Track(ctx, "thrift.Operation.Results", "stmt=%s", o.StatementID())() + + var directResults *cli_service.TSparkDirectResults + var opHandle *cli_service.TOperationHandle + if o.exStmtResp != nil { + directResults = o.exStmtResp.DirectResults + opHandle = o.exStmtResp.OperationHandle + } + return dbsqlrows.NewRows(ctx, opHandle, o.backend.client, o.backend.cfg, directResults, callbacks) +} + +// IsStaging reports whether this operation is a staging operation. It reuses the +// result-set metadata from direct results when present, otherwise issues the +// GetResultSetMetadata RPC — the exact branch conn.execStagingOperation used. +// Returns false (no error) when there is no handle. The raw RPC error is returned +// unwrapped; the caller owns the "error performing staging operation" wrapping. +func (o *thriftOperation) IsStaging(ctx context.Context) (bool, error) { + defer debuglog.Track(ctx, "thrift.Operation.IsStaging", "stmt=%s", o.StatementID())() + + if !o.hasHandle() { + return false, nil + } + if o.exStmtResp.DirectResults != nil && o.exStmtResp.DirectResults.ResultSetMetadata != nil { + md := o.exStmtResp.DirectResults.ResultSetMetadata + return md.IsStagingOperation != nil && *md.IsStagingOperation, nil + } + + debuglog.Logf(ctx, "thrift.Operation.IsStaging", "GetResultSetMetadata") + resp, err := o.backend.client.GetResultSetMetadata(ctx, &cli_service.TGetResultSetMetadataReq{ + OperationHandle: o.exStmtResp.OperationHandle, + }) + if err != nil { + return false, err + } + return resp.IsStagingOperation != nil && *resp.IsStagingOperation, nil +} + +// Close best-effort closes the server operation. It skips the CloseOperation RPC +// when there is no handle, when direct results already closed it, or when the +// operation is already in the CLOSED state; closed reports whether the RPC was +// actually sent, so the caller records CLOSE_STATEMENT telemetry only in that case. +func (o *thriftOperation) Close(ctx context.Context) (bool, error) { + defer debuglog.Track(ctx, "thrift.Operation.Close", "stmt=%s", o.StatementID())() + + if !o.hasHandle() { + return false, nil + } + alreadyClosed := o.exStmtResp.DirectResults != nil && o.exStmtResp.DirectResults.CloseOperation != nil + if alreadyClosed { + return false, nil + } + if o.opStatusResp != nil && o.opStatusResp.GetOperationState() == cli_service.TOperationState_CLOSED_STATE { + return false, nil + } + + debuglog.Logf(ctx, "thrift.Operation.Close", "CloseOperation") + _, err := o.backend.client.CloseOperation(ctx, &cli_service.TCloseOperationReq{ + OperationHandle: o.exStmtResp.OperationHandle, + }) + if err != nil { + return true, err + } + return true, nil +} + +// ExecutionError wraps cause as the driver's execution error, attaching this +// operation's terminal sqlstate from the status response. Returns nil when cause +// is nil. +func (o *thriftOperation) ExecutionError(ctx context.Context, cause error) error { + if cause == nil { + return nil + } + return dbsqlerrint.NewExecutionError(ctx, dbsqlerr.ErrQueryExecution, cause, o.opStatusResp) +} diff --git a/internal/backend/thrift/operation_test.go b/internal/backend/thrift/operation_test.go new file mode 100644 index 00000000..172db425 --- /dev/null +++ b/internal/backend/thrift/operation_test.go @@ -0,0 +1,175 @@ +package thrift + +import ( + "context" + "errors" + "testing" + + "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/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// opHandle builds a minimal operation handle for tests. +func opHandle() *cli_service.TOperationHandle { + return &cli_service.TOperationHandle{ + OperationId: &cli_service.THandleIdentifier{ + GUID: []byte{1, 2, 3, 4, 2, 23, 4, 2, 3, 1, 2, 3, 4, 4, 223, 34}, + Secret: []byte("s"), + }, + } +} + +// TestOperationClose pins the close-decision matrix: Close must issue the +// CloseOperation RPC (closed=true) only for a live terminal op with an open +// handle, and skip it (closed=false) for no-handle / already-closed-via-direct- +// results / already-CLOSED_STATE. The returned `closed` bool gates CLOSE_STATEMENT +// telemetry in conn, so the skip branches are load-bearing. +func TestOperationClose(t *testing.T) { + makeOp := func(exResp *cli_service.TExecuteStatementResp, opStatus *cli_service.TGetOperationStatusResp, onClose func()) *thriftOperation { + tc := &client.TestClient{ + FnCloseOperation: func(ctx context.Context, req *cli_service.TCloseOperationReq) (*cli_service.TCloseOperationResp, error) { + if onClose != nil { + onClose() + } + return &cli_service.TCloseOperationResp{}, nil + }, + } + be := newTestBackend(tc, getTestSession(), config.WithDefaults()) + return be.OperationForTest(exResp, opStatus) + } + + t.Run("no handle -> skip RPC, closed=false", func(t *testing.T) { + called := false + op := makeOp(&cli_service.TExecuteStatementResp{}, nil, func() { called = true }) + closed, err := op.Close(context.Background()) + assert.False(t, closed) + assert.NoError(t, err) + assert.False(t, called, "CloseOperation must not be called without a handle") + }) + + t.Run("already closed via direct results -> skip RPC, closed=false", func(t *testing.T) { + called := false + exResp := &cli_service.TExecuteStatementResp{ + OperationHandle: opHandle(), + DirectResults: &cli_service.TSparkDirectResults{ + CloseOperation: &cli_service.TCloseOperationResp{}, + }, + } + op := makeOp(exResp, nil, func() { called = true }) + closed, err := op.Close(context.Background()) + assert.False(t, closed) + assert.NoError(t, err) + assert.False(t, called, "must not re-close an operation the server already closed") + }) + + t.Run("already in CLOSED_STATE -> skip RPC, closed=false", func(t *testing.T) { + called := false + exResp := &cli_service.TExecuteStatementResp{OperationHandle: opHandle()} + opStatus := &cli_service.TGetOperationStatusResp{ + OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_CLOSED_STATE), + } + op := makeOp(exResp, opStatus, func() { called = true }) + closed, err := op.Close(context.Background()) + assert.False(t, closed) + assert.NoError(t, err) + assert.False(t, called, "must not re-close an operation already in CLOSED_STATE") + }) + + // Live terminal states with an open handle: Close issues the RPC. + for _, state := range []cli_service.TOperationState{ + cli_service.TOperationState_FINISHED_STATE, + cli_service.TOperationState_ERROR_STATE, + cli_service.TOperationState_CANCELED_STATE, + cli_service.TOperationState_TIMEDOUT_STATE, + } { + t.Run("open handle in "+state.String()+" -> issue RPC, closed=true", func(t *testing.T) { + called := false + exResp := &cli_service.TExecuteStatementResp{OperationHandle: opHandle()} + opStatus := &cli_service.TGetOperationStatusResp{OperationState: cli_service.TOperationStatePtr(state)} + op := makeOp(exResp, opStatus, func() { called = true }) + closed, err := op.Close(context.Background()) + assert.True(t, closed, "a live op with an open handle must be closed") + assert.NoError(t, err) + assert.True(t, called, "CloseOperation RPC must be issued") + }) + } + + t.Run("nil opStatus with open handle -> issue RPC (state defaults, not CLOSED)", func(t *testing.T) { + called := false + exResp := &cli_service.TExecuteStatementResp{OperationHandle: opHandle()} + op := makeOp(exResp, nil, func() { called = true }) + closed, err := op.Close(context.Background()) + assert.True(t, closed) + assert.NoError(t, err) + assert.True(t, called) + }) + + t.Run("RPC error is returned with closed=true", func(t *testing.T) { + tc := &client.TestClient{ + FnCloseOperation: func(ctx context.Context, req *cli_service.TCloseOperationReq) (*cli_service.TCloseOperationResp, error) { + return nil, errors.New("boom") + }, + } + be := newTestBackend(tc, getTestSession(), config.WithDefaults()) + op := be.OperationForTest(&cli_service.TExecuteStatementResp{OperationHandle: opHandle()}, nil) + closed, err := op.Close(context.Background()) + assert.True(t, closed, "closed=true reflects that the RPC was attempted, so telemetry records it") + assert.Error(t, err) + }) +} + +// TestOperationIsStaging covers the two branches: answer from direct-results +// metadata (no RPC) vs. GetResultSetMetadata RPC when direct results are absent. +func TestOperationIsStaging(t *testing.T) { + truePtr := true + falsePtr := false + + t.Run("from direct results, no RPC", func(t *testing.T) { + metaCalls := 0 + tc := &client.TestClient{ + FnGetResultSetMetadata: func(ctx context.Context, req *cli_service.TGetResultSetMetadataReq) (*cli_service.TGetResultSetMetadataResp, error) { + metaCalls++ + return &cli_service.TGetResultSetMetadataResp{}, nil + }, + } + be := newTestBackend(tc, getTestSession(), config.WithDefaults()) + exResp := &cli_service.TExecuteStatementResp{ + OperationHandle: opHandle(), + DirectResults: &cli_service.TSparkDirectResults{ + ResultSetMetadata: &cli_service.TGetResultSetMetadataResp{IsStagingOperation: &truePtr}, + }, + } + op := be.OperationForTest(exResp, nil) + staging, err := op.IsStaging(context.Background()) + assert.True(t, staging) + assert.NoError(t, err) + assert.Equal(t, 0, metaCalls, "direct-results metadata must be used without an RPC") + }) + + t.Run("falls back to GetResultSetMetadata RPC", func(t *testing.T) { + metaCalls := 0 + tc := &client.TestClient{ + FnGetResultSetMetadata: func(ctx context.Context, req *cli_service.TGetResultSetMetadataReq) (*cli_service.TGetResultSetMetadataResp, error) { + metaCalls++ + return &cli_service.TGetResultSetMetadataResp{IsStagingOperation: &falsePtr}, nil + }, + } + be := newTestBackend(tc, getTestSession(), config.WithDefaults()) + op := be.OperationForTest(&cli_service.TExecuteStatementResp{OperationHandle: opHandle()}, nil) + staging, err := op.IsStaging(context.Background()) + assert.False(t, staging) + assert.NoError(t, err) + assert.Equal(t, 1, metaCalls, "absent direct-results metadata must trigger one RPC") + }) + + t.Run("no handle -> false, no RPC", func(t *testing.T) { + be := newTestBackend(&client.TestClient{}, getTestSession(), config.WithDefaults()) + op := be.OperationForTest(&cli_service.TExecuteStatementResp{}, nil) + staging, err := op.IsStaging(context.Background()) + assert.False(t, staging) + require.NoError(t, err) + }) +} diff --git a/internal/backend/thrift/params.go b/internal/backend/thrift/params.go new file mode 100644 index 00000000..3d6f51ca --- /dev/null +++ b/internal/backend/thrift/params.go @@ -0,0 +1,41 @@ +package thrift + +import ( + "github.com/databricks/databricks-sql-go/internal/backend" + "github.com/databricks/databricks-sql-go/internal/cli_service" +) + +// toSparkParameters maps the backend-neutral, already-stringified parameters +// onto Thrift's TSparkParameter wire type. A nil Param.Value denotes SQL NULL +// and produces a TSparkParameter whose Value is nil. +func toSparkParameters(params []backend.Param) []*cli_service.TSparkParameter { + if len(params) == 0 { + return nil + } + sparkParams := make([]*cli_service.TSparkParameter, 0, len(params)) + for i := range params { + p := params[i] + + var sparkValue *cli_service.TSparkParameterValue + if p.Value != nil { + // Copy into a fresh local so &v doesn't alias the loop variable's + // slot across iterations. + v := *p.Value + sparkValue = &cli_service.TSparkParameterValue{StringValue: &v} + } + + var sparkName *string + if p.Name != "" { + name := p.Name + sparkName = &name + } + + sparkType := p.Type + sparkParams = append(sparkParams, &cli_service.TSparkParameter{ + Name: sparkName, + Type: &sparkType, + Value: sparkValue, + }) + } + return sparkParams +} diff --git a/internal/backend/thrift/params_test.go b/internal/backend/thrift/params_test.go new file mode 100644 index 00000000..d0107653 --- /dev/null +++ b/internal/backend/thrift/params_test.go @@ -0,0 +1,69 @@ +package thrift + +import ( + "testing" + + "github.com/databricks/databricks-sql-go/internal/backend" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// Tests the stage-2 wire mapping in isolation: backend.Param -> TSparkParameter. +// Type inference and value stringification (stage 1) live in the dbsql package +// and are tested there; this covers only the mechanical shape mapping — named vs +// positional, the nil-Value -> SQL NULL case, and pointer independence across +// entries. strPtr is the shared helper from backend_test.go. +func TestToSparkParameters(t *testing.T) { + t.Run("empty input yields nil", func(t *testing.T) { + assert.Nil(t, toSparkParameters(nil)) + assert.Nil(t, toSparkParameters([]backend.Param{})) + }) + + t.Run("named parameter maps name, type and value", func(t *testing.T) { + out := toSparkParameters([]backend.Param{{Name: "p1", Type: "BIGINT", Value: strPtr("5")}}) + require.Len(t, out, 1) + require.NotNil(t, out[0].Name) + assert.Equal(t, "p1", *out[0].Name) + require.NotNil(t, out[0].Type) + assert.Equal(t, "BIGINT", *out[0].Type) + require.NotNil(t, out[0].Value) + require.NotNil(t, out[0].Value.StringValue) + assert.Equal(t, "5", *out[0].Value.StringValue) + }) + + t.Run("positional parameter has nil name", func(t *testing.T) { + out := toSparkParameters([]backend.Param{{Name: "", Type: "STRING", Value: strPtr("x")}}) + require.Len(t, out, 1) + assert.Nil(t, out[0].Name, "empty Name must map to a nil TSparkParameter.Name (positional)") + assert.Equal(t, "STRING", *out[0].Type) + assert.Equal(t, "x", *out[0].Value.StringValue) + }) + + t.Run("nil Value maps to SQL NULL (nil TSparkParameterValue)", func(t *testing.T) { + out := toSparkParameters([]backend.Param{{Name: "n", Type: "VOID", Value: nil}}) + require.Len(t, out, 1) + assert.Nil(t, out[0].Value, "nil Param.Value must produce a nil TSparkParameter.Value") + require.NotNil(t, out[0].Type) + assert.Equal(t, "VOID", *out[0].Type) + }) + + t.Run("multiple params keep independent value pointers", func(t *testing.T) { + // Guards the loop-variable aliasing fix: each entry must point at its own + // copy of the value/type, not a shared final iteration's slot. + out := toSparkParameters([]backend.Param{ + {Name: "a", Type: "INT", Value: strPtr("1")}, + {Name: "b", Type: "INT", Value: strPtr("2")}, + {Name: "c", Type: "INT", Value: strPtr("3")}, + }) + require.Len(t, out, 3) + assert.Equal(t, "1", *out[0].Value.StringValue) + assert.Equal(t, "2", *out[1].Value.StringValue) + assert.Equal(t, "3", *out[2].Value.StringValue) + assert.Equal(t, "a", *out[0].Name) + assert.Equal(t, "b", *out[1].Name) + assert.Equal(t, "c", *out[2].Name) + // Distinct backing pointers, not aliases of one slot. + assert.NotSame(t, out[0].Value.StringValue, out[1].Value.StringValue) + assert.NotSame(t, out[0].Name, out[1].Name) + }) +} diff --git a/internal/debuglog/debuglog.go b/internal/debuglog/debuglog.go new file mode 100644 index 00000000..3afbc05e --- /dev/null +++ b/internal/debuglog/debuglog.go @@ -0,0 +1,162 @@ +// Package debuglog is the step-level debug tracer for the databricks-sql-go +// SEA/kernel integration. +// +// The SEA-via-kernel path crosses a Go-GC <-> Rust-ownership FFI boundary where +// failures are silent and non-local (thread migration, use-after-free, thread +// blowup, cross-region latency that looks like "the protocol is slow"). A Go +// stack trace stops at the cgo edge, so the way to see which step failed or +// slowed down is an explicit, ordered, timed, per-function log at each step. +// +// Emission is through the driver's existing logger (logger.DebugLogger, a +// trace-level sibling that shares the main logger's sink and format), so there +// is a single logging system and one place to redirect output. It has its own +// DBSQL_DEBUG_LOG flag (below) so it is enabled independently of the package log +// level. +// +// Every event carries a process-global monotonic sequence number (a total order +// even when timestamps tie), the conn/correlation/query ids from the context, +// the function the step is in, and — for Track — the elapsed time. +// +// Gating: OFF by default (a single atomic-bool load per call site when off). +// Enable via the DBSQL_DEBUG_LOG env var (1/true/yes/on) or SetEnabled. +package debuglog + +import ( + "context" + "os" + "strings" + "sync/atomic" + "time" + + "github.com/rs/zerolog" + + "github.com/databricks/databricks-sql-go/driverctx" + "github.com/databricks/databricks-sql-go/logger" +) + +// enabled is the single gate. Kept as an atomic so SetEnabled is safe to call +// concurrently with logging (e.g. a test flipping it while a goroutine logs). +var enabled atomic.Bool + +// seq is the process-global sequence counter — the primary ordering key, since +// timestamps can tie or go backwards across cores. +var seq atomic.Uint64 + +// clockOverride optionally replaces the time source so tests can make timing +// deterministic. nil (the default) means use time.Now. Kept atomic — like +// enabled and seq — so an override is safe to set concurrently with logging. +var clockOverride atomic.Pointer[func() time.Time] + +// now returns the current time from the override if one is set, else time.Now. +func now() time.Time { + if f := clockOverride.Load(); f != nil { + return (*f)() + } + return time.Now() +} + +func init() { + if envEnabled(os.Getenv("DBSQL_DEBUG_LOG")) { + enabled.Store(true) + } +} + +// envEnabled reports whether an env-var value means "on". Accepts 1/true/yes/on +// (case-insensitive, trimmed); everything else (including empty) is off. +func envEnabled(v string) bool { + switch strings.ToLower(strings.TrimSpace(v)) { + case "1", "true", "yes", "on": + return true + default: + return false + } +} + +// Enabled reports whether step logging is currently on. Cheap enough to call at +// every step; callers use it to skip building expensive log arguments. +func Enabled() bool { return enabled.Load() } + +// SetEnabled turns step logging on or off at runtime. Returns the previous value +// so a test can save and restore it. +func SetEnabled(on bool) (previous bool) { return enabled.Swap(on) } + +// setClock overrides the time source, or resets to time.Now when f is nil +// (test-only; unexported). +func setClock(f func() time.Time) { + if f == nil { + clockOverride.Store(nil) + return + } + clockOverride.Store(&f) +} + +// Logf emits a single ordered, function-tagged step event. fn is the function +// the step is in (e.g. "thrift.Backend.Execute"); the format/args describe the +// step. No-op when logging is disabled — guard the call with Enabled() if the +// args are expensive to build. +func Logf(ctx context.Context, fn string, format string, args ...any) { + if !enabled.Load() { + return + } + event(ctx, fn, "").Msgf(format, args...) +} + +// Track marks the start of a step and returns a function that logs the step's +// elapsed time. Intended for `defer`: +// +// defer debuglog.Track(ctx, "thrift.Backend.Execute", "sql=%q", query)() +// +// The enter event is emitted immediately; the returned closure emits a matching +// done event carrying the elapsed duration. Both are no-ops when disabled, and +// the returned closure is always safe to call. +func Track(ctx context.Context, fn string, format string, args ...any) func() { + if !enabled.Load() { + return func() {} + } + start := now() + event(ctx, fn, "enter").Msgf(format, args...) + return func() { + // If logging was disabled between enter and here, skip the done line. This + // can leave an orphan enter with no done, which is acceptable for a debug + // trace — SetEnabled is not toggled mid-query in normal use, and emitting a + // done through a now-disabled logger would be worse. + if !enabled.Load() { + return + } + event(ctx, fn, "done").Dur("elapsed", now().Sub(start)).Msg("") + } +} + +// event builds a trace-level zerolog event pre-populated with a nanosecond +// timestamp, the sequence number, the step's function, the phase +// ("enter"/"done", or "" for a point log), and the conn/correlation/query ids +// present on the context. The caller finishes it with Msg/Msgf. Emission goes +// through logger.DebugLogger so the step trace shares the driver's single sink. +// +// The ts field is written as an explicit RFC3339Nano string rather than via +// zerolog's Time(), because zerolog's time rendering is governed by the global +// zerolog.TimeFieldFormat (RFC3339, whole seconds). Latency work needs sub-second +// resolution to correlate steps with server-side traces and to see inter-step +// gaps, so ts carries nanosecond precision independent of the shared logger's +// coarser default `time` field. +func event(ctx context.Context, fn, phase string) *zerolog.Event { + e := logger.DebugLogger().Log(). + Str("ts", now().Format(time.RFC3339Nano)). + Uint64("seq", seq.Add(1)). + Str("fn", fn) + if phase != "" { + e = e.Str("phase", phase) + } + if ctx != nil { + if v := driverctx.ConnIdFromContext(ctx); v != "" { + e = e.Str("connId", v) + } + if v := driverctx.CorrelationIdFromContext(ctx); v != "" { + e = e.Str("corrId", v) + } + if v := driverctx.QueryIdFromContext(ctx); v != "" { + e = e.Str("queryId", v) + } + } + return e +} diff --git a/internal/debuglog/debuglog_test.go b/internal/debuglog/debuglog_test.go new file mode 100644 index 00000000..f02f3a77 --- /dev/null +++ b/internal/debuglog/debuglog_test.go @@ -0,0 +1,322 @@ +package debuglog + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "strings" + "sync" + "testing" + "time" + + "github.com/databricks/databricks-sql-go/driverctx" + "github.com/databricks/databricks-sql-go/logger" +) + +// entry is one parsed zerolog line from the debug sink. +type entry struct { + TS string `json:"ts"` + Seq uint64 `json:"seq"` + Fn string `json:"fn"` + Phase string `json:"phase"` + Message string `json:"message"` + Elapsed float64 `json:"elapsed"` + ConnID string `json:"connId"` + CorrID string `json:"corrId"` + QueryID string `json:"queryId"` + raw map[string]any +} + +// withCapture enables logging into a fresh buffer for the duration of fn, then +// restores the prior enabled state, output, and clock. Every test goes through +// this so global state never leaks between tests. +func withCapture(t *testing.T, fn func(entries func() []entry)) { + t.Helper() + prevEnabled := SetEnabled(true) + var buf syncBuf + logger.SetLogOutput(&buf) + t.Cleanup(func() { + SetEnabled(prevEnabled) + logger.SetLogOutput(nil) + setClock(nil) // reset to time.Now + }) + fn(func() []entry { return parse(t, buf.String()) }) +} + +// parse decodes the captured JSON lines into entries. +func parse(t *testing.T, out string) []entry { + t.Helper() + var entries []entry + for _, line := range strings.Split(strings.TrimRight(out, "\n"), "\n") { + if line == "" { + continue + } + var raw map[string]any + if err := json.Unmarshal([]byte(line), &raw); err != nil { + t.Fatalf("line is not JSON: %q (%v)", line, err) + } + var e entry + if err := json.Unmarshal([]byte(line), &e); err != nil { + t.Fatalf("cannot decode entry: %q (%v)", line, err) + } + e.raw = raw + entries = append(entries, e) + } + return entries +} + +func TestDisabledByDefaultIsSilent(t *testing.T) { + prev := SetEnabled(false) // force disabled regardless of ambient env / prior tests + defer SetEnabled(prev) + + var buf syncBuf + logger.SetLogOutput(&buf) + defer logger.SetLogOutput(nil) + + Logf(context.Background(), "pkg.Fn", "should not appear") + done := Track(context.Background(), "pkg.Fn", "step") + done() + + if buf.String() != "" { + t.Fatalf("expected no output when disabled, got: %q", buf.String()) + } + if Enabled() { + t.Fatal("Enabled() should be false") + } +} + +func TestEnvEnabled(t *testing.T) { + cases := map[string]bool{ + "1": true, "true": true, "TRUE": true, " yes ": true, "On": true, + "": false, "0": false, "false": false, "no": false, "nope": false, + } + for in, want := range cases { + if got := envEnabled(in); got != want { + t.Errorf("envEnabled(%q) = %v, want %v", in, got, want) + } + } +} + +func TestLogfCarriesFunctionMessageAndSeq(t *testing.T) { + withCapture(t, func(entries func() []entry) { + Logf(context.Background(), "thrift.Backend.Execute", "sql=%q rows=%d", "select 1", 42) + es := entries() + if len(es) != 1 { + t.Fatalf("expected 1 entry, got %d: %+v", len(es), es) + } + e := es[0] + if e.Fn != "thrift.Backend.Execute" { + t.Errorf("fn = %q, want thrift.Backend.Execute", e.Fn) + } + if e.Message != `sql="select 1" rows=42` { + t.Errorf("message = %q, want formatted step", e.Message) + } + if e.Seq == 0 { + t.Errorf("seq should be a positive counter, got %d", e.Seq) + } + if _, ok := e.raw["phase"]; ok { + t.Errorf("point log should have no phase field: %v", e.raw) + } + }) +} + +func TestTimestampIsNanosecondPrecise(t *testing.T) { + withCapture(t, func(entries func() []entry) { + // A fixed instant carrying sub-second nanoseconds — the coarse RFC3339 + // (whole-second) default would drop the fractional part. + fixed := time.Date(2026, 7, 5, 12, 0, 0, 123456789, time.UTC) + setClock(func() time.Time { return fixed }) + + Logf(context.Background(), "pkg.Fn", "step") + e := entries()[0] + + if e.TS == "" { + t.Fatal("ts field missing") + } + parsed, err := time.Parse(time.RFC3339Nano, e.TS) + if err != nil { + t.Fatalf("ts %q is not RFC3339Nano: %v", e.TS, err) + } + if parsed.Nanosecond() != 123456789 { + t.Errorf("ts lost sub-second precision: got %d ns, want 123456789 (ts=%q)", parsed.Nanosecond(), e.TS) + } + }) +} + +func TestTrackEmitsEnterAndDoneWithElapsed(t *testing.T) { + withCapture(t, func(entries func() []entry) { + // Deterministic clock: enter at t0, done 5ms later. + t0 := time.Date(2026, 7, 3, 12, 0, 0, 0, time.UTC) + times := []time.Time{t0, t0.Add(5 * time.Millisecond)} + i := 0 + setClock(func() time.Time { + idx := i + if idx >= len(times) { + idx = len(times) - 1 + } + i++ + return times[idx] + }) + + done := Track(context.Background(), "kernel.Session.Open", "host=%s", "example") + done() + + es := entries() + if len(es) != 2 { + t.Fatalf("expected enter+done (2 entries), got %d: %+v", len(es), es) + } + if es[0].Phase != "enter" || es[1].Phase != "done" { + t.Fatalf("phases = %q,%q want enter,done", es[0].Phase, es[1].Phase) + } + if es[0].Message != "host=example" { + t.Errorf("enter message = %q, want host=example", es[0].Message) + } + // zerolog renders Dur in milliseconds by default; 5ms -> 5. + if es[1].Elapsed != 5 { + t.Errorf("done elapsed = %v ms, want 5", es[1].Elapsed) + } + }) +} + +func TestSequenceIsMonotonic(t *testing.T) { + withCapture(t, func(entries func() []entry) { + for n := 0; n < 5; n++ { + Logf(context.Background(), "pkg.Fn", "line %d", n) + } + es := entries() + if len(es) != 5 { + t.Fatalf("expected 5 entries, got %d", len(es)) + } + for i := 1; i < len(es); i++ { + if es[i].Seq <= es[i-1].Seq { + t.Errorf("sequence not strictly increasing: %d then %d", es[i-1].Seq, es[i].Seq) + } + } + }) +} + +func TestContextIDsAppear(t *testing.T) { + withCapture(t, func(entries func() []entry) { + ctx := driverctx.NewContextWithConnId(context.Background(), "conn-123") + ctx = driverctx.NewContextWithCorrelationId(ctx, "corr-456") + ctx = driverctx.NewContextWithQueryId(ctx, "query-789") + Logf(ctx, "pkg.Fn", "step") + e := entries()[0] + if e.ConnID != "conn-123" || e.CorrID != "corr-456" || e.QueryID != "query-789" { + t.Errorf("ids = %q/%q/%q, want conn-123/corr-456/query-789", e.ConnID, e.CorrID, e.QueryID) + } + }) +} + +func TestEmptyContextIDsAreOmitted(t *testing.T) { + withCapture(t, func(entries func() []entry) { + Logf(context.Background(), "pkg.Fn", "step") + raw := entries()[0].raw + for _, unwanted := range []string{"connId", "corrId", "queryId"} { + if _, ok := raw[unwanted]; ok { + t.Errorf("empty id field %q should be omitted: %v", unwanted, raw) + } + } + }) +} + +func TestNilContextIsSafe(t *testing.T) { + withCapture(t, func(entries func() []entry) { + //nolint:staticcheck // intentionally passing nil ctx to prove it is safe + Logf(nil, "pkg.Fn", "step") + if es := entries(); len(es) != 1 || es[0].Fn != "pkg.Fn" { + t.Errorf("nil context should still log a pkg.Fn entry, got %+v", es) + } + }) +} + +func TestDisableMidStepSuppressesDone(t *testing.T) { + withCapture(t, func(entries func() []entry) { + done := Track(context.Background(), "pkg.Fn", "step") + SetEnabled(false) // caller flips it off before the deferred done fires + done() + SetEnabled(true) // restore so cleanup is consistent + es := entries() + if len(es) != 1 || es[0].Phase != "enter" { + t.Fatalf("expected only the enter entry, got %+v", es) + } + }) +} + +// TestConcurrentLoggingIsRaceFree exercises the atomic gate and the atomic clock +// override under -race: many goroutines log while others flip SetEnabled and +// setClock concurrently. +func TestConcurrentLoggingIsRaceFree(t *testing.T) { + prev := SetEnabled(true) + defer SetEnabled(prev) + logger.SetLogOutput(&syncBuf{}) + defer logger.SetLogOutput(nil) + defer setClock(nil) + + var wg sync.WaitGroup + for g := 0; g < 8; g++ { + wg.Add(1) + go func(g int) { + defer wg.Done() + for i := 0; i < 100; i++ { + Logf(context.Background(), "pkg.Fn", "g=%d i=%d", g, i) + done := Track(context.Background(), "pkg.Fn", "step") + done() + } + }(g) + } + wg.Add(1) + go func() { + defer wg.Done() + for i := 0; i < 50; i++ { + SetEnabled(i%2 == 0) + } + SetEnabled(true) + }() + // Flip the clock override concurrently with logging — reads of clockOverride + // happen in every logged event, so this proves the override is race-free too. + wg.Add(1) + go func() { + defer wg.Done() + fixed := func() time.Time { return time.Unix(0, 0) } + for i := 0; i < 50; i++ { + if i%2 == 0 { + setClock(fixed) + } else { + setClock(nil) + } + } + }() + wg.Wait() +} + +// syncBuf is a concurrency-safe buffer for capturing log output under -race. +type syncBuf struct { + mu sync.Mutex + buf bytes.Buffer +} + +func (s *syncBuf) Write(p []byte) (int, error) { + s.mu.Lock() + defer s.mu.Unlock() + return s.buf.Write(p) +} + +func (s *syncBuf) String() string { + s.mu.Lock() + defer s.mu.Unlock() + return s.buf.String() +} + +// ExampleTrack shows the intended defer-Track idiom. +func ExampleTrack() { + SetEnabled(true) + defer SetEnabled(false) + fn := func(ctx context.Context) { + defer Track(ctx, "pkg.example", "doing work")() + } + fn(context.Background()) + fmt.Println("ran") + // Output: ran +} diff --git a/internal/querytags/querytags.go b/internal/querytags/querytags.go new file mode 100644 index 00000000..d7e3d2d0 --- /dev/null +++ b/internal/querytags/querytags.go @@ -0,0 +1,40 @@ +// Package querytags holds the query-tag wire serialization shared by the public +// dbsql API and the execution backends. +// +// It lives in internal/ so that internal/backend/thrift (and the future kernel +// backend) can serialize per-statement query tags without importing the public +// dbsql package, which would create an import cycle (dbsql -> backend -> dbsql). +// The public dbsql.SerializeQueryTags forwards here so the exported behavior and +// symbol are unchanged. +package querytags + +import "strings" + +// Serialize converts a map of query tags to the wire format string. The format +// is comma-separated key:value pairs (e.g., "team:engineering,app:etl"). +// +// Escaping rules (consistent with the Python and NodeJS connectors): +// - Keys: only backslashes are escaped +// - Values: backslashes, colons, and commas are escaped with a leading backslash +// - Empty string values result in just the key being emitted (no colon) +// +// Returns empty string if the map is nil or empty. +func Serialize(tags map[string]string) string { + if len(tags) == 0 { + return "" + } + + parts := make([]string, 0, len(tags)) + for k, v := range tags { + escapedKey := strings.ReplaceAll(k, `\`, `\\`) + if v == "" { + parts = append(parts, escapedKey) + } else { + escapedValue := strings.ReplaceAll(v, `\`, `\\`) + escapedValue = strings.ReplaceAll(escapedValue, `:`, `\:`) + escapedValue = strings.ReplaceAll(escapedValue, `,`, `\,`) + parts = append(parts, escapedKey+":"+escapedValue) + } + } + return strings.Join(parts, ",") +} diff --git a/logger/logger.go b/logger/logger.go index 683501a1..ab10a6c6 100644 --- a/logger/logger.go +++ b/logger/logger.go @@ -40,11 +40,31 @@ var Logger = &DBSQLLogger{ zerolog.New(os.Stderr).With().Timestamp().Logger(), } +// out is the destination shared by Logger and the debug-trace sibling, so +// SetLogOutput redirects both together. +var out io.Writer = os.Stderr + +// debugLogger is a sibling of Logger that shares its output and format but is +// held at trace level, so the step tracer (internal/debuglog) is gated by its +// own toggle rather than by the package log level. +var debugLogger = newDebugLogger(out) + +func newDebugLogger(w io.Writer) zerolog.Logger { + return zerolog.New(w).With().Timestamp().Logger().Level(zerolog.TraceLevel) +} + +// DebugLogger returns the trace-level sibling logger used for step tracing. Its +// events are emitted regardless of the package log level; callers gate emission +// with their own flag (see internal/debuglog). +func DebugLogger() *zerolog.Logger { return &debugLogger } + // 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" { - Logger = &DBSQLLogger{Logger.Output(zerolog.ConsoleWriter{Out: os.Stderr})} + out = zerolog.ConsoleWriter{Out: os.Stderr} + Logger = &DBSQLLogger{Logger.Output(out)} + debugLogger = newDebugLogger(out) } // by default only log warns or above loglvl := zerolog.WarnLevel @@ -71,8 +91,11 @@ func SetLogLevel(l string) error { } // Sets logging output. Default is os.Stderr. If in terminal, pretty logs are enabled. +// Redirects both the main logger and the debug-trace sibling so they stay on one sink. func SetLogOutput(w io.Writer) { + out = w Logger.Logger = Logger.Output(w) + debugLogger = newDebugLogger(w) } // Sets log to trace. -1 diff --git a/parameter_test.go b/parameter_test.go index 13e531bb..786623f0 100644 --- a/parameter_test.go +++ b/parameter_test.go @@ -2,16 +2,29 @@ package dbsql import ( "database/sql/driver" - dbsqlerr "github.com/databricks/databricks-sql-go/errors" - "github.com/stretchr/testify/require" "strconv" "testing" "time" + dbsqlerr "github.com/databricks/databricks-sql-go/errors" + thriftbackend "github.com/databricks/databricks-sql-go/internal/backend/thrift" "github.com/databricks/databricks-sql-go/internal/cli_service" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) +// convertNamedValuesToSparkParams composes dbsql-side type inference +// (convertNamedValuesToParams) with the Thrift-side wire mapping +// (thriftbackend.ParamsToSparkForTest) so these tests can assert on the +// end-to-end TSparkParameter output. +func convertNamedValuesToSparkParams(values []driver.NamedValue) ([]*cli_service.TSparkParameter, error) { + params, err := convertNamedValuesToParams(values) + if err != nil { + return nil, err + } + return thriftbackend.ParamsToSparkForTest(params), nil +} + func TestParameter_Inference(t *testing.T) { t.Run("Should infer types correctly", func(t *testing.T) { values := [7]driver.NamedValue{ @@ -214,3 +227,56 @@ func TestParameters_ConvertToSpark(t *testing.T) { require.Equal(t, err.Error(), dbsqlerr.ErrMixedNamedAndPositionalParameters) }) } + +// TestConvertNamedValuesToParams checks the neutral backend.Param output of the +// dbsql-side conversion directly — the boundary the backend consumes. In +// particular a nil value must map to type VOID with a nil Value (SQL NULL), and +// a decimal must carry its derived DECIMAL(p,s) type. This complements the +// end-to-end TSparkParameter assertions above. +func TestConvertNamedValuesToParams(t *testing.T) { + t.Run("infers types and renders neutral params", func(t *testing.T) { + values := []driver.NamedValue{ + {Name: "a", Value: int64(5)}, + {Name: "b", Value: "hello"}, + {Name: "", Value: Parameter{Name: "c", Value: "6.2", Type: SqlDecimal}}, + } + params, err := convertNamedValuesToParams(values) + require.NoError(t, err) + require.Len(t, params, 3) + + require.Equal(t, "a", params[0].Name) + require.Equal(t, "BIGINT", params[0].Type) + require.NotNil(t, params[0].Value) + require.Equal(t, "5", *params[0].Value) + + require.Equal(t, "STRING", params[1].Type) + require.Equal(t, "hello", *params[1].Value) + + require.Equal(t, "DECIMAL(2,1)", params[2].Type) + require.Equal(t, "6.2", *params[2].Value) + }) + + t.Run("nil value maps to VOID with nil Value", func(t *testing.T) { + params, err := convertNamedValuesToParams([]driver.NamedValue{{Value: nil}}) + require.NoError(t, err) + require.Len(t, params, 1) + require.Equal(t, "VOID", params[0].Type) + require.Nil(t, params[0].Value, "SQL NULL must have a nil Value") + }) + + t.Run("empty input yields no params", func(t *testing.T) { + params, err := convertNamedValuesToParams(nil) + require.NoError(t, err) + require.Empty(t, params) + }) + + t.Run("mixed named and positional params error", func(t *testing.T) { + values := []driver.NamedValue{ + {Name: "a", Value: int(1)}, + {Value: int(2)}, + } + _, err := convertNamedValuesToParams(values) + require.Error(t, err) + require.Equal(t, dbsqlerr.ErrMixedNamedAndPositionalParameters, err.Error()) + }) +} diff --git a/parameters.go b/parameters.go index c24027bd..8f777581 100644 --- a/parameters.go +++ b/parameters.go @@ -9,7 +9,7 @@ import ( "time" dbsqlerr "github.com/databricks/databricks-sql-go/errors" - "github.com/databricks/databricks-sql-go/internal/cli_service" + "github.com/databricks/databricks-sql-go/internal/backend" "github.com/pkg/errors" ) @@ -164,8 +164,13 @@ func inferType(param *Parameter) { } } -func convertNamedValuesToSparkParams(values []driver.NamedValue) ([]*cli_service.TSparkParameter, error) { - var sparkParams []*cli_service.TSparkParameter +// convertNamedValuesToParams infers each parameter's SQL type and renders its +// value to the string wire form, returning backend-neutral params. Inference and +// stringification depend on the public Parameter/SqlType surface, so they live +// here; each backend maps the neutral result onto its own wire type. A SqlVoid +// parameter yields a nil Value (SQL NULL). +func convertNamedValuesToParams(values []driver.NamedValue) ([]backend.Param, error) { + var params []backend.Param sqlParams := valuesToParameters(values) @@ -175,10 +180,9 @@ func convertNamedValuesToSparkParams(values []driver.NamedValue) ([]*cli_service inferTypes(sqlParams) for i := range sqlParams { sqlParam := sqlParams[i] - sparkValue := new(cli_service.TSparkParameterValue) - if sqlParam.Type == SqlVoid { - sparkValue = nil - } else { + + var value *string + if sqlParam.Type != SqlVoid { var stringValue string switch v := sqlParam.Value.(type) { case string: @@ -194,22 +198,26 @@ func convertNamedValuesToSparkParams(values []driver.NamedValue) ([]*cli_service default: stringValue = fmt.Sprintf("%v", sqlParam.Value) } - sparkValue = &cli_service.TSparkParameterValue{StringValue: &stringValue} + value = &stringValue } - var sparkParamType string + var paramType string if sqlParam.Type == SqlDecimal { - sparkParamType = inferDecimalType(sparkValue.GetStringValue()) + // The original derived the decimal type from the rendered string + // value; for a nil (VOID) value that string is empty, matching + // TSparkParameterValue.GetStringValue()'s zero value. + var s string + if value != nil { + s = *value + } + paramType = inferDecimalType(s) } else { - sparkParamType = sqlParam.Type.String() + paramType = sqlParam.Type.String() } - var sparkParamName *string if sqlParam.Name != "" { - sparkParamName = &sqlParam.Name hasNamedParams = true } else { - sparkParamName = nil hasPositionalParams = true } @@ -217,10 +225,13 @@ func convertNamedValuesToSparkParams(values []driver.NamedValue) ([]*cli_service return nil, errors.New(dbsqlerr.ErrMixedNamedAndPositionalParameters) } - sparkParam := cli_service.TSparkParameter{Name: sparkParamName, Type: &sparkParamType, Value: sparkValue} - sparkParams = append(sparkParams, &sparkParam) + params = append(params, backend.Param{ + Name: sqlParam.Name, + Type: paramType, + Value: value, + }) } - return sparkParams, nil + return params, nil } func inferDecimalType(d string) (t string) { diff --git a/query_tags.go b/query_tags.go index 6f7ec844..8fa7cca4 100644 --- a/query_tags.go +++ b/query_tags.go @@ -1,6 +1,6 @@ package dbsql -import "strings" +import "github.com/databricks/databricks-sql-go/internal/querytags" // SerializeQueryTags converts a map of query tags to the wire format string. // The format is comma-separated key:value pairs (e.g., "team:engineering,app:etl"). @@ -11,22 +11,9 @@ import "strings" // - Empty string values result in just the key being emitted (no colon) // // Returns empty string if the map is nil or empty. +// +// The implementation lives in internal/querytags so the execution backends can +// share it without importing this package; this remains the public entry point. func SerializeQueryTags(tags map[string]string) string { - if len(tags) == 0 { - return "" - } - - parts := make([]string, 0, len(tags)) - for k, v := range tags { - escapedKey := strings.ReplaceAll(k, `\`, `\\`) - if v == "" { - parts = append(parts, escapedKey) - } else { - escapedValue := strings.ReplaceAll(v, `\`, `\\`) - escapedValue = strings.ReplaceAll(escapedValue, `:`, `\:`) - escapedValue = strings.ReplaceAll(escapedValue, `,`, `\,`) - parts = append(parts, escapedKey+":"+escapedValue) - } - } - return strings.Join(parts, ",") + return querytags.Serialize(tags) } diff --git a/statement_test.go b/statement_test.go index 18a5d41c..73032398 100644 --- a/statement_test.go +++ b/statement_test.go @@ -6,6 +6,7 @@ import ( "testing" "github.com/apache/thrift/lib/go/thrift" + thriftbackend "github.com/databricks/databricks-sql-go/internal/backend/thrift" "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" @@ -99,9 +100,8 @@ func TestStmt_ExecContext(t *testing.T) { FnGetResultSetMetadata: getResultSetMetadata, } testConn := &conn{ - session: getTestSession(), - client: testClient, cfg: config.WithDefaults(), + backend: thriftbackend.NewForTest(testClient, getTestSession(), config.WithDefaults()), } testQuery := "insert 10" testStmt := &stmt{ @@ -154,9 +154,8 @@ func TestStmt_QueryContext(t *testing.T) { FnGetOperationStatus: getOperationStatus, } testConn := &conn{ - session: getTestSession(), - client: testClient, cfg: config.WithDefaults(), + backend: thriftbackend.NewForTest(testClient, getTestSession(), config.WithDefaults()), } testQuery := "select 1" testStmt := &stmt{ From e566e28d2bfeb58dcd907a038388b1bdbd2ed28a Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Sun, 5 Jul 2026 15:55:25 +0000 Subject: [PATCH 2/6] fix: preserve queryId enrichment semantics on the execute path The Layer-1 refactor replaced client.LoggerAndContext (which enriched the context queryId with fill-if-empty semantics) with an inline "if op.StatementID() != \"\" { NewContextWithQueryId(...) }". That changed two observable behaviors for callers who manage their own query id via driverctx: - it overwrote a caller-set queryId with the handle GUID (old code kept the caller's value), and - it skipped NewContextWithQueryId entirely on the no-handle path, so a registered QueryIdCallback never fired (old code always called it, even with an empty id). Restore both via a small enrichQueryId helper used at the ExecContext and QueryContext call sites, and pin the behavior with TestEnrichQueryId. This makes the refactor's "zero behavior change" claim hold on this path. Signed-off-by: Mani Kaustubh Mathur --- connection.go | 30 ++++++++++++++++++------- enrich_queryid_test.go | 50 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 72 insertions(+), 8 deletions(-) create mode 100644 enrich_queryid_test.go diff --git a/connection.go b/connection.go index 27c206e2..596385fc 100644 --- a/connection.go +++ b/connection.go @@ -133,10 +133,11 @@ func (c *conn) ExecContext(ctx context.Context, query string, args []driver.Name log.Err(err).Msgf("databricks: failed to execute query: query %s", query) return nil, err } - // Enrich the logger/context with the statement id. - if sid := op.StatementID(); sid != "" { - ctx = driverctx.NewContextWithQueryId(ctx, sid) - } + // Enrich the logger/context with the statement id, matching the removed + // client.LoggerAndContext semantics: fill the queryId only if the caller + // hasn't set one, and always call NewContextWithQueryId on the empty branch + // so a registered QueryIdCallback fires (even with an empty id). + ctx = enrichQueryId(ctx, op.StatementID()) log = logger.WithContext(driverctx.ConnIdFromContext(ctx), corrId, driverctx.QueryIdFromContext(ctx)) // Telemetry: set up metric context BEFORE staging operation so that the @@ -265,10 +266,11 @@ func (c *conn) QueryContext(ctx context.Context, query string, args []driver.Nam log.Duration(msg, start) return nil, err } - // Enrich the logger/context with the statement id. - if sid := op.StatementID(); sid != "" { - ctx = driverctx.NewContextWithQueryId(ctx, sid) - } + // Enrich the logger/context with the statement id, matching the removed + // client.LoggerAndContext semantics: fill the queryId only if the caller + // hasn't set one, and always call NewContextWithQueryId on the empty branch + // so a registered QueryIdCallback fires (even with an empty id). + ctx = enrichQueryId(ctx, op.StatementID()) log = logger.WithContext(driverctx.ConnIdFromContext(ctx), corrId, driverctx.QueryIdFromContext(ctx)) defer log.Duration(msg, start) @@ -400,6 +402,18 @@ func (c *conn) runQuery(ctx context.Context, query string, args []driver.NamedVa return c.backend.Execute(ctx, backend.ExecRequest{Query: query, Params: params}) } +// enrichQueryId sets the operation's statement id as the context queryId, +// preserving the pre-refactor client.LoggerAndContext behavior: a queryId the +// caller already put on the context is kept (not overwritten), and when the +// context has no queryId, NewContextWithQueryId is always called with the +// derived id — even if it is empty — so any registered QueryIdCallback fires. +func enrichQueryId(ctx context.Context, statementID string) context.Context { + if driverctx.QueryIdFromContext(ctx) != "" { + return ctx + } + return driverctx.NewContextWithQueryId(ctx, statementID) +} + func (c *conn) CheckNamedValue(nv *driver.NamedValue) error { var err error if parameter, ok := nv.Value.(Parameter); ok { diff --git a/enrich_queryid_test.go b/enrich_queryid_test.go new file mode 100644 index 00000000..ac6c4f19 --- /dev/null +++ b/enrich_queryid_test.go @@ -0,0 +1,50 @@ +package dbsql + +import ( + "context" + "testing" + + "github.com/databricks/databricks-sql-go/driverctx" + "github.com/stretchr/testify/assert" +) + +// TestEnrichQueryId pins the queryId-enrichment semantics preserved from the +// removed client.LoggerAndContext: (1) a caller-set queryId is never +// overwritten, and (2) when the context has no queryId, the derived id is always +// applied via NewContextWithQueryId — even when empty — so a registered +// QueryIdCallback fires. The pre-refactor code relied on both; a naive +// "if id != ” { set }" rewrite broke both. +func TestEnrichQueryId(t *testing.T) { + t.Run("preserves a caller-set queryId", func(t *testing.T) { + ctx := driverctx.NewContextWithQueryId(context.Background(), "caller-set") + ctx = enrichQueryId(ctx, "statement-guid") + assert.Equal(t, "caller-set", driverctx.QueryIdFromContext(ctx), + "a queryId the caller already set must not be overwritten") + }) + + t.Run("fills the statement id when the context has none", func(t *testing.T) { + ctx := enrichQueryId(context.Background(), "statement-guid") + assert.Equal(t, "statement-guid", driverctx.QueryIdFromContext(ctx)) + }) + + t.Run("fires the QueryIdCallback even with an empty derived id", func(t *testing.T) { + var fired bool + var got string + ctx := driverctx.NewContextWithQueryIdCallback(context.Background(), func(id string) { + fired = true + got = id + }) + // No handle -> empty statement id. The callback must still fire (old + // LoggerAndContext always called NewContextWithQueryId on the empty branch). + _ = enrichQueryId(ctx, "") + assert.True(t, fired, "QueryIdCallback must fire on the no-queryId path even with an empty id") + assert.Equal(t, "", got) + }) + + t.Run("fires the QueryIdCallback with the derived id", func(t *testing.T) { + var got string + ctx := driverctx.NewContextWithQueryIdCallback(context.Background(), func(id string) { got = id }) + _ = enrichQueryId(ctx, "statement-guid") + assert.Equal(t, "statement-guid", got) + }) +} From e24e903f6f951497d9cf7c68825e057bd3b3a210 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Mon, 6 Jul 2026 06:07:09 +0000 Subject: [PATCH 3/6] fix: guard debugLogger swap with atomic.Pointer MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit SetLogOutput reassigned the package-global debugLogger (a plain zerolog.Logger) while the step tracer read it via DebugLogger() on every event from query goroutines — an unsynchronized read/write of a multi-field struct that -race flags and that can surface a torn logger value. Hold debugLogger in an atomic.Pointer[zerolog.Logger] instead: SetLogOutput stores a fresh logger, DebugLogger() loads a stable snapshot. Matches the atomic treatment already given to the sibling globals (enabled, seq, clockOverride). Add TestConcurrentSetLogOutputIsRaceFree, which reproduces the race (concurrent SetLogOutput during active tracing) and fails under -race on the plain-var version. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- internal/debuglog/debuglog_test.go | 34 ++++++++++++++++++++++++++++++ logger/logger.go | 24 +++++++++++++++------ 2 files changed, 52 insertions(+), 6 deletions(-) diff --git a/internal/debuglog/debuglog_test.go b/internal/debuglog/debuglog_test.go index f02f3a77..77360c32 100644 --- a/internal/debuglog/debuglog_test.go +++ b/internal/debuglog/debuglog_test.go @@ -291,6 +291,40 @@ func TestConcurrentLoggingIsRaceFree(t *testing.T) { wg.Wait() } +// TestConcurrentSetLogOutputIsRaceFree exercises the debug logger swap under +// -race: goroutines emit step traces (reading logger.DebugLogger()) while +// another goroutine calls logger.SetLogOutput, which reinstalls the sibling +// logger. A plain (non-atomic) global would be a torn read/write of the +// multi-field zerolog.Logger struct here. +func TestConcurrentSetLogOutputIsRaceFree(t *testing.T) { + prev := SetEnabled(true) + defer SetEnabled(prev) + defer logger.SetLogOutput(nil) + + var wg sync.WaitGroup + for g := 0; g < 8; g++ { + wg.Add(1) + go func(g int) { + defer wg.Done() + for i := 0; i < 100; i++ { + Logf(context.Background(), "pkg.Fn", "g=%d i=%d", g, i) + done := Track(context.Background(), "pkg.Fn", "step") + done() + } + }(g) + } + // Swap the output sink repeatedly while the tracers above are reading the + // sibling logger on every event. + wg.Add(1) + go func() { + defer wg.Done() + for i := 0; i < 100; i++ { + logger.SetLogOutput(&syncBuf{}) + } + }() + wg.Wait() +} + // syncBuf is a concurrency-safe buffer for capturing log output under -race. type syncBuf struct { mu sync.Mutex diff --git a/logger/logger.go b/logger/logger.go index ab10a6c6..093b18b9 100644 --- a/logger/logger.go +++ b/logger/logger.go @@ -4,6 +4,7 @@ import ( "io" "os" "runtime" + "sync/atomic" "time" "github.com/mattn/go-isatty" @@ -46,17 +47,27 @@ var out io.Writer = os.Stderr // debugLogger is a sibling of Logger that shares its output and format but is // held at trace level, so the step tracer (internal/debuglog) is gated by its -// own toggle rather than by the package log level. -var debugLogger = newDebugLogger(out) +// own toggle rather than by the package log level. Held in an atomic.Pointer +// because SetLogOutput can swap it while the step tracer reads it from query +// goroutines; a plain struct assignment would be a torn read/write under -race. +var debugLogger atomic.Pointer[zerolog.Logger] func newDebugLogger(w io.Writer) zerolog.Logger { return zerolog.New(w).With().Timestamp().Logger().Level(zerolog.TraceLevel) } +// storeDebugLogger atomically installs a fresh trace-level logger over w. +func storeDebugLogger(w io.Writer) { + l := newDebugLogger(w) + debugLogger.Store(&l) +} + // DebugLogger returns the trace-level sibling logger used for step tracing. Its // events are emitted regardless of the package log level; callers gate emission -// with their own flag (see internal/debuglog). -func DebugLogger() *zerolog.Logger { return &debugLogger } +// with their own flag (see internal/debuglog). The returned pointer is a stable +// snapshot — a concurrent SetLogOutput swaps in a new logger without mutating +// the one already handed out. +func DebugLogger() *zerolog.Logger { return debugLogger.Load() } // Enable pretty printing for interactive terminals and json for production. func init() { @@ -64,8 +75,9 @@ func init() { if isatty.IsTerminal(os.Stdout.Fd()) && runtime.GOOS != "windows" { out = zerolog.ConsoleWriter{Out: os.Stderr} Logger = &DBSQLLogger{Logger.Output(out)} - debugLogger = newDebugLogger(out) } + // Install the debug-trace sibling over the (possibly tty-adjusted) sink. + storeDebugLogger(out) // by default only log warns or above loglvl := zerolog.WarnLevel if lvst := os.Getenv("DATABRICKS_LOG_LEVEL"); lvst != "" { @@ -95,7 +107,7 @@ func SetLogLevel(l string) error { func SetLogOutput(w io.Writer) { out = w Logger.Logger = Logger.Output(w) - debugLogger = newDebugLogger(w) + storeDebugLogger(w) } // Sets log to trace. -1 From 0c04d92b65b686ef8d820e09a02602fbcb975d4e Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Mon, 6 Jul 2026 13:03:14 +0000 Subject: [PATCH 4/6] refactor: emit step trace through the main logger, drop the sibling logger The step tracer wrote through a second, trace-level logger (logger.DebugLogger) gated by its own DBSQL_DEBUG_LOG env var. That was a parallel logging knob other maintainers would have to learn, and it forced the atomic.Pointer swap in SetLogOutput to keep the sibling race-free. Route internal/debuglog through the existing logger.Logger at debug level instead, gated by the driver's own SetLogLevel / DATABRICKS_LOG_LEVEL. Step lines are now ordinary level=debug lines on the one sink, in the same JSON/pretty format as every other driver log, so Go and kernel-side stderr logs read as one stream. The helper keeps its real value: the before-the-step enter line (so a step that hangs at the FFI edge is still visible), enter/done elapsed timing, and the fn/phase/conn/corr/query fields. This removes the second logger entirely, so the debugLogger data race is gone by construction rather than patched: logger.go reverts to its pre-refactor shape (no atomic.Pointer, no shared out var), and the concurrent-SetLogOutput race test is dropped as moot. Also drops the now-unused seq counter and clock override (seq only ordered Go-vs-Go lines; cross-boundary interleave rides the shared logger's wall-clock time field). Net -230 lines. go test ./... green; -race clean on debuglog/thrift/logger; CGO_ENABLED=0 build preserved. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- .../thrift/debuglog_integration_test.go | 18 +- internal/debuglog/debuglog.go | 149 ++++-------- internal/debuglog/debuglog_test.go | 226 ++++-------------- logger/logger.go | 37 +-- 4 files changed, 100 insertions(+), 330 deletions(-) diff --git a/internal/backend/thrift/debuglog_integration_test.go b/internal/backend/thrift/debuglog_integration_test.go index 91676b48..097c1510 100644 --- a/internal/backend/thrift/debuglog_integration_test.go +++ b/internal/backend/thrift/debuglog_integration_test.go @@ -7,6 +7,7 @@ package thrift import ( "context" "encoding/json" + "os" "strings" "sync" "testing" @@ -15,14 +16,12 @@ import ( "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/debuglog" "github.com/databricks/databricks-sql-go/logger" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) type logEntry struct { - Seq uint64 `json:"seq"` Fn string `json:"fn"` Phase string `json:"phase"` } @@ -46,12 +45,13 @@ func (c *captureBuf) String() string { } func TestBackend_DebugLoggingOrderedAndNested(t *testing.T) { - prev := debuglog.SetEnabled(true) + prevLevel := logger.Logger.GetLevel() var buf captureBuf logger.SetLogOutput(&buf) + require.NoError(t, logger.SetLogLevel("debug")) t.Cleanup(func() { - debuglog.SetEnabled(prev) - logger.SetLogOutput(nil) + logger.SetLogOutput(os.Stderr) + logger.Logger.Logger = logger.Logger.Level(prevLevel) }) // A successful direct-results execute so we traverse Execute -> runQuery -> @@ -81,12 +81,10 @@ func TestBackend_DebugLoggingOrderedAndNested(t *testing.T) { assert.True(t, hasStep(entries, fn, "done"), "expected a done for %s", fn) } - // Sequence numbers strictly increase across the whole stream — the total - // order that lets Go and kernel-Rust lines interleave into one ordered log. + // Enter+done for all three nested steps land as ordered lines on the shared + // logger — the single stream that kernel-side stderr logs interleave into by + // execution order. require.GreaterOrEqual(t, len(entries), 6, "expected at least enter+done for 3 nested steps") - for i := 1; i < len(entries); i++ { - assert.Greater(t, entries[i].Seq, entries[i-1].Seq, "sequence must strictly increase") - } // The outermost Execute enters before the inner executeStatement enters and // finishes after it — the nesting is reflected in the stream order. diff --git a/internal/debuglog/debuglog.go b/internal/debuglog/debuglog.go index 3afbc05e..ef37fb80 100644 --- a/internal/debuglog/debuglog.go +++ b/internal/debuglog/debuglog.go @@ -5,27 +5,27 @@ // failures are silent and non-local (thread migration, use-after-free, thread // blowup, cross-region latency that looks like "the protocol is slow"). A Go // stack trace stops at the cgo edge, so the way to see which step failed or -// slowed down is an explicit, ordered, timed, per-function log at each step. +// slowed down is an explicit, ordered, timed, per-function log at each step — +// in particular an "enter" line emitted *before* a step runs, so a step that +// hangs at the FFI edge (and never returns) is still visible. // -// Emission is through the driver's existing logger (logger.DebugLogger, a -// trace-level sibling that shares the main logger's sink and format), so there -// is a single logging system and one place to redirect output. It has its own -// DBSQL_DEBUG_LOG flag (below) so it is enabled independently of the package log -// level. +// This is a thin helper over the driver's single logger (logger.Logger): every +// event is an ordinary debug-level line on the same sink, in the same format +// (JSON in production, pretty in a terminal) as the rest of the driver's logs. +// There is no second logger and no separate output knob — turn it on the way +// you turn on any driver logging, with logger.SetLogLevel("debug") or +// DATABRICKS_LOG_LEVEL=debug. That keeps these lines readable next to the +// existing driver logs and lets kernel-side logs (written to the same stderr) +// interleave into one stream in execution order. // -// Every event carries a process-global monotonic sequence number (a total order -// even when timestamps tie), the conn/correlation/query ids from the context, -// the function the step is in, and — for Track — the elapsed time. -// -// Gating: OFF by default (a single atomic-bool load per call site when off). -// Enable via the DBSQL_DEBUG_LOG env var (1/true/yes/on) or SetEnabled. +// Each event carries the function the step is in, the phase ("enter"/"done", +// omitted for a point log), the elapsed time (on "done"), and the +// conn/correlation/query ids from the context. The shared logger supplies the +// timestamp and level fields, so the format matches every other driver line. package debuglog import ( "context" - "os" - "strings" - "sync/atomic" "time" "github.com/rs/zerolog" @@ -34,116 +34,51 @@ import ( "github.com/databricks/databricks-sql-go/logger" ) -// enabled is the single gate. Kept as an atomic so SetEnabled is safe to call -// concurrently with logging (e.g. a test flipping it while a goroutine logs). -var enabled atomic.Bool - -// seq is the process-global sequence counter — the primary ordering key, since -// timestamps can tie or go backwards across cores. -var seq atomic.Uint64 - -// clockOverride optionally replaces the time source so tests can make timing -// deterministic. nil (the default) means use time.Now. Kept atomic — like -// enabled and seq — so an override is safe to set concurrently with logging. -var clockOverride atomic.Pointer[func() time.Time] - -// now returns the current time from the override if one is set, else time.Now. -func now() time.Time { - if f := clockOverride.Load(); f != nil { - return (*f)() - } - return time.Now() -} - -func init() { - if envEnabled(os.Getenv("DBSQL_DEBUG_LOG")) { - enabled.Store(true) - } -} - -// envEnabled reports whether an env-var value means "on". Accepts 1/true/yes/on -// (case-insensitive, trimmed); everything else (including empty) is off. -func envEnabled(v string) bool { - switch strings.ToLower(strings.TrimSpace(v)) { - case "1", "true", "yes", "on": - return true - default: - return false - } -} - -// Enabled reports whether step logging is currently on. Cheap enough to call at -// every step; callers use it to skip building expensive log arguments. -func Enabled() bool { return enabled.Load() } - -// SetEnabled turns step logging on or off at runtime. Returns the previous value -// so a test can save and restore it. -func SetEnabled(on bool) (previous bool) { return enabled.Swap(on) } - -// setClock overrides the time source, or resets to time.Now when f is nil -// (test-only; unexported). -func setClock(f func() time.Time) { - if f == nil { - clockOverride.Store(nil) - return - } - clockOverride.Store(&f) -} +// Enabled reports whether debug-level logging is currently on, i.e. whether +// events would be emitted. Cheap enough to call at every step; callers use it +// to skip building expensive log arguments. It simply reflects the driver's log +// level, so debug tracing follows SetLogLevel like every other driver log. +func Enabled() bool { return logger.Logger.GetLevel() <= zerolog.DebugLevel } -// Logf emits a single ordered, function-tagged step event. fn is the function -// the step is in (e.g. "thrift.Backend.Execute"); the format/args describe the -// step. No-op when logging is disabled — guard the call with Enabled() if the -// args are expensive to build. +// Logf emits a single function-tagged step event at debug level. fn is the +// function the step is in (e.g. "thrift.Backend.Execute"); the format/args +// describe the step. No-op when debug logging is off (the underlying event is +// nil and discards cheaply) — guard the call with Enabled() if the args are +// expensive to build. func Logf(ctx context.Context, fn string, format string, args ...any) { - if !enabled.Load() { - return - } event(ctx, fn, "").Msgf(format, args...) } // Track marks the start of a step and returns a function that logs the step's // elapsed time. Intended for `defer`: // -// defer debuglog.Track(ctx, "thrift.Backend.Execute", "sql=%q", query)() +// defer debuglog.Track(ctx, "thrift.Backend.Execute", "sql.len=%d", len(query))() // -// The enter event is emitted immediately; the returned closure emits a matching -// done event carrying the elapsed duration. Both are no-ops when disabled, and -// the returned closure is always safe to call. +// The enter event is emitted immediately (so a step that never returns is still +// visible); the returned closure emits a matching done event carrying the +// elapsed duration. Both are no-ops when debug logging is off, and the returned +// closure is always safe to call. func Track(ctx context.Context, fn string, format string, args ...any) func() { - if !enabled.Load() { + if !Enabled() { return func() {} } - start := now() + start := time.Now() event(ctx, fn, "enter").Msgf(format, args...) return func() { - // If logging was disabled between enter and here, skip the done line. This - // can leave an orphan enter with no done, which is acceptable for a debug - // trace — SetEnabled is not toggled mid-query in normal use, and emitting a - // done through a now-disabled logger would be worse. - if !enabled.Load() { - return - } - event(ctx, fn, "done").Dur("elapsed", now().Sub(start)).Msg("") + event(ctx, fn, "done").Dur("elapsed", time.Now().Sub(start)).Msg("") } } -// event builds a trace-level zerolog event pre-populated with a nanosecond -// timestamp, the sequence number, the step's function, the phase -// ("enter"/"done", or "" for a point log), and the conn/correlation/query ids -// present on the context. The caller finishes it with Msg/Msgf. Emission goes -// through logger.DebugLogger so the step trace shares the driver's single sink. -// -// The ts field is written as an explicit RFC3339Nano string rather than via -// zerolog's Time(), because zerolog's time rendering is governed by the global -// zerolog.TimeFieldFormat (RFC3339, whole seconds). Latency work needs sub-second -// resolution to correlate steps with server-side traces and to see inter-step -// gaps, so ts carries nanosecond precision independent of the shared logger's -// coarser default `time` field. +// event builds a debug-level zerolog event pre-populated with the step's +// function, the phase ("enter"/"done", or "" for a point log), and the +// conn/correlation/query ids present on the context. The caller finishes it +// with Msg/Msgf. Emission goes through logger.Logger so the step trace shares +// the driver's single sink, format, and timestamp — one ordered stream that +// kernel-side stderr logs interleave into by execution order. Returns a nil +// event (a cheap no-op that ignores all chained calls) when debug logging is +// off. func event(ctx context.Context, fn, phase string) *zerolog.Event { - e := logger.DebugLogger().Log(). - Str("ts", now().Format(time.RFC3339Nano)). - Uint64("seq", seq.Add(1)). - Str("fn", fn) + e := logger.Logger.Debug().Str("fn", fn) if phase != "" { e = e.Str("phase", phase) } diff --git a/internal/debuglog/debuglog_test.go b/internal/debuglog/debuglog_test.go index 77360c32..02b4533e 100644 --- a/internal/debuglog/debuglog_test.go +++ b/internal/debuglog/debuglog_test.go @@ -5,19 +5,19 @@ import ( "context" "encoding/json" "fmt" + "os" "strings" "sync" "testing" - "time" "github.com/databricks/databricks-sql-go/driverctx" "github.com/databricks/databricks-sql-go/logger" ) -// entry is one parsed zerolog line from the debug sink. +// entry is one parsed zerolog line from the shared logger's sink. type entry struct { - TS string `json:"ts"` - Seq uint64 `json:"seq"` + Level string `json:"level"` + Time string `json:"time"` Fn string `json:"fn"` Phase string `json:"phase"` Message string `json:"message"` @@ -28,23 +28,26 @@ type entry struct { raw map[string]any } -// withCapture enables logging into a fresh buffer for the duration of fn, then -// restores the prior enabled state, output, and clock. Every test goes through -// this so global state never leaks between tests. +// withCapture routes the shared logger into a fresh buffer at debug level for +// the duration of fn, then restores the prior level and output. Every test goes +// through this so global logger state never leaks between tests. func withCapture(t *testing.T, fn func(entries func() []entry)) { t.Helper() - prevEnabled := SetEnabled(true) + prevLevel := logger.Logger.GetLevel() var buf syncBuf logger.SetLogOutput(&buf) + if err := logger.SetLogLevel("debug"); err != nil { + t.Fatalf("SetLogLevel: %v", err) + } t.Cleanup(func() { - SetEnabled(prevEnabled) - logger.SetLogOutput(nil) - setClock(nil) // reset to time.Now + logger.SetLogOutput(os.Stderr) + logger.Logger.Logger = logger.Logger.Level(prevLevel) }) fn(func() []entry { return parse(t, buf.String()) }) } -// parse decodes the captured JSON lines into entries. +// parse decodes the captured JSON lines into entries, ignoring any line without +// an fn field (e.g. the logger's own init/level lines). func parse(t *testing.T, out string) []entry { t.Helper() var entries []entry @@ -56,6 +59,9 @@ func parse(t *testing.T, out string) []entry { if err := json.Unmarshal([]byte(line), &raw); err != nil { t.Fatalf("line is not JSON: %q (%v)", line, err) } + if _, ok := raw["fn"]; !ok { + continue + } var e entry if err := json.Unmarshal([]byte(line), &e); err != nil { t.Fatalf("cannot decode entry: %q (%v)", line, err) @@ -66,39 +72,39 @@ func parse(t *testing.T, out string) []entry { return entries } -func TestDisabledByDefaultIsSilent(t *testing.T) { - prev := SetEnabled(false) // force disabled regardless of ambient env / prior tests - defer SetEnabled(prev) - +func TestSilentWhenLevelAboveDebug(t *testing.T) { + prevLevel := logger.Logger.GetLevel() var buf syncBuf logger.SetLogOutput(&buf) - defer logger.SetLogOutput(nil) + if err := logger.SetLogLevel("warn"); err != nil { + t.Fatalf("SetLogLevel: %v", err) + } + t.Cleanup(func() { + logger.SetLogOutput(os.Stderr) + logger.Logger.Logger = logger.Logger.Level(prevLevel) + }) Logf(context.Background(), "pkg.Fn", "should not appear") done := Track(context.Background(), "pkg.Fn", "step") done() - if buf.String() != "" { - t.Fatalf("expected no output when disabled, got: %q", buf.String()) + if es := parse(t, buf.String()); len(es) != 0 { + t.Fatalf("expected no step entries at warn level, got %+v", es) } if Enabled() { - t.Fatal("Enabled() should be false") + t.Fatal("Enabled() should be false at warn level") } } -func TestEnvEnabled(t *testing.T) { - cases := map[string]bool{ - "1": true, "true": true, "TRUE": true, " yes ": true, "On": true, - "": false, "0": false, "false": false, "no": false, "nope": false, - } - for in, want := range cases { - if got := envEnabled(in); got != want { - t.Errorf("envEnabled(%q) = %v, want %v", in, got, want) +func TestEnabledFollowsLogLevel(t *testing.T) { + withCapture(t, func(_ func() []entry) { + if !Enabled() { + t.Fatal("Enabled() should be true at debug level") } - } + }) } -func TestLogfCarriesFunctionMessageAndSeq(t *testing.T) { +func TestLogfCarriesFunctionAndMessage(t *testing.T) { withCapture(t, func(entries func() []entry) { Logf(context.Background(), "thrift.Backend.Execute", "sql=%q rows=%d", "select 1", 42) es := entries() @@ -112,53 +118,30 @@ func TestLogfCarriesFunctionMessageAndSeq(t *testing.T) { if e.Message != `sql="select 1" rows=42` { t.Errorf("message = %q, want formatted step", e.Message) } - if e.Seq == 0 { - t.Errorf("seq should be a positive counter, got %d", e.Seq) - } if _, ok := e.raw["phase"]; ok { t.Errorf("point log should have no phase field: %v", e.raw) } }) } -func TestTimestampIsNanosecondPrecise(t *testing.T) { +// TestFormatMatchesDriverLogs pins the parity the collapse buys: step lines +// carry the same level and time fields as every other driver log line, because +// they go through the same logger. +func TestFormatMatchesDriverLogs(t *testing.T) { withCapture(t, func(entries func() []entry) { - // A fixed instant carrying sub-second nanoseconds — the coarse RFC3339 - // (whole-second) default would drop the fractional part. - fixed := time.Date(2026, 7, 5, 12, 0, 0, 123456789, time.UTC) - setClock(func() time.Time { return fixed }) - Logf(context.Background(), "pkg.Fn", "step") e := entries()[0] - - if e.TS == "" { - t.Fatal("ts field missing") - } - parsed, err := time.Parse(time.RFC3339Nano, e.TS) - if err != nil { - t.Fatalf("ts %q is not RFC3339Nano: %v", e.TS, err) + if e.Level != "debug" { + t.Errorf("level = %q, want debug (shared-logger parity)", e.Level) } - if parsed.Nanosecond() != 123456789 { - t.Errorf("ts lost sub-second precision: got %d ns, want 123456789 (ts=%q)", parsed.Nanosecond(), e.TS) + if e.Time == "" { + t.Error("time field missing — should come from the shared logger") } }) } func TestTrackEmitsEnterAndDoneWithElapsed(t *testing.T) { withCapture(t, func(entries func() []entry) { - // Deterministic clock: enter at t0, done 5ms later. - t0 := time.Date(2026, 7, 3, 12, 0, 0, 0, time.UTC) - times := []time.Time{t0, t0.Add(5 * time.Millisecond)} - i := 0 - setClock(func() time.Time { - idx := i - if idx >= len(times) { - idx = len(times) - 1 - } - i++ - return times[idx] - }) - done := Track(context.Background(), "kernel.Session.Open", "host=%s", "example") done() @@ -172,26 +155,8 @@ func TestTrackEmitsEnterAndDoneWithElapsed(t *testing.T) { if es[0].Message != "host=example" { t.Errorf("enter message = %q, want host=example", es[0].Message) } - // zerolog renders Dur in milliseconds by default; 5ms -> 5. - if es[1].Elapsed != 5 { - t.Errorf("done elapsed = %v ms, want 5", es[1].Elapsed) - } - }) -} - -func TestSequenceIsMonotonic(t *testing.T) { - withCapture(t, func(entries func() []entry) { - for n := 0; n < 5; n++ { - Logf(context.Background(), "pkg.Fn", "line %d", n) - } - es := entries() - if len(es) != 5 { - t.Fatalf("expected 5 entries, got %d", len(es)) - } - for i := 1; i < len(es); i++ { - if es[i].Seq <= es[i-1].Seq { - t.Errorf("sequence not strictly increasing: %d then %d", es[i-1].Seq, es[i].Seq) - } + if _, ok := es[1].raw["elapsed"]; !ok { + t.Errorf("done entry should carry an elapsed field: %v", es[1].raw) } }) } @@ -231,100 +196,6 @@ func TestNilContextIsSafe(t *testing.T) { }) } -func TestDisableMidStepSuppressesDone(t *testing.T) { - withCapture(t, func(entries func() []entry) { - done := Track(context.Background(), "pkg.Fn", "step") - SetEnabled(false) // caller flips it off before the deferred done fires - done() - SetEnabled(true) // restore so cleanup is consistent - es := entries() - if len(es) != 1 || es[0].Phase != "enter" { - t.Fatalf("expected only the enter entry, got %+v", es) - } - }) -} - -// TestConcurrentLoggingIsRaceFree exercises the atomic gate and the atomic clock -// override under -race: many goroutines log while others flip SetEnabled and -// setClock concurrently. -func TestConcurrentLoggingIsRaceFree(t *testing.T) { - prev := SetEnabled(true) - defer SetEnabled(prev) - logger.SetLogOutput(&syncBuf{}) - defer logger.SetLogOutput(nil) - defer setClock(nil) - - var wg sync.WaitGroup - for g := 0; g < 8; g++ { - wg.Add(1) - go func(g int) { - defer wg.Done() - for i := 0; i < 100; i++ { - Logf(context.Background(), "pkg.Fn", "g=%d i=%d", g, i) - done := Track(context.Background(), "pkg.Fn", "step") - done() - } - }(g) - } - wg.Add(1) - go func() { - defer wg.Done() - for i := 0; i < 50; i++ { - SetEnabled(i%2 == 0) - } - SetEnabled(true) - }() - // Flip the clock override concurrently with logging — reads of clockOverride - // happen in every logged event, so this proves the override is race-free too. - wg.Add(1) - go func() { - defer wg.Done() - fixed := func() time.Time { return time.Unix(0, 0) } - for i := 0; i < 50; i++ { - if i%2 == 0 { - setClock(fixed) - } else { - setClock(nil) - } - } - }() - wg.Wait() -} - -// TestConcurrentSetLogOutputIsRaceFree exercises the debug logger swap under -// -race: goroutines emit step traces (reading logger.DebugLogger()) while -// another goroutine calls logger.SetLogOutput, which reinstalls the sibling -// logger. A plain (non-atomic) global would be a torn read/write of the -// multi-field zerolog.Logger struct here. -func TestConcurrentSetLogOutputIsRaceFree(t *testing.T) { - prev := SetEnabled(true) - defer SetEnabled(prev) - defer logger.SetLogOutput(nil) - - var wg sync.WaitGroup - for g := 0; g < 8; g++ { - wg.Add(1) - go func(g int) { - defer wg.Done() - for i := 0; i < 100; i++ { - Logf(context.Background(), "pkg.Fn", "g=%d i=%d", g, i) - done := Track(context.Background(), "pkg.Fn", "step") - done() - } - }(g) - } - // Swap the output sink repeatedly while the tracers above are reading the - // sibling logger on every event. - wg.Add(1) - go func() { - defer wg.Done() - for i := 0; i < 100; i++ { - logger.SetLogOutput(&syncBuf{}) - } - }() - wg.Wait() -} - // syncBuf is a concurrency-safe buffer for capturing log output under -race. type syncBuf struct { mu sync.Mutex @@ -345,8 +216,9 @@ func (s *syncBuf) String() string { // ExampleTrack shows the intended defer-Track idiom. func ExampleTrack() { - SetEnabled(true) - defer SetEnabled(false) + prevLevel := logger.Logger.GetLevel() + _ = logger.SetLogLevel("debug") + defer func() { logger.Logger.Logger = logger.Logger.Level(prevLevel) }() fn := func(ctx context.Context) { defer Track(ctx, "pkg.example", "doing work")() } diff --git a/logger/logger.go b/logger/logger.go index 093b18b9..683501a1 100644 --- a/logger/logger.go +++ b/logger/logger.go @@ -4,7 +4,6 @@ import ( "io" "os" "runtime" - "sync/atomic" "time" "github.com/mattn/go-isatty" @@ -41,43 +40,12 @@ var Logger = &DBSQLLogger{ zerolog.New(os.Stderr).With().Timestamp().Logger(), } -// out is the destination shared by Logger and the debug-trace sibling, so -// SetLogOutput redirects both together. -var out io.Writer = os.Stderr - -// debugLogger is a sibling of Logger that shares its output and format but is -// held at trace level, so the step tracer (internal/debuglog) is gated by its -// own toggle rather than by the package log level. Held in an atomic.Pointer -// because SetLogOutput can swap it while the step tracer reads it from query -// goroutines; a plain struct assignment would be a torn read/write under -race. -var debugLogger atomic.Pointer[zerolog.Logger] - -func newDebugLogger(w io.Writer) zerolog.Logger { - return zerolog.New(w).With().Timestamp().Logger().Level(zerolog.TraceLevel) -} - -// storeDebugLogger atomically installs a fresh trace-level logger over w. -func storeDebugLogger(w io.Writer) { - l := newDebugLogger(w) - debugLogger.Store(&l) -} - -// DebugLogger returns the trace-level sibling logger used for step tracing. Its -// events are emitted regardless of the package log level; callers gate emission -// with their own flag (see internal/debuglog). The returned pointer is a stable -// snapshot — a concurrent SetLogOutput swaps in a new logger without mutating -// the one already handed out. -func DebugLogger() *zerolog.Logger { return debugLogger.Load() } - // 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" { - out = zerolog.ConsoleWriter{Out: os.Stderr} - Logger = &DBSQLLogger{Logger.Output(out)} + Logger = &DBSQLLogger{Logger.Output(zerolog.ConsoleWriter{Out: os.Stderr})} } - // Install the debug-trace sibling over the (possibly tty-adjusted) sink. - storeDebugLogger(out) // by default only log warns or above loglvl := zerolog.WarnLevel if lvst := os.Getenv("DATABRICKS_LOG_LEVEL"); lvst != "" { @@ -103,11 +71,8 @@ func SetLogLevel(l string) error { } // Sets logging output. Default is os.Stderr. If in terminal, pretty logs are enabled. -// Redirects both the main logger and the debug-trace sibling so they stay on one sink. func SetLogOutput(w io.Writer) { - out = w Logger.Logger = Logger.Output(w) - storeDebugLogger(w) } // Sets log to trace. -1 From fe9f43b18df906ef884d75ebc306ce0473cc6405 Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Mon, 6 Jul 2026 13:17:09 +0000 Subject: [PATCH 5/6] fix: use time.Since in Track done closure (gosimple S1012) golangci-lint's gosimple flagged time.Now().Sub(start) in the Track elapsed computation. Use time.Since(start), which is identical in behavior. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- internal/debuglog/debuglog.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/internal/debuglog/debuglog.go b/internal/debuglog/debuglog.go index ef37fb80..86d8b7dd 100644 --- a/internal/debuglog/debuglog.go +++ b/internal/debuglog/debuglog.go @@ -65,7 +65,7 @@ func Track(ctx context.Context, fn string, format string, args ...any) func() { start := time.Now() event(ctx, fn, "enter").Msgf(format, args...) return func() { - event(ctx, fn, "done").Dur("elapsed", time.Now().Sub(start)).Msg("") + event(ctx, fn, "done").Dur("elapsed", time.Since(start)).Msg("") } } From 276f40911ff4aa9dfc6ed89320221a0a0240ce8d Mon Sep 17 00:00:00 2001 From: Mani Kaustubh Mathur Date: Tue, 7 Jul 2026 15:00:52 +0000 Subject: [PATCH 6/6] refactor: address PR #382 code-review-squad findings Doc + test-only response to the multi-perspective review; no seam-behavior change beyond Close idempotency. - backend.Operation: doc close-ownership split (exec path uses Close, query path uses Rows.Close) and require both to be independently idempotent; carve row decoding out of the package-doc neutrality claim. - thriftOperation.Close: add a closed memo so a second call is a no-op (never re-issues CloseOperation, never re-fires CLOSE_STATEMENT). - Restore the end-to-end composed-error assertion the move dropped (unexpected-state + sqlstate across ERROR/CANCELED/TIMEDOUT/CLOSED). - Pin the Execute non-nil-Operation-on-error contract (pre-handle and post-handle terminal-error branches). - Add StatementID and AffectedRows accessor tests. - Collapse the byte-identical newTestBackend into NewForTest (~37 sites). - Rewrite "removed LoggerAndContext" comments to state the contract directly instead of narrating history. Co-authored-by: Isaac Signed-off-by: Mani Kaustubh Mathur --- connection.go | 21 +-- enrich_queryid_test.go | 13 +- internal/backend/backend.go | 31 +++- internal/backend/thrift/backend_test.go | 138 ++++++++++++---- .../thrift/debuglog_integration_test.go | 2 +- internal/backend/thrift/operation.go | 15 +- internal/backend/thrift/operation_test.go | 156 +++++++++++++++++- 7 files changed, 308 insertions(+), 68 deletions(-) diff --git a/connection.go b/connection.go index 596385fc..e88e6c61 100644 --- a/connection.go +++ b/connection.go @@ -133,10 +133,8 @@ func (c *conn) ExecContext(ctx context.Context, query string, args []driver.Name log.Err(err).Msgf("databricks: failed to execute query: query %s", query) return nil, err } - // Enrich the logger/context with the statement id, matching the removed - // client.LoggerAndContext semantics: fill the queryId only if the caller - // hasn't set one, and always call NewContextWithQueryId on the empty branch - // so a registered QueryIdCallback fires (even with an empty id). + // Attach the statement id to the context queryId and refresh the logger so + // subsequent lines carry it. ctx = enrichQueryId(ctx, op.StatementID()) log = logger.WithContext(driverctx.ConnIdFromContext(ctx), corrId, driverctx.QueryIdFromContext(ctx)) @@ -266,10 +264,8 @@ func (c *conn) QueryContext(ctx context.Context, query string, args []driver.Nam log.Duration(msg, start) return nil, err } - // Enrich the logger/context with the statement id, matching the removed - // client.LoggerAndContext semantics: fill the queryId only if the caller - // hasn't set one, and always call NewContextWithQueryId on the empty branch - // so a registered QueryIdCallback fires (even with an empty id). + // Attach the statement id to the context queryId and refresh the logger so + // subsequent lines carry it. ctx = enrichQueryId(ctx, op.StatementID()) log = logger.WithContext(driverctx.ConnIdFromContext(ctx), corrId, driverctx.QueryIdFromContext(ctx)) defer log.Duration(msg, start) @@ -402,11 +398,10 @@ func (c *conn) runQuery(ctx context.Context, query string, args []driver.NamedVa return c.backend.Execute(ctx, backend.ExecRequest{Query: query, Params: params}) } -// enrichQueryId sets the operation's statement id as the context queryId, -// preserving the pre-refactor client.LoggerAndContext behavior: a queryId the -// caller already put on the context is kept (not overwritten), and when the -// context has no queryId, NewContextWithQueryId is always called with the -// derived id — even if it is empty — so any registered QueryIdCallback fires. +// enrichQueryId sets the operation's statement id as the context queryId. A +// caller-set queryId is never overwritten. When the context has no queryId, +// NewContextWithQueryId is always called with the derived id — even if it is +// empty — so any registered QueryIdCallback fires. func enrichQueryId(ctx context.Context, statementID string) context.Context { if driverctx.QueryIdFromContext(ctx) != "" { return ctx diff --git a/enrich_queryid_test.go b/enrich_queryid_test.go index ac6c4f19..e9d06656 100644 --- a/enrich_queryid_test.go +++ b/enrich_queryid_test.go @@ -8,12 +8,10 @@ import ( "github.com/stretchr/testify/assert" ) -// TestEnrichQueryId pins the queryId-enrichment semantics preserved from the -// removed client.LoggerAndContext: (1) a caller-set queryId is never -// overwritten, and (2) when the context has no queryId, the derived id is always -// applied via NewContextWithQueryId — even when empty — so a registered -// QueryIdCallback fires. The pre-refactor code relied on both; a naive -// "if id != ” { set }" rewrite broke both. +// TestEnrichQueryId pins the two non-obvious queryId-enrichment semantics: +// (1) a caller-set queryId is never overwritten, and (2) when the context has +// no queryId, the derived id is always applied via NewContextWithQueryId — +// even when empty — so a registered QueryIdCallback fires. func TestEnrichQueryId(t *testing.T) { t.Run("preserves a caller-set queryId", func(t *testing.T) { ctx := driverctx.NewContextWithQueryId(context.Background(), "caller-set") @@ -34,8 +32,7 @@ func TestEnrichQueryId(t *testing.T) { fired = true got = id }) - // No handle -> empty statement id. The callback must still fire (old - // LoggerAndContext always called NewContextWithQueryId on the empty branch). + // No handle -> empty statement id. The callback must still fire. _ = enrichQueryId(ctx, "") assert.True(t, fired, "QueryIdCallback must fire on the no-queryId path even with an empty id") assert.Equal(t, "", got) diff --git a/internal/backend/backend.go b/internal/backend/backend.go index 79d83cab..37fa1a43 100644 --- a/internal/backend/backend.go +++ b/internal/backend/backend.go @@ -6,7 +6,9 @@ // The interfaces are expressed in neutral terms — plain Go request structs, a // driver.Rows result, and small accessors for statement id, affected rows, and // sqlstate-aware error wrapping — so that only a Backend implementation imports -// its protocol's client (e.g. the Thrift backend imports internal/cli_service). +// its protocol's client. Row decoding is not part of this neutral seam: the +// internal/rows constructor is coupled to Thrift result types, so each Backend +// supplies driver.Rows through its own construction path. package backend import ( @@ -50,6 +52,16 @@ type Backend interface { // Operation is a submitted statement that has reached (or attempted to reach) a // terminal state. It exposes only neutral accessors, keeping the concrete // protocol types inside the backend implementation. +// +// Server-side close ownership is path-dependent. On the exec path (ExecContext, +// staging) the caller closes the server-side operation via Close. On the query +// path (QueryContext) the caller does NOT call Close; instead the driver.Rows +// returned by Results owns closing the server-side operation on its Close. An +// implementation MUST therefore make both this Operation.Close and the +// Rows.Close returned by Results independently idempotent: either one may be +// the sole closer, and both may be invoked on the same operation (e.g. an +// error-path defer plus a normal-path close), so a second invocation of either +// closer must be a no-op rather than issuing a duplicate teardown RPC. type Operation interface { // StatementID is the server statement/operation id (empty if the server // returned no handle). @@ -65,6 +77,11 @@ type Operation interface { // carry the telemetry hooks (chunk timing, close, CloudFetch file); the caller // supplies whichever subset it needs (a full query wires all three, the // staging path wires only chunk timing). + // + // On the query path the returned Rows owns closing the server-side operation: + // the caller does not invoke Operation.Close on that path and relies on + // Rows.Close to issue any needed teardown RPC. See the type-level Operation + // comment for the idempotency contract this places on both closers. Results(ctx context.Context, callbacks *dbsqlrows.TelemetryCallbacks) (driver.Rows, error) // IsStaging reports whether this is a staging (PUT/GET/REMOVE) operation. It @@ -73,10 +90,16 @@ type Operation interface { IsStaging(ctx context.Context) (bool, error) // Close best-effort closes the server-side operation, if one is still open. + // It is called only on the exec/staging paths — on the query path the + // returned Rows owns close (see the Operation and Results doc comments). // It is idempotent and safe to call regardless of whether execution - // succeeded. closed reports whether a close RPC was actually issued (false - // when the operation had no handle or was already closed), so the caller can - // record CLOSE_STATEMENT telemetry only when a real close happened. + // succeeded: a second call on the same Operation, or a call after Rows.Close + // already tore down the server operation, must be a no-op rather than + // issuing a duplicate close RPC. closed reports whether a close RPC was + // actually issued (false when the operation had no handle, was already + // closed via direct results, was already closed by a prior invocation, or + // was already closed via Rows.Close), so the caller can record + // CLOSE_STATEMENT telemetry only when a real close happened. Close(ctx context.Context) (closed bool, err error) // ExecutionError wraps cause as the driver's execution error, attaching this diff --git a/internal/backend/thrift/backend_test.go b/internal/backend/thrift/backend_test.go index 0ad13c3d..ed76c3bc 100644 --- a/internal/backend/thrift/backend_test.go +++ b/internal/backend/thrift/backend_test.go @@ -6,6 +6,7 @@ package thrift import ( "context" "database/sql/driver" + "errors" "fmt" "testing" "time" @@ -17,18 +18,9 @@ import ( "github.com/databricks/databricks-sql-go/internal/config" "github.com/databricks/databricks-sql-go/internal/thrift_protocol" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) -// newTestBackend builds a Backend around a mock client and canned session, -// bypassing New/OpenSession. -func newTestBackend(cli cli_service.TCLIService, session *cli_service.TOpenSessionResp, cfg *config.Config) *Backend { - b := &Backend{cfg: cfg, client: cli, session: session} - if session != nil && session.SessionHandle != nil { - b.sessionID = client.SprintGuid(session.SessionHandle.GetSessionId().GUID) - } - return b -} - // namedToParams converts driver.NamedValues to backend.Params with a minimal // string-typed mapping — enough for these tests, which only exercise param // passthrough, not the full type inference (that lives in the dbsql package). @@ -66,7 +58,7 @@ func TestBackend_executeStatement(t *testing.T) { testClient := &client.TestClient{ FnExecuteStatement: executeStatement, } - be := newTestBackend(testClient, getTestSession(), config.WithDefaults()) + be := NewForTest(testClient, getTestSession(), config.WithDefaults()) _, err := be.executeStatement(context.Background(), backend.ExecRequest{Query: "select 1", Params: namedToParams([]driver.NamedValue{})}) assert.Error(t, err) assert.Equal(t, 1, executeStatementCount) @@ -112,7 +104,7 @@ func TestBackend_executeStatement(t *testing.T) { testClient := &client.TestClient{ FnExecuteStatement: executeStatement, } - be := newTestBackend(testClient, getTestSession(), config.WithDefaults()) + be := NewForTest(testClient, getTestSession(), config.WithDefaults()) _, err := be.executeStatement(context.Background(), backend.ExecRequest{Query: "select 1", Params: namedToParams([]driver.NamedValue{})}) assert.NoError(t, err) @@ -177,7 +169,7 @@ func TestBackend_executeStatement(t *testing.T) { FnExecuteStatement: executeStatement, FnCancelOperation: cancelOperation, } - be := newTestBackend(testClient, getTestSession(), config.WithDefaults()) + be := NewForTest(testClient, getTestSession(), config.WithDefaults()) ctx := context.Background() ctx, cancel = context.WithCancel(ctx) @@ -241,7 +233,7 @@ func TestBackend_executeStatement(t *testing.T) { FnExecuteStatement: executeStatement, FnCancelOperation: cancelOperation, } - be := newTestBackend(testClient, getTestSession(), config.WithDefaults()) + be := NewForTest(testClient, getTestSession(), config.WithDefaults()) ctx := context.Background() ctx, cancel = context.WithCancel(ctx) defer cancel() @@ -348,7 +340,7 @@ func TestBackend_executeStatement_ProtocolFeatures(t *testing.T) { FnExecuteStatement: executeStatement, } - be := newTestBackend(testClient, session, tc.cfg) + be := NewForTest(testClient, session, tc.cfg) var args []driver.NamedValue if tc.hasParameters { @@ -438,7 +430,7 @@ func TestBackend_executeStatement_QueryTags(t *testing.T) { }, nil } - return newTestBackend( + return NewForTest( &client.TestClient{FnExecuteStatement: executeStatement}, getTestSession(), config.WithDefaults(), @@ -588,7 +580,7 @@ func TestBackend_executeStatement_QueryTags(t *testing.T) { // Create conn with session that has session-level tags cfg := config.WithDefaults() cfg.SessionParams = sessionParams - be := newTestBackend(testClient, session, cfg) + be := NewForTest(testClient, session, cfg) // Execute with statement-level tags ctx := driverctx.NewContextWithQueryTags(context.Background(), map[string]string{ @@ -622,7 +614,7 @@ func TestBackend_pollOperation(t *testing.T) { testClient := &client.TestClient{ FnGetOperationStatus: getOperationStatus, } - be := newTestBackend(testClient, getTestSession(), config.WithDefaults()) + be := NewForTest(testClient, getTestSession(), config.WithDefaults()) res, err := be.pollOperation(context.Background(), &cli_service.TOperationHandle{ OperationId: &cli_service.THandleIdentifier{ GUID: []byte{1, 2, 3, 4, 2, 23, 4, 2, 3, 2, 4, 7, 8, 223, 34, 54}, @@ -648,7 +640,7 @@ func TestBackend_pollOperation(t *testing.T) { testClient := &client.TestClient{ FnGetOperationStatus: getOperationStatus, } - be := newTestBackend(testClient, getTestSession(), config.WithDefaults()) + be := NewForTest(testClient, getTestSession(), config.WithDefaults()) res, err := be.pollOperation(context.Background(), &cli_service.TOperationHandle{ OperationId: &cli_service.THandleIdentifier{ GUID: []byte{1, 2, 3, 4, 2, 23, 4, 2, 3, 2, 3, 4, 4, 223, 34, 54}, @@ -674,7 +666,7 @@ func TestBackend_pollOperation(t *testing.T) { testClient := &client.TestClient{ FnGetOperationStatus: getOperationStatus, } - be := newTestBackend(testClient, getTestSession(), config.WithDefaults()) + be := NewForTest(testClient, getTestSession(), config.WithDefaults()) res, err := be.pollOperation(context.Background(), &cli_service.TOperationHandle{ OperationId: &cli_service.THandleIdentifier{ GUID: []byte{1, 2, 3, 4, 2, 23, 4, 2, 3, 2, 3, 4, 4, 223, 34, 54}, @@ -700,7 +692,7 @@ func TestBackend_pollOperation(t *testing.T) { testClient := &client.TestClient{ FnGetOperationStatus: getOperationStatus, } - be := newTestBackend(testClient, getTestSession(), config.WithDefaults()) + be := NewForTest(testClient, getTestSession(), config.WithDefaults()) res, err := be.pollOperation(context.Background(), &cli_service.TOperationHandle{ OperationId: &cli_service.THandleIdentifier{ GUID: []byte{1, 2, 3, 4, 2, 23, 4, 2, 3, 1, 2, 3, 4, 223, 34, 54}, @@ -726,7 +718,7 @@ func TestBackend_pollOperation(t *testing.T) { testClient := &client.TestClient{ FnGetOperationStatus: getOperationStatus, } - be := newTestBackend(testClient, getTestSession(), config.WithDefaults()) + be := NewForTest(testClient, getTestSession(), config.WithDefaults()) res, err := be.pollOperation(context.Background(), &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}, @@ -754,7 +746,7 @@ func TestBackend_pollOperation(t *testing.T) { testClient := &client.TestClient{ FnGetOperationStatus: getOperationStatus, } - be := newTestBackend(testClient, getTestSession(), config.WithDefaults()) + be := NewForTest(testClient, getTestSession(), config.WithDefaults()) res, err := be.pollOperation(context.Background(), &cli_service.TOperationHandle{ OperationId: &cli_service.THandleIdentifier{ GUID: []byte{1, 2, 3, 4, 2, 23, 4, 2, 3, 1, 3, 4, 4, 223, 34, 54}, @@ -792,7 +784,7 @@ func TestBackend_pollOperation(t *testing.T) { FnGetOperationStatus: getOperationStatus, FnCancelOperation: cancelOperation, } - be := newTestBackend(testClient, getTestSession(), config.WithDefaults()) + be := NewForTest(testClient, getTestSession(), config.WithDefaults()) ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) defer cancel() res, err := be.pollOperation(ctx, &cli_service.TOperationHandle{ @@ -834,7 +826,7 @@ func TestBackend_pollOperation(t *testing.T) { } cfg := config.WithDefaults() cfg.PollInterval = 100 * time.Millisecond - be := newTestBackend(testClient, getTestSession(), cfg) + be := NewForTest(testClient, getTestSession(), cfg) ctx, cancel := context.WithTimeout(context.Background(), 250*time.Millisecond) defer cancel() res, err := be.pollOperation(ctx, &cli_service.TOperationHandle{ @@ -876,7 +868,7 @@ func TestBackend_pollOperation(t *testing.T) { } cfg := config.WithDefaults() cfg.PollInterval = 100 * time.Millisecond - be := newTestBackend(testClient, getTestSession(), cfg) + be := NewForTest(testClient, getTestSession(), cfg) ctx, cancel := context.WithCancel(context.Background()) defer cancel() go func() { @@ -907,7 +899,7 @@ func TestBackend_runQuery(t *testing.T) { testClient := &client.TestClient{ FnExecuteStatement: executeStatement, } - be := newTestBackend(testClient, getTestSession(), config.WithDefaults()) + be := NewForTest(testClient, getTestSession(), config.WithDefaults()) exStmtResp, opStatusResp, err := be.runQuery(context.Background(), backend.ExecRequest{Query: "select 1", Params: namedToParams([]driver.NamedValue{})}) assert.Error(t, err) assert.Nil(t, exStmtResp) @@ -945,7 +937,7 @@ func TestBackend_runQuery(t *testing.T) { FnExecuteStatement: executeStatement, FnGetOperationStatus: getOperationStatus, } - be := newTestBackend(testClient, getTestSession(), config.WithDefaults()) + be := NewForTest(testClient, getTestSession(), config.WithDefaults()) exStmtResp, opStatusResp, err := be.runQuery(context.Background(), backend.ExecRequest{Query: "select 1", Params: namedToParams([]driver.NamedValue{})}) assert.Error(t, err) @@ -987,7 +979,7 @@ func TestBackend_runQuery(t *testing.T) { FnExecuteStatement: executeStatement, FnGetOperationStatus: getOperationStatus, } - be := newTestBackend(testClient, getTestSession(), config.WithDefaults()) + be := NewForTest(testClient, getTestSession(), config.WithDefaults()) exStmtResp, opStatusResp, err := be.runQuery(context.Background(), backend.ExecRequest{Query: "select 1", Params: namedToParams([]driver.NamedValue{})}) assert.NoError(t, err) @@ -1030,7 +1022,7 @@ func TestBackend_runQuery(t *testing.T) { FnExecuteStatement: executeStatement, FnGetOperationStatus: getOperationStatus, } - be := newTestBackend(testClient, getTestSession(), config.WithDefaults()) + be := NewForTest(testClient, getTestSession(), config.WithDefaults()) exStmtResp, opStatusResp, err := be.runQuery(context.Background(), backend.ExecRequest{Query: "select 1", Params: namedToParams([]driver.NamedValue{})}) assert.Error(t, err) @@ -1079,7 +1071,7 @@ func TestBackend_runQuery(t *testing.T) { FnExecuteStatement: executeStatement, FnGetOperationStatus: getOperationStatus, } - be := newTestBackend(testClient, getTestSession(), config.WithDefaults()) + be := NewForTest(testClient, getTestSession(), config.WithDefaults()) exStmtResp, opStatusResp, err := be.runQuery(context.Background(), backend.ExecRequest{Query: "select 1", Params: namedToParams([]driver.NamedValue{})}) assert.NoError(t, err) @@ -1127,7 +1119,7 @@ func TestBackend_runQuery(t *testing.T) { FnExecuteStatement: executeStatement, FnGetOperationStatus: getOperationStatus, } - be := newTestBackend(testClient, getTestSession(), config.WithDefaults()) + be := NewForTest(testClient, getTestSession(), config.WithDefaults()) exStmtResp, opStatusResp, err := be.runQuery(context.Background(), backend.ExecRequest{Query: "select 1", Params: namedToParams([]driver.NamedValue{})}) assert.Error(t, err) @@ -1176,7 +1168,7 @@ func TestBackend_runQuery(t *testing.T) { FnExecuteStatement: executeStatement, FnGetOperationStatus: getOperationStatus, } - be := newTestBackend(testClient, getTestSession(), config.WithDefaults()) + be := NewForTest(testClient, getTestSession(), config.WithDefaults()) exStmtResp, opStatusResp, err := be.runQuery(context.Background(), backend.ExecRequest{Query: "select 1", Params: namedToParams([]driver.NamedValue{})}) assert.NoError(t, err) @@ -1225,7 +1217,7 @@ func TestBackend_runQuery(t *testing.T) { FnExecuteStatement: executeStatement, FnGetOperationStatus: getOperationStatus, } - be := newTestBackend(testClient, getTestSession(), config.WithDefaults()) + be := NewForTest(testClient, getTestSession(), config.WithDefaults()) exStmtResp, opStatusResp, err := be.runQuery(context.Background(), backend.ExecRequest{Query: "select 1", Params: namedToParams([]driver.NamedValue{})}) assert.Error(t, err) @@ -1235,3 +1227,81 @@ func TestBackend_runQuery(t *testing.T) { assert.NotNil(t, opStatusResp) }) } + +// TestBackend_Execute pins the load-bearing contract from backend.Backend.Execute: +// even when err is non-nil, the returned Operation MUST be non-nil so the caller +// can uniformly reach the operation-close and telemetry paths. This is what lets +// conn.runQuery use a single `if op == nil` guard to distinguish pre-backend +// failures (param conversion) from post-backend errors that still need close. +func TestBackend_Execute(t *testing.T) { + // Case 1: the pre-handle error path — ExecuteStatement RPC itself fails. + // The returned Operation has no handle, so its handle-less accessor contract + // (StatementID "", Close closed=false with no RPC) must hold. + t.Run("pre-handle failure returns non-nil Operation with handle-less accessors", func(t *testing.T) { + closeCalls := 0 + be := NewForTest( + &client.TestClient{ + FnExecuteStatement: func(ctx context.Context, req *cli_service.TExecuteStatementReq) (*cli_service.TExecuteStatementResp, error) { + return nil, errors.New("boom: rpc failed before a handle was issued") + }, + FnCloseOperation: func(ctx context.Context, req *cli_service.TCloseOperationReq) (*cli_service.TCloseOperationResp, error) { + closeCalls++ + return &cli_service.TCloseOperationResp{}, nil + }, + }, + getTestSession(), + config.WithDefaults(), + ) + op, err := be.Execute(context.Background(), backend.ExecRequest{Query: "select 1"}) + require.Error(t, err) + require.NotNil(t, op, "Execute MUST return a non-nil Operation even on error") + assert.Equal(t, "", op.StatementID(), "handle-less op returns empty statement id") + closed, closeErr := op.Close(context.Background()) + assert.False(t, closed, "handle-less Close must not issue an RPC") + assert.NoError(t, closeErr) + assert.Equal(t, 0, closeCalls, "no CloseOperation RPC on a handle-less op") + }) + + // Case 2: the post-handle error path — the server returns a handle but the + // operation lands in a non-FINISHED terminal state. The Operation must still + // be non-nil (the caller uses it for Close + ExecutionError + telemetry) and + // carry the handle so it can close the server op on the error path. + t.Run("post-handle terminal-error returns non-nil Operation carrying handle+status", func(t *testing.T) { + displayMsg := "server said no" + closeCalls := 0 + be := NewForTest( + &client.TestClient{ + FnExecuteStatement: func(ctx context.Context, req *cli_service.TExecuteStatementReq) (*cli_service.TExecuteStatementResp, error) { + return &cli_service.TExecuteStatementResp{ + Status: &cli_service.TStatus{StatusCode: cli_service.TStatusCode_SUCCESS_STATUS}, + OperationHandle: opHandle(), + DirectResults: &cli_service.TSparkDirectResults{ + OperationStatus: &cli_service.TGetOperationStatusResp{ + OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_ERROR_STATE), + DisplayMessage: &displayMsg, + }, + }, + }, nil + }, + FnCloseOperation: func(ctx context.Context, req *cli_service.TCloseOperationReq) (*cli_service.TCloseOperationResp, error) { + closeCalls++ + return &cli_service.TCloseOperationResp{}, nil + }, + }, + getTestSession(), + config.WithDefaults(), + ) + op, err := be.Execute(context.Background(), backend.ExecRequest{Query: "select 1"}) + require.Error(t, err) + require.NotNil(t, op, "Execute MUST return a non-nil Operation even on error") + assert.NotEmpty(t, op.StatementID(), "post-handle op carries a statement id") + // The caller drives Close on the error path; it must issue the RPC because + // the operation is still open on the server (ERROR_STATE, not CLOSED_STATE). + closed, closeErr := op.Close(context.Background()) + assert.True(t, closed, "post-handle Close on a live ERROR_STATE op must issue the RPC") + assert.NoError(t, closeErr) + assert.Equal(t, 1, closeCalls) + // ExecutionError must also work off the non-nil op and pull sqlstate from opStatusResp. + require.Error(t, op.ExecutionError(context.Background(), err)) + }) +} diff --git a/internal/backend/thrift/debuglog_integration_test.go b/internal/backend/thrift/debuglog_integration_test.go index 097c1510..89de954b 100644 --- a/internal/backend/thrift/debuglog_integration_test.go +++ b/internal/backend/thrift/debuglog_integration_test.go @@ -68,7 +68,7 @@ func TestBackend_DebugLoggingOrderedAndNested(t *testing.T) { }, }, nil } - be := newTestBackend(&client.TestClient{FnExecuteStatement: executeStatement}, getTestSession(), config.WithDefaults()) + be := NewForTest(&client.TestClient{FnExecuteStatement: executeStatement}, getTestSession(), config.WithDefaults()) _, err := be.Execute(context.Background(), backend.ExecRequest{Query: "select 1"}) assert.NoError(t, err) diff --git a/internal/backend/thrift/operation.go b/internal/backend/thrift/operation.go index e2adff83..99b2ffcc 100644 --- a/internal/backend/thrift/operation.go +++ b/internal/backend/thrift/operation.go @@ -26,6 +26,10 @@ type thriftOperation struct { // goroutine at a time (pool discipline), so a plain memo needs no lock. statementID string statementIDSet bool + // closed records whether Close has already issued the CloseOperation RPC on + // this Operation, so a second call is a no-op per the backend.Operation + // idempotency contract. + closed bool } var _ backend.Operation = (*thriftOperation)(nil) @@ -99,12 +103,16 @@ func (o *thriftOperation) IsStaging(ctx context.Context) (bool, error) { } // Close best-effort closes the server operation. It skips the CloseOperation RPC -// when there is no handle, when direct results already closed it, or when the -// operation is already in the CLOSED state; closed reports whether the RPC was -// actually sent, so the caller records CLOSE_STATEMENT telemetry only in that case. +// when there is no handle, when direct results already closed it, when the +// operation is already in the CLOSED state, or when Close has already been +// invoked on this Operation; closed reports whether the RPC was actually sent, +// so the caller records CLOSE_STATEMENT telemetry only in that case. func (o *thriftOperation) Close(ctx context.Context) (bool, error) { defer debuglog.Track(ctx, "thrift.Operation.Close", "stmt=%s", o.StatementID())() + if o.closed { + return false, nil + } if !o.hasHandle() { return false, nil } @@ -120,6 +128,7 @@ func (o *thriftOperation) Close(ctx context.Context) (bool, error) { _, err := o.backend.client.CloseOperation(ctx, &cli_service.TCloseOperationReq{ OperationHandle: o.exStmtResp.OperationHandle, }) + o.closed = true if err != nil { return true, err } diff --git a/internal/backend/thrift/operation_test.go b/internal/backend/thrift/operation_test.go index 172db425..a27a6765 100644 --- a/internal/backend/thrift/operation_test.go +++ b/internal/backend/thrift/operation_test.go @@ -5,6 +5,7 @@ import ( "errors" "testing" + 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" @@ -37,7 +38,7 @@ func TestOperationClose(t *testing.T) { return &cli_service.TCloseOperationResp{}, nil }, } - be := newTestBackend(tc, getTestSession(), config.WithDefaults()) + be := NewForTest(tc, getTestSession(), config.WithDefaults()) return be.OperationForTest(exResp, opStatus) } @@ -113,7 +114,7 @@ func TestOperationClose(t *testing.T) { return nil, errors.New("boom") }, } - be := newTestBackend(tc, getTestSession(), config.WithDefaults()) + be := NewForTest(tc, getTestSession(), config.WithDefaults()) op := be.OperationForTest(&cli_service.TExecuteStatementResp{OperationHandle: opHandle()}, nil) closed, err := op.Close(context.Background()) assert.True(t, closed, "closed=true reflects that the RPC was attempted, so telemetry records it") @@ -121,6 +122,108 @@ func TestOperationClose(t *testing.T) { }) } +// TestOperationStatementID pins the accessor contract: empty string when the +// Operation has no server handle, and the formatted GUID when it does. The +// value is memoized after the first read, so repeated calls return the same +// string cheaply. +func TestOperationStatementID(t *testing.T) { + be := NewForTest(&client.TestClient{}, getTestSession(), config.WithDefaults()) + + t.Run("empty when there is no handle", func(t *testing.T) { + op := be.OperationForTest(&cli_service.TExecuteStatementResp{}, nil) + assert.Equal(t, "", op.StatementID()) + }) + + t.Run("formatted GUID when the handle is present", func(t *testing.T) { + op := be.OperationForTest(&cli_service.TExecuteStatementResp{OperationHandle: opHandle()}, nil) + got := op.StatementID() + // opHandle's GUID is 16 bytes, so SprintGuid renders it dashed. + assert.Equal(t, "01020304-0217-0402-0301-02030404df22", got) + // Repeated calls hit the memo and return the same value. + assert.Equal(t, got, op.StatementID()) + }) +} + +// TestOperationAffectedRows pins the accessor contract on the success path: +// the value is read from opStatusResp.NumModifiedRows via the Thrift getter, +// which returns 0 when the field is unset. AffectedRows is defined only on +// the success path (per backend.Operation), where runQuery has guaranteed a +// non-nil opStatusResp, so a nil-opStatusResp input is out of contract and +// not tested. +func TestOperationAffectedRows(t *testing.T) { + be := NewForTest(&client.TestClient{}, getTestSession(), config.WithDefaults()) + + t.Run("returns NumModifiedRows on the success path", func(t *testing.T) { + rows := int64(42) + op := be.OperationForTest( + &cli_service.TExecuteStatementResp{OperationHandle: opHandle()}, + &cli_service.TGetOperationStatusResp{NumModifiedRows: &rows}, + ) + assert.Equal(t, int64(42), op.AffectedRows()) + }) + + t.Run("returns 0 when NumModifiedRows is unset", func(t *testing.T) { + op := be.OperationForTest( + &cli_service.TExecuteStatementResp{OperationHandle: opHandle()}, + &cli_service.TGetOperationStatusResp{}, + ) + assert.Equal(t, int64(0), op.AffectedRows()) + }) +} + +// TestOperationClose_Idempotency pins the interface contract: a second Close +// call on the same Operation must not re-issue the CloseOperation RPC. This +// guards a defer + explicit-close double-invocation pattern (and any consumer +// that calls Close twice) from re-firing CLOSE_STATEMENT telemetry or wasting +// an RPC. +func TestOperationClose_Idempotency(t *testing.T) { + t.Run("second Close on a live op is a no-op", func(t *testing.T) { + closeCalls := 0 + tc := &client.TestClient{ + FnCloseOperation: func(ctx context.Context, req *cli_service.TCloseOperationReq) (*cli_service.TCloseOperationResp, error) { + closeCalls++ + return &cli_service.TCloseOperationResp{}, nil + }, + } + be := NewForTest(tc, getTestSession(), config.WithDefaults()) + exResp := &cli_service.TExecuteStatementResp{OperationHandle: opHandle()} + opStatus := &cli_service.TGetOperationStatusResp{OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_FINISHED_STATE)} + op := be.OperationForTest(exResp, opStatus) + + closed1, err1 := op.Close(context.Background()) + assert.True(t, closed1) + require.NoError(t, err1) + + closed2, err2 := op.Close(context.Background()) + assert.False(t, closed2, "second Close must not re-issue the RPC") + require.NoError(t, err2) + + assert.Equal(t, 1, closeCalls, "CloseOperation RPC must be issued exactly once") + }) + + t.Run("second Close after RPC error is a no-op", func(t *testing.T) { + closeCalls := 0 + tc := &client.TestClient{ + FnCloseOperation: func(ctx context.Context, req *cli_service.TCloseOperationReq) (*cli_service.TCloseOperationResp, error) { + closeCalls++ + return nil, errors.New("boom") + }, + } + be := NewForTest(tc, getTestSession(), config.WithDefaults()) + op := be.OperationForTest(&cli_service.TExecuteStatementResp{OperationHandle: opHandle()}, nil) + + closed1, err1 := op.Close(context.Background()) + assert.True(t, closed1) + assert.Error(t, err1) + + closed2, err2 := op.Close(context.Background()) + assert.False(t, closed2, "second Close after a failed first Close must not retry the RPC") + require.NoError(t, err2) + + assert.Equal(t, 1, closeCalls, "a failed close attempt is not retried by a second Close call") + }) +} + // TestOperationIsStaging covers the two branches: answer from direct-results // metadata (no RPC) vs. GetResultSetMetadata RPC when direct results are absent. func TestOperationIsStaging(t *testing.T) { @@ -135,7 +238,7 @@ func TestOperationIsStaging(t *testing.T) { return &cli_service.TGetResultSetMetadataResp{}, nil }, } - be := newTestBackend(tc, getTestSession(), config.WithDefaults()) + be := NewForTest(tc, getTestSession(), config.WithDefaults()) exResp := &cli_service.TExecuteStatementResp{ OperationHandle: opHandle(), DirectResults: &cli_service.TSparkDirectResults{ @@ -157,7 +260,7 @@ func TestOperationIsStaging(t *testing.T) { return &cli_service.TGetResultSetMetadataResp{IsStagingOperation: &falsePtr}, nil }, } - be := newTestBackend(tc, getTestSession(), config.WithDefaults()) + be := NewForTest(tc, getTestSession(), config.WithDefaults()) op := be.OperationForTest(&cli_service.TExecuteStatementResp{OperationHandle: opHandle()}, nil) staging, err := op.IsStaging(context.Background()) assert.False(t, staging) @@ -166,10 +269,53 @@ func TestOperationIsStaging(t *testing.T) { }) t.Run("no handle -> false, no RPC", func(t *testing.T) { - be := newTestBackend(&client.TestClient{}, getTestSession(), config.WithDefaults()) + be := NewForTest(&client.TestClient{}, getTestSession(), config.WithDefaults()) op := be.OperationForTest(&cli_service.TExecuteStatementResp{}, nil) staging, err := op.IsStaging(context.Background()) assert.False(t, staging) require.NoError(t, err) }) } + +// TestOperationExecutionError pins the exact user-visible composed error string +// returned when a statement finishes in a non-FINISHED terminal state, plus the +// sqlstate carried on the wrapped DBExecutionError. The string is the driver's +// stable contract for callers that match on it; the composition spans four +// pieces (databricksError.Error's "databricks: execution error: " prefix, +// ErrQueryExecution "failed to execute query", unexpectedOperationState's +// "unexpected operation state STATE" wrap, and opStatusResp.DisplayMessage) so +// any of them regressing must fail this test. +func TestOperationExecutionError(t *testing.T) { + t.Run("nil cause -> nil error", func(t *testing.T) { + be := NewForTest(&client.TestClient{}, getTestSession(), config.WithDefaults()) + op := be.OperationForTest(&cli_service.TExecuteStatementResp{}, nil) + assert.NoError(t, op.ExecutionError(context.Background(), nil)) + }) + + for _, state := range []cli_service.TOperationState{ + cli_service.TOperationState_ERROR_STATE, + cli_service.TOperationState_CANCELED_STATE, + cli_service.TOperationState_TIMEDOUT_STATE, + cli_service.TOperationState_CLOSED_STATE, + } { + t.Run("composed message + sqlstate for "+state.String(), func(t *testing.T) { + displayMsg := "the server said no" + sqlState := "42000" + opStatus := &cli_service.TGetOperationStatusResp{ + OperationState: cli_service.TOperationStatePtr(state), + DisplayMessage: &displayMsg, + SqlState: &sqlState, + } + be := NewForTest(&client.TestClient{}, getTestSession(), config.WithDefaults()) + op := be.OperationForTest(&cli_service.TExecuteStatementResp{OperationHandle: opHandle()}, opStatus) + err := op.ExecutionError(context.Background(), unexpectedOperationState(opStatus)) + require.Error(t, err) + assert.EqualError(t, err, + "databricks: execution error: failed to execute query: unexpected operation state "+ + state.String()+": "+displayMsg) + var execErr dbsqlerr.DBExecutionError + require.True(t, errors.As(err, &execErr), "wrapped error must satisfy DBExecutionError") + assert.Equal(t, sqlState, execErr.SqlState(), "sqlstate must propagate from opStatusResp") + }) + } +}