refactor: extract Backend/Operation seam for SEA integration

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
 }

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)
+	})
+}

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
+}