diff --git a/connection.go b/connection.go index 10d4a0e2..e88e6c61 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,36 @@ 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 + } + // 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)) // 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 +165,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 +246,35 @@ 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 + } + // 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) // 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 +287,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 +375,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 +384,29 @@ 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) - +// 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 { - 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) - } + return nil, dbsqlerrint.NewExecutionError(ctx, dbsqlerr.ErrQueryExecution, err, nil) } + return c.backend.Execute(ctx, backend.ExecRequest{Query: query, Params: params}) } -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), - } +// 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 } - - // 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.ArrowConfig.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) - 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 + return driverctx.NewContextWithQueryId(ctx, statementID) } func (c *conn) CheckNamedValue(nv *driver.NamedValue) error { @@ -916,28 +700,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 +725,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 e549bd24..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.ArrowConfig.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.ArrowConfig.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 f61d6c07..e1701fcb 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/enrich_queryid_test.go b/enrich_queryid_test.go new file mode 100644 index 00000000..e9d06656 --- /dev/null +++ b/enrich_queryid_test.go @@ -0,0 +1,47 @@ +package dbsql + +import ( + "context" + "testing" + + "github.com/databricks/databricks-sql-go/driverctx" + "github.com/stretchr/testify/assert" +) + +// 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") + 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. + _ = 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) + }) +} diff --git a/internal/backend/backend.go b/internal/backend/backend.go new file mode 100644 index 00000000..37fa1a43 --- /dev/null +++ b/internal/backend/backend.go @@ -0,0 +1,132 @@ +// 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. 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 ( + "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. +// +// 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). + 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). + // + // 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 + // 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 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: 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 + // 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..ed76c3bc --- /dev/null +++ b/internal/backend/thrift/backend_test.go @@ -0,0 +1,1307 @@ +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" + "errors" + "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" + "github.com/stretchr/testify/require" +) + +// 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 := 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) + }) + + 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 := NewForTest(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 := NewForTest(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 := NewForTest(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 := NewForTest(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 NewForTest( + &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 := NewForTest(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 := 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}, + 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 := 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}, + 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 := 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}, + 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 := 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}, + 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 := 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}, + 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 := 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}, + 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 := NewForTest(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 := NewForTest(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 := NewForTest(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 := 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) + 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 := 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.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 := 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) + 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 := 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.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 := 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) + 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 := 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.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 := 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) + 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 := 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.Equal(t, 1, executeStatementCount) + assert.Equal(t, 1, getOperationStatusCount) + assert.NotNil(t, exStmtResp) + 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 new file mode 100644 index 00000000..89de954b --- /dev/null +++ b/internal/backend/thrift/debuglog_integration_test.go @@ -0,0 +1,128 @@ +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" + "os" + "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/logger" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +type logEntry struct { + 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) { + prevLevel := logger.Logger.GetLevel() + var buf captureBuf + logger.SetLogOutput(&buf) + require.NoError(t, logger.SetLogLevel("debug")) + t.Cleanup(func() { + logger.SetLogOutput(os.Stderr) + logger.Logger.Logger = logger.Logger.Level(prevLevel) + }) + + // 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 := NewForTest(&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) + } + + // 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") + + // 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..99b2ffcc --- /dev/null +++ b/internal/backend/thrift/operation.go @@ -0,0 +1,146 @@ +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 + // 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) + +// 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, 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 + } + 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, + }) + o.closed = true + 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..a27a6765 --- /dev/null +++ b/internal/backend/thrift/operation_test.go @@ -0,0 +1,321 @@ +package thrift + +import ( + "context" + "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" + "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 := NewForTest(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 := 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") + assert.Error(t, err) + }) +} + +// 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) { + 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 := NewForTest(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 := NewForTest(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 := 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") + }) + } +} 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..86d8b7dd --- /dev/null +++ b/internal/debuglog/debuglog.go @@ -0,0 +1,97 @@ +// 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 — +// 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. +// +// 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. +// +// 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" + "time" + + "github.com/rs/zerolog" + + "github.com/databricks/databricks-sql-go/driverctx" + "github.com/databricks/databricks-sql-go/logger" +) + +// 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 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) { + 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.len=%d", len(query))() +// +// 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() { + return func() {} + } + start := time.Now() + event(ctx, fn, "enter").Msgf(format, args...) + return func() { + event(ctx, fn, "done").Dur("elapsed", time.Since(start)).Msg("") + } +} + +// 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.Logger.Debug().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..02b4533e --- /dev/null +++ b/internal/debuglog/debuglog_test.go @@ -0,0 +1,228 @@ +package debuglog + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "os" + "strings" + "sync" + "testing" + + "github.com/databricks/databricks-sql-go/driverctx" + "github.com/databricks/databricks-sql-go/logger" +) + +// entry is one parsed zerolog line from the shared logger's sink. +type entry struct { + Level string `json:"level"` + Time string `json:"time"` + 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 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() + 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() { + 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, 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 + 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) + } + 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) + } + e.raw = raw + entries = append(entries, e) + } + return entries +} + +func TestSilentWhenLevelAboveDebug(t *testing.T) { + prevLevel := logger.Logger.GetLevel() + var buf syncBuf + logger.SetLogOutput(&buf) + 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 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 at warn level") + } +} + +func TestEnabledFollowsLogLevel(t *testing.T) { + withCapture(t, func(_ func() []entry) { + if !Enabled() { + t.Fatal("Enabled() should be true at debug level") + } + }) +} + +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() + 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 _, ok := e.raw["phase"]; ok { + t.Errorf("point log should have no phase field: %v", e.raw) + } + }) +} + +// 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) { + Logf(context.Background(), "pkg.Fn", "step") + e := entries()[0] + if e.Level != "debug" { + t.Errorf("level = %q, want debug (shared-logger parity)", e.Level) + } + 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) { + 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) + } + if _, ok := es[1].raw["elapsed"]; !ok { + t.Errorf("done entry should carry an elapsed field: %v", es[1].raw) + } + }) +} + +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) + } + }) +} + +// 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() { + 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")() + } + 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/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{