From a229ad0bd7ef070dfc7315cb6aebb88c1e90e239 Mon Sep 17 00:00:00 2001 From: Rahul Singhal Date: Wed, 9 Sep 2026 21:51:37 +0000 Subject: [PATCH 1/6] feat: auto-recover Reyden Thrift connections onto the kernel backend An unconfigured connection to a Reyden / Real-Time SQL warehouse defaults to the Thrift backend, which the SQL Gateway proxy rejects with SQLSTATE KP001. Detect that rejection at OpenSession (CheckStatus) and transparently re-open the session on the SEA/kernel backend, remembering the warehouse in a process-wide cache keyed by (host, warehouse_id) with a ~6h TTL so later connects skip the doomed Thrift attempt. Only the default path auto-recovers; an explicit WithUseKernel is always honored. On a double failure the kernel error is surfaced with the Thrift rejection preserved via errors.Join. Recovery requires a databricks_kernel build, since the kernel backend is otherwise not linked in (the fallback then surfaces the not-compiled error joined with the Reyden rejection). Co-authored-by: Isaac Signed-off-by: Rahul Singhal --- CHANGELOG.md | 3 + connector.go | 174 +++++++++--- connector_reyden_test.go | 368 +++++++++++++++++++++++++ errors/errors.go | 7 + internal/client/client.go | 6 + internal/errors/err.go | 32 +++ internal/warehouse_cache/cache.go | 138 ++++++++++ internal/warehouse_cache/cache_test.go | 168 +++++++++++ 8 files changed, 864 insertions(+), 32 deletions(-) create mode 100644 connector_reyden_test.go create mode 100644 internal/warehouse_cache/cache.go create mode 100644 internal/warehouse_cache/cache_test.go diff --git a/CHANGELOG.md b/CHANGELOG.md index 7f321ff5..fee7a63a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,8 @@ # Release History +## Unreleased +- Transparently auto-recover Thrift connections to Reyden / Real-Time warehouses: when a warehouse rejects the default Thrift protocol (SQLSTATE `KP001`), the session is re-opened on the SEA/kernel backend and the warehouse is remembered so later connections skip Thrift. Applies only when no backend was chosen explicitly (`WithUseKernel`). Note: recovery requires a `databricks_kernel` build, since the kernel backend is otherwise not linked in. + ## v1.15.1 (2026-09-01) - Pin the seven per-platform kernel bindings modules to v1.0.0. - Disable kernel telemetry by default when `enableTelemetry` is unset; explicit `true` and `false` values are unchanged (databricks/databricks-sql-go#464). diff --git a/connector.go b/connector.go index 6e502b20..382af388 100644 --- a/connector.go +++ b/connector.go @@ -4,6 +4,7 @@ import ( "context" "crypto/tls" "database/sql/driver" + "errors" "fmt" "net/http" "net/url" @@ -22,13 +23,26 @@ import ( "github.com/databricks/databricks-sql-go/internal/client" "github.com/databricks/databricks-sql-go/internal/config" "github.com/databricks/databricks-sql-go/internal/debuglog" + "github.com/databricks/databricks-sql-go/internal/warehouse_cache" "github.com/databricks/databricks-sql-go/logger" "github.com/databricks/databricks-sql-go/telemetry" ) +// backendFactory is a function type for creating kernel backends. This seam allows +// tests to inject a fake kernel backend without requiring the build-tag gated +// newKernelBackend. In production, the factory is nil and newKernelBackend is used directly. +type backendFactory func(ctx context.Context, cfg *config.Config) (backend.Backend, error) + +// thriftBackendFactory is the analogous seam for the Thrift backend, letting tests +// inject a fake Thrift backend (e.g. one that rejects OpenSession with the Reyden +// marker) instead of the real thrift.New. In production the factory is nil. +type thriftBackendFactory func(ctx context.Context, cfg *config.Config, client *http.Client) (backend.Backend, error) + type connector struct { - cfg *config.Config - client *http.Client + cfg *config.Config + client *http.Client + kernelBackendFactory backendFactory // Seam for testing; nil in production + thriftBackendFactory thriftBackendFactory // Seam for testing; nil in production } func skipDriverTelemetry(cfg *config.Config) bool { @@ -46,40 +60,14 @@ type federatedTokenAuthenticator struct { func (c *connector) Connect(ctx context.Context) (driver.Conn, error) { defer debuglog.Track(ctx, "connector.Connect", "host=%s", c.cfg.Host)() - // Build the execution backend. Thrift is the default; the SEA-via-kernel - // backend is selected when UseKernel is set. newKernelBackend is build-tag - // gated: in the default pure-Go build it returns a clear "not linked in" - // error, so the kernel path compiles and links only under -tags - // databricks_kernel + CGO_ENABLED=1. - var be backend.Backend - var err error - if c.cfg.UseKernel { - be, err = newKernelBackend(ctx, c.cfg) - } else { - // The experimental WithKernel* options have no Thrift-path equivalent — reject - // them loudly rather than silently ignore, so a caller who sets one (a - // trusted-CA bundle, a hostname-verify skip, a proxy, a retry budget, or a - // CloudFetch chunk cap) and forgets WithUseKernel learns the option had no - // effect instead of connecting as if it were never set. Every WithKernel* - // option allocates KernelExperimental, so this one gate covers them all; the - // message names the family rather than a stale subset that drifts as options - // are added. - if c.cfg.KernelExperimental != nil { - return nil, fmt.Errorf("databricks: a WithKernel* option %w; "+ - "add WithUseKernel(true) or remove it", dbsqlerr.ErrRequiresKernelBackend) - } - be, err = thrift.New(ctx, c.cfg, c.client) - } + // openSessionWithReydenFallback handles the session opening with automatic + // recovery for Reyden / Real-Time warehouses that reject Thrift. It returns + // the backend, latency, and error. + be, sessionLatencyMs, err := c.openSessionWithReydenFallback(ctx) if err != nil { return nil, err } - sessionStart := time.Now() - if err := be.OpenSession(ctx); err != nil { - return nil, err - } - sessionLatencyMs := time.Since(sessionStart).Milliseconds() - conn := &conn{ id: be.SessionID(), cfg: c.cfg, @@ -139,6 +127,24 @@ func (c *connector) Driver() driver.Driver { return &databricksDriver{} } +// getKernelBackend returns a kernel backend, using the injected factory if available +// (for tests), otherwise using the build-tag-gated newKernelBackend (production). +func (c *connector) getKernelBackend(ctx context.Context) (backend.Backend, error) { + if c.kernelBackendFactory != nil { + return c.kernelBackendFactory(ctx, c.cfg) + } + return newKernelBackend(ctx, c.cfg) +} + +// getThriftBackend returns a Thrift backend, using the injected factory if available +// (for tests), otherwise the real thrift.New (production). +func (c *connector) getThriftBackend(ctx context.Context) (backend.Backend, error) { + if c.thriftBackendFactory != nil { + return c.thriftBackendFactory(ctx, c.cfg, c.client) + } + return thrift.New(ctx, c.cfg, c.client) +} + var _ driver.Connector = (*connector)(nil) type ConnOption func(*config.Config) @@ -802,3 +808,107 @@ func WithTokenCache(enabled bool) ConnOption { kernelExperimental(c).TokenCacheEnabled = enabled } } + +// openSessionWithReydenFallback opens a session with automatic recovery for +// Reyden / Real-Time warehouses that reject the Thrift protocol. When a +// warehouse rejects the default Thrift OpenSession (SQLSTATE KP001), it +// transparently re-opens on the kernel backend and remembers the warehouse +// so later connections skip the doomed Thrift attempt. +// +// Auto-recovery applies only when no backend was chosen explicitly (neither +// UseKernel nor other backend-selecting options). On success or failure, it +// returns the backend, session latency, and error. +func (c *connector) openSessionWithReydenFallback(ctx context.Context) (backend.Backend, int64, error) { + // Extract warehouse ID from HTTPPath for cache lookups. + warehouseID := warehouse_cache.ExtractWarehouseID(c.cfg.HTTPPath) + + // Pre-check: if this warehouse is already known to reject Thrift, open + // directly on the kernel backend and skip the doomed Thrift attempt. + if warehouseID != "" && warehouse_cache.IsKnownReyden(c.cfg.Host, warehouseID) { + logger.Debug().Msgf( + "warehouse %s on %s is known to require kernel backend; skipping Thrift", + warehouseID, c.cfg.Host) + var be backend.Backend + var err error + if be, err = c.getKernelBackend(ctx); err != nil { + return nil, 0, err + } + sessionStart := time.Now() + if err := be.OpenSession(ctx); err != nil { + return nil, 0, err + } + return be, time.Since(sessionStart).Milliseconds(), nil + } + + // Build the execution backend. Thrift is the default; the SEA-via-kernel + // backend is selected when UseKernel is set. newKernelBackend is build-tag + // gated: in the default pure-Go build it returns a clear "not linked in" + // error, so the kernel path compiles and links only under -tags + // databricks_kernel + CGO_ENABLED=1. + var be backend.Backend + var err error + if c.cfg.UseKernel { + be, err = c.getKernelBackend(ctx) + } else { + // The experimental WithKernel* options have no Thrift-path equivalent — reject + // them loudly rather than silently ignore, so a caller who sets one (a + // trusted-CA bundle, a hostname-verify skip, a proxy, a retry budget, or a + // CloudFetch chunk cap) and forgets WithUseKernel learns the option had no + // effect instead of connecting as if it were never set. Every WithKernel* + // option allocates KernelExperimental, so this one gate covers them all; the + // message names the family rather than a stale subset that drifts as options + // are added. + if c.cfg.KernelExperimental != nil { + return nil, 0, fmt.Errorf("databricks: a WithKernel* option %w; "+ + "add WithUseKernel(true) or remove it", dbsqlerr.ErrRequiresKernelBackend) + } + be, err = c.getThriftBackend(ctx) + } + if err != nil { + return nil, 0, err + } + + // Attempt to open the session. If an explicit backend was selected, honor it + // even on a Reyden rejection — auto-recovery applies only on the default path. + sessionStart := time.Now() + err = be.OpenSession(ctx) + sessionLatency := time.Since(sessionStart).Milliseconds() + + // Check for Reyden Thrift rejection on the default (non-UseKernel) path. + // If detected, mark the warehouse and retry on the kernel backend. + if err != nil && !c.cfg.UseKernel && errors.Is(err, dbsqlerr.ErrReydenThriftUnsupported) { + logger.Info().Msg("Thrift is not supported for this Reyden/Real-Time warehouse; " + + "transparently re-opening the session on the kernel backend") + + // Remember the rejection regardless of the retry's outcome — the + // warehouse is Reyden either way, so future connects should skip Thrift; + // a kernel failure below is a separate, orthogonal problem. + if warehouseID != "" { + warehouse_cache.MarkReyden(c.cfg.Host, warehouseID) + } + + // Retry on the kernel backend. + var kernelConstructErr error + be, kernelConstructErr = c.getKernelBackend(ctx) + if kernelConstructErr != nil { + // Surface the kernel construction error, wrapped with the original + // Thrift rejection in the chain for diagnosis. + return nil, 0, errors.Join(kernelConstructErr, err) + } + + kernelStart := time.Now() + kernelOpenErr := be.OpenSession(ctx) + kernelLatency := time.Since(kernelStart).Milliseconds() + + if kernelOpenErr != nil { + // Surface the kernel failure (the actionable one) while keeping + // the original Thrift rejection in the chain for diagnosis. + return nil, kernelLatency, errors.Join(kernelOpenErr, err) + } + return be, kernelLatency, nil + } + + // Either the session opened successfully or failed for a reason other than + // Reyden rejection on the default Thrift path. Return the outcome as-is. + return be, sessionLatency, err +} diff --git a/connector_reyden_test.go b/connector_reyden_test.go new file mode 100644 index 00000000..ec2303b1 --- /dev/null +++ b/connector_reyden_test.go @@ -0,0 +1,368 @@ +package dbsql + +import ( + "context" + "errors" + "net/http" + "testing" + + 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/config" + dbsqlerrint "github.com/databricks/databricks-sql-go/internal/errors" + "github.com/databricks/databricks-sql-go/internal/warehouse_cache" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestCheckStatusReydenDetection tests that CheckStatus detects SQLSTATE KP001 +// and returns a distinct error vs. other ERROR_STATUS codes. +func TestCheckStatusReydenDetection(t *testing.T) { + t.Run("KP001 on ERROR_STATUS returns ReydenThriftUnsupported marker", func(t *testing.T) { + sqlState := "KP001" + errMsg := "Lakehouse/RT is not supported for Thrift protocol" + reydenResp := &cli_service.TOpenSessionResp{ + Status: &cli_service.TStatus{ + StatusCode: cli_service.TStatusCode_ERROR_STATUS, + SqlState: &sqlState, + ErrorMessage: &errMsg, + }, + } + + err := client.CheckStatus(reydenResp) + require.Error(t, err) + // Verify it's the Reyden marker. + assert.True(t, errors.Is(err, dbsqlerr.ErrReydenThriftUnsupported), + "error should satisfy errors.Is for ErrReydenThriftUnsupported") + }) + + t.Run("non-KP001 ERROR_STATUS returns generic error, NOT Reyden marker", func(t *testing.T) { + sqlState := "42000" // Syntax error + errMsg := "a syntax error" + syntaxResp := &cli_service.TOpenSessionResp{ + Status: &cli_service.TStatus{ + StatusCode: cli_service.TStatusCode_ERROR_STATUS, + SqlState: &sqlState, + ErrorMessage: &errMsg, + }, + } + + err := client.CheckStatus(syntaxResp) + require.Error(t, err) + // Verify it's NOT the Reyden marker. + assert.False(t, errors.Is(err, dbsqlerr.ErrReydenThriftUnsupported), + "generic error should NOT satisfy errors.Is for ErrReydenThriftUnsupported") + }) + + t.Run("SUCCESS_STATUS returns no error", func(t *testing.T) { + successResp := &cli_service.TOpenSessionResp{ + Status: &cli_service.TStatus{ + StatusCode: cli_service.TStatusCode_SUCCESS_STATUS, + }, + } + + err := client.CheckStatus(successResp) + assert.NoError(t, err) + }) +} + +// fakeKernelBackend is a mock kernel backend for testing the fallback logic. +type fakeKernelBackend struct { + openSessionErr error + sessionID string +} + +var _ backend.Backend = (*fakeKernelBackend)(nil) + +func (f *fakeKernelBackend) OpenSession(ctx context.Context) error { + if f.openSessionErr != nil { + return f.openSessionErr + } + f.sessionID = "kernel-sess-id" + return nil +} + +func (f *fakeKernelBackend) CloseSession(ctx context.Context) error { + return nil +} + +func (f *fakeKernelBackend) SessionValid() bool { + return f.sessionID != "" +} + +func (f *fakeKernelBackend) SessionID() string { + return f.sessionID +} + +func (f *fakeKernelBackend) Execute(ctx context.Context, req backend.ExecRequest) (backend.Operation, error) { + return nil, errors.New("not implemented") +} + +// fakeThriftBackend is a mock Thrift backend that can be configured to fail with KP001. +type fakeThriftBackend struct { + openSessionErr error + sessionID string +} + +var _ backend.Backend = (*fakeThriftBackend)(nil) + +func (f *fakeThriftBackend) OpenSession(ctx context.Context) error { + return f.openSessionErr +} + +func (f *fakeThriftBackend) CloseSession(ctx context.Context) error { + return nil +} + +func (f *fakeThriftBackend) SessionValid() bool { + return f.sessionID != "" +} + +func (f *fakeThriftBackend) SessionID() string { + return f.sessionID +} + +func (f *fakeThriftBackend) Execute(ctx context.Context, req backend.ExecRequest) (backend.Operation, error) { + return nil, errors.New("not implemented") +} + +// TestReydenFallback tests the connector's openSessionWithReydenFallback logic. +type TestReydenFallback struct { + host string + warehousePath string + warehouseID string +} + +func NewTestReydenFallback() *TestReydenFallback { + return &TestReydenFallback{ + host: "reyden.example.com", + warehousePath: "/sql/1.0/warehouses/wh-reyden", + warehouseID: "wh-reyden", + } +} + +// makeConnector builds a connector with injected backend factories. Either +// factory may be nil (the path that isn't exercised by a given test). +func (t *TestReydenFallback) makeConnector( + thriftFactory thriftBackendFactory, + kernelFactory backendFactory, +) *connector { + cfg := config.WithDefaults() + cfg.Host = t.host + cfg.HTTPPath = t.warehousePath + return &connector{ + cfg: cfg, + thriftBackendFactory: thriftFactory, + kernelBackendFactory: kernelFactory, + } +} + +func TestReydenReactiveRecovery(t *testing.T) { + t.Run("Thrift KP001 rejection recovers onto the kernel backend", func(t *testing.T) { + defer warehouse_cache.ClearCache() + + tt := NewTestReydenFallback() + + // Fake Thrift backend that rejects OpenSession with the Reyden marker. + fakeThrift := &fakeThriftBackend{ + openSessionErr: dbsqlerrint.NewReydenThriftUnsupportedError( + "Lakehouse/RT is not supported for Thrift protocol"), + } + kernelBackend := &fakeKernelBackend{} + + conn := tt.makeConnector( + func(ctx context.Context, cfg *config.Config, client *http.Client) (backend.Backend, error) { + return fakeThrift, nil + }, + func(ctx context.Context, cfg *config.Config) (backend.Backend, error) { + return kernelBackend, nil + }, + ) + + be, _, err := conn.openSessionWithReydenFallback(context.Background()) + require.NoError(t, err) + + // The returned backend is the kernel one, with an open session. + kb, ok := be.(*fakeKernelBackend) + require.True(t, ok, "recovery should return the kernel backend, got %T", be) + assert.Same(t, kernelBackend, kb) + assert.True(t, kb.SessionValid(), "kernel session should be open") + + // The rejection is remembered for future connects. + assert.True(t, warehouse_cache.IsKnownReyden(tt.host, tt.warehouseID), + "warehouse should be marked Reyden after the rejection") + }) +} + +func TestReydenPreCheck(t *testing.T) { + t.Run("Known Reyden warehouse opens kernel directly without Thrift attempt", func(t *testing.T) { + defer warehouse_cache.ClearCache() + + tt := NewTestReydenFallback() + + // Mark the warehouse as known-Reyden so the pre-check fires. + warehouse_cache.MarkReyden(tt.host, tt.warehouseID) + + thriftCalled := false + kernelCalled := false + conn := tt.makeConnector( + func(ctx context.Context, cfg *config.Config, client *http.Client) (backend.Backend, error) { + thriftCalled = true + return &fakeThriftBackend{}, nil + }, + func(ctx context.Context, cfg *config.Config) (backend.Backend, error) { + kernelCalled = true + return &fakeKernelBackend{}, nil + }, + ) + + be, latency, err := conn.openSessionWithReydenFallback(context.Background()) + + assert.NoError(t, err) + assert.NotNil(t, be) + assert.True(t, kernelCalled, "pre-check should open the kernel backend") + assert.False(t, thriftCalled, "pre-check must skip the Thrift OpenSession round-trip") + assert.GreaterOrEqual(t, latency, int64(0), "latency should be non-negative") + }) +} + +func TestReydenCacheMarking(t *testing.T) { + t.Run("Rejection marks warehouse in cache", func(t *testing.T) { + defer warehouse_cache.ClearCache() + + tt := NewTestReydenFallback() + + // Verify warehouse is not initially marked. + assert.False(t, warehouse_cache.IsKnownReyden(tt.host, tt.warehouseID)) + + // Mark it. + warehouse_cache.MarkReyden(tt.host, tt.warehouseID) + + // Verify it's now marked. + assert.True(t, warehouse_cache.IsKnownReyden(tt.host, tt.warehouseID)) + }) + + t.Run("Warehouse ID extraction from HTTP path", func(t *testing.T) { + tests := []struct { + path string + expected string + }{ + {"/sql/1.0/warehouses/wh-123", "wh-123"}, + {"/sql/1.0/endpoints/ep-456", "ep-456"}, + {"/sql/1.0/warehouses/wh-123?o=789", "wh-123"}, + {"/sql/protocolv1/o/123/cluster", ""}, // Cluster path, no warehouse ID + } + + for _, tc := range tests { + t.Run(tc.path, func(t *testing.T) { + got := warehouse_cache.ExtractWarehouseID(tc.path) + assert.Equal(t, tc.expected, got) + }) + } + }) +} + +func TestReydenGuardrail(t *testing.T) { + t.Run("Explicit UseKernel goes straight to kernel, never touching Thrift", func(t *testing.T) { + defer warehouse_cache.ClearCache() + + tt := NewTestReydenFallback() + + thriftCalled := false + kernelBackend := &fakeKernelBackend{} + conn := tt.makeConnector( + func(ctx context.Context, cfg *config.Config, client *http.Client) (backend.Backend, error) { + thriftCalled = true + return &fakeThriftBackend{}, nil + }, + func(ctx context.Context, cfg *config.Config) (backend.Backend, error) { + return kernelBackend, nil + }, + ) + conn.cfg.UseKernel = true // explicit backend selection + + be, _, err := conn.openSessionWithReydenFallback(context.Background()) + require.NoError(t, err) + assert.False(t, thriftCalled, "explicit UseKernel must not attempt Thrift") + _, ok := be.(*fakeKernelBackend) + assert.True(t, ok, "explicit UseKernel should return the kernel backend, got %T", be) + }) +} + +func TestReydenNonReydenError(t *testing.T) { + t.Run("Non-KP001 Thrift error propagates unchanged, no kernel fallback", func(t *testing.T) { + defer warehouse_cache.ClearCache() + + tt := NewTestReydenFallback() + + genericErr := errors.New("some server error") + kernelCalled := false + conn := tt.makeConnector( + func(ctx context.Context, cfg *config.Config, client *http.Client) (backend.Backend, error) { + return &fakeThriftBackend{openSessionErr: genericErr}, nil + }, + func(ctx context.Context, cfg *config.Config) (backend.Backend, error) { + kernelCalled = true + return &fakeKernelBackend{}, nil + }, + ) + + _, _, err := conn.openSessionWithReydenFallback(context.Background()) + require.Error(t, err) + assert.Same(t, genericErr, err, "the original Thrift error should propagate unchanged") + assert.False(t, kernelCalled, "a non-Reyden error must not trigger the kernel fallback") + assert.False(t, warehouse_cache.IsKnownReyden(tt.host, tt.warehouseID), + "a non-Reyden error must not mark the warehouse") + }) +} + +func TestReydenDoubleFailureChaining(t *testing.T) { + t.Run("Kernel open failure after Thrift rejection chains both errors", func(t *testing.T) { + defer warehouse_cache.ClearCache() + + tt := NewTestReydenFallback() + + thriftErr := dbsqlerrint.NewReydenThriftUnsupportedError( + "Lakehouse/RT is not supported for Thrift protocol") + kernelErr := errors.New("kernel open failed") + conn := tt.makeConnector( + func(ctx context.Context, cfg *config.Config, client *http.Client) (backend.Backend, error) { + return &fakeThriftBackend{openSessionErr: thriftErr}, nil + }, + func(ctx context.Context, cfg *config.Config) (backend.Backend, error) { + return &fakeKernelBackend{openSessionErr: kernelErr}, nil + }, + ) + + _, _, err := conn.openSessionWithReydenFallback(context.Background()) + require.Error(t, err) + // Both the kernel failure (the actionable one) and the original Thrift + // rejection must be reachable in the returned error chain. + assert.ErrorIs(t, err, kernelErr, "chain should contain the kernel failure") + assert.ErrorIs(t, err, dbsqlerr.ErrReydenThriftUnsupported, + "chain should preserve the Thrift rejection marker") + }) +} + +func TestReydenDefaultBuildKernelNotCompiled(t *testing.T) { + t.Run("Default build's newKernelBackend returns 'not compiled' error", func(t *testing.T) { + // In the default pure-Go build, newKernelBackend returns an error indicating + // the kernel backend is not compiled in. This naturally exercises the + // error-chaining path for double-failure scenarios. + cfg := config.WithDefaults() + cfg.UseKernel = true + + // Attempt to create a kernel backend in the default build. + // (This would fail with ErrKernelNotCompiled outside of a + // databricks_kernel+CGO_ENABLED=1 build.) + be, err := newKernelBackend(context.Background(), cfg) + + // In the default build, this should fail. + assert.Nil(t, be) + assert.Error(t, err) + assert.True(t, errors.Is(err, dbsqlerr.ErrKernelNotCompiled), + "default build should not have kernel backend compiled in") + }) +} diff --git a/errors/errors.go b/errors/errors.go index c90f3b1a..b3d07fa7 100644 --- a/errors/errors.go +++ b/errors/errors.go @@ -91,6 +91,13 @@ var ErrRequiresKernelBackend error = errors.New("requires the SEA-via-kernel bac // retryable error instead of matching on message text. var ErrInvalidKernelConfig error = errors.New("invalid kernel backend configuration") +// value to be used with errors.Is() to determine that a Reyden / Real-Time +// warehouse rejected the legacy Thrift protocol (SQLSTATE KP001). The connection +// layer transparently recovers by re-opening the session on the kernel backend. +// Subclassing under a distinct sentinel lets a caller detect this specific +// recoverable condition programmatically instead of matching on error message. +var ErrReydenThriftUnsupported error = errors.New("Reyden warehouse rejects legacy Thrift protocol") + // Base interface for driver errors type DBError interface { // Descriptive message describing the error diff --git a/internal/client/client.go b/internal/client/client.go index 01d98010..b650d2e5 100644 --- a/internal/client/client.go +++ b/internal/client/client.go @@ -398,6 +398,12 @@ func CheckStatus(resp interface{}) error { if ok { status := rpcresp.GetStatus() if status.StatusCode == cli_service.TStatusCode_ERROR_STATUS { + // Detect Reyden warehouse rejection: SQLSTATE KP001 indicates the warehouse + // rejects the legacy Thrift protocol. Surface a distinct marker so the + // connection layer can transparently re-open on the kernel backend. + if status.GetSqlState() == "KP001" { + return dbsqlerrint.NewReydenThriftUnsupportedError(status.GetErrorMessage()) + } return errors.New(status.GetErrorMessage()) } if status.StatusCode == cli_service.TStatusCode_INVALID_HANDLE_STATUS { diff --git a/internal/errors/err.go b/internal/errors/err.go index f6745f41..b581ed55 100644 --- a/internal/errors/err.go +++ b/internal/errors/err.go @@ -311,3 +311,35 @@ func (e badConnectionError) Error() string { func NewBadConnectionError(err error) error { return badConnectionError{err: err} } + +// reydenThriftUnsupportedError signals that a Reyden / Real-Time warehouse +// rejected the legacy Thrift protocol (SQLSTATE KP001). The connection layer +// transparently recovers by re-opening the session on the kernel backend. +type reydenThriftUnsupportedError struct { + err error +} + +func (e reydenThriftUnsupportedError) Is(err error) bool { + return err == dbsqlerr.ErrReydenThriftUnsupported +} + +func (e reydenThriftUnsupportedError) Unwrap() error { + return e.err +} + +func (e reydenThriftUnsupportedError) Error() string { + if e.err != nil { + return e.err.Error() + } + return "Reyden warehouse rejects legacy Thrift protocol" +} + +// NewReydenThriftUnsupportedError creates a marker error for Reyden warehouse +// rejection of Thrift protocol (SQLSTATE KP001). +func NewReydenThriftUnsupportedError(msg string) error { + var err error + if msg != "" { + err = errors.New(msg) + } + return reydenThriftUnsupportedError{err: err} +} diff --git a/internal/warehouse_cache/cache.go b/internal/warehouse_cache/cache.go new file mode 100644 index 00000000..904fab63 --- /dev/null +++ b/internal/warehouse_cache/cache.go @@ -0,0 +1,138 @@ +// Package warehouse_cache maintains a process-wide cache of warehouses known to +// reject the legacy Thrift protocol. +// +// A Reyden / Real-Time SQL warehouse rejects a Thrift OpenSession — the SQL +// Gateway proxy stamps SQLSTATE KP001 on the rejection. When the driver +// auto-recovers by re-opening on the kernel backend, it records the warehouse +// here so later connections to the same warehouse skip the doomed Thrift attempt +// and open on the kernel directly. +// +// Keyed by (host, warehouse_id) — the host is part of the key so the same +// warehouse id observed on two different workspaces never collides. Entries +// expire after TTLSeconds so a warehouse later reconfigured to accept Thrift +// is eventually retried. +package warehouse_cache + +import ( + "regexp" + "strings" + "sync" + "time" +) + +const ( + // TTLSeconds controls how long a cached Reyden warehouse entry remains valid. + // Matches the ADBC driver's 6-hour horizon. + TTLSeconds = 6 * 60 * 60 +) + +var ( + // warehousePathRE matches warehouse and endpoint paths like + // /sql/1.0/warehouses/ or .../endpoints/; the id stops at the + // next /, ?, or & (e.g. a ?o= SPOG routing param). All-purpose-compute + // cluster paths carry no warehouse id and never match — they are never + // Reyden warehouses. + warehousePathRE = regexp.MustCompile(`(?:/|^)(?:warehouses|endpoints)/([^?&/]+)`) +) + +// ExtractWarehouseID returns the warehouse/endpoint id embedded in httpPath, or empty string. +func ExtractWarehouseID(httpPath string) string { + if httpPath == "" { + return "" + } + match := warehousePathRE.FindStringSubmatch(httpPath) + if match != nil && len(match) > 1 { + return match[1] + } + return "" +} + +// Cache is a thread-safe cache of warehouses known to reject Thrift. +type Cache struct { + mu sync.RWMutex + // (host_lowercased, warehouse_id) -> expiry deadline (monotonic time) + expiry map[[2]string]time.Time +} + +// NewCache returns a new thread-safe cache. +func NewCache() *Cache { + return &Cache{ + expiry: make(map[[2]string]time.Time), + } +} + +// makeKey returns the cache key for (host, warehouse_id), with host lowercased. +func makeKey(host, warehouseID string) [2]string { + return [2]string{strings.ToLower(host), warehouseID} +} + +// MarkReyden records that warehouseID on host rejects the Thrift protocol. +// Performs opportunistic sweep of expired entries. +func (c *Cache) MarkReyden(host, warehouseID string) { + if warehouseID == "" { + return + } + now := time.Now() + c.mu.Lock() + defer c.mu.Unlock() + + // Opportunistic sweep: mark_reyden only runs on an actual Thrift + // rejection (rare), so purging every expired entry here is near-free + // and bounds the cache to warehouses seen within the TTL window + // rather than every warehouse ever seen. + for key := range c.expiry { + if c.expiry[key].Before(now) { + delete(c.expiry, key) + } + } + + key := makeKey(host, warehouseID) + c.expiry[key] = now.Add(time.Duration(TTLSeconds) * time.Second) +} + +// IsKnownReyden returns whether warehouseID on host is known (unexpired) to reject Thrift. +func (c *Cache) IsKnownReyden(host, warehouseID string) bool { + if warehouseID == "" { + return false + } + now := time.Now() + c.mu.RLock() + defer c.mu.RUnlock() + + key := makeKey(host, warehouseID) + deadline, ok := c.expiry[key] + if !ok { + return false + } + if deadline.Before(now) { + // Entry is expired but not yet lazily evicted. + // Lazy eviction happens on the next MarkReyden call. + return false + } + return true +} + +// Clear resets the cache. Intended for tests. +func (c *Cache) Clear() { + c.mu.Lock() + defer c.mu.Unlock() + c.expiry = make(map[[2]string]time.Time) +} + +// Global singleton cache, multi-tenant safe via the host component of the key. +var globalCache = NewCache() + +// MarkReyden records that warehouseID on host rejects the Thrift protocol. +func MarkReyden(host, warehouseID string) { + globalCache.MarkReyden(host, warehouseID) +} + +// IsKnownReyden returns whether warehouseID on host is known to reject Thrift. +func IsKnownReyden(host, warehouseID string) bool { + return globalCache.IsKnownReyden(host, warehouseID) +} + +// ClearCache resets the global cache. Intended for tests. +func ClearCache() { + globalCache.Clear() +} diff --git a/internal/warehouse_cache/cache_test.go b/internal/warehouse_cache/cache_test.go new file mode 100644 index 00000000..9b5d9877 --- /dev/null +++ b/internal/warehouse_cache/cache_test.go @@ -0,0 +1,168 @@ +package warehouse_cache + +import ( + "testing" + "time" +) + +func TestExtractWarehouseID(t *testing.T) { + tests := []struct { + path string + expected string + }{ + {"/sql/1.0/warehouses/abc123", "abc123"}, + {"/sql/1.0/endpoints/def456", "def456"}, + {"/sql/1.0/warehouses/abc123?o=42", "abc123"}, + {"/sql/1.0/warehouses/abc123?o=42&other=val", "abc123"}, + {"sql/1.0/warehouses/wh", "wh"}, + {"/sql/protocolv1/o/1234567890/0101-cluster", ""}, + {"", ""}, + } + + for _, tt := range tests { + t.Run(tt.path, func(t *testing.T) { + result := ExtractWarehouseID(tt.path) + if result != tt.expected { + t.Errorf("ExtractWarehouseID(%q) = %q, want %q", tt.path, result, tt.expected) + } + }) + } +} + +func TestCacheMarkAndIsKnown(t *testing.T) { + cache := NewCache() + host := "host.example.com" + warehouseID := "wh-123" + + // Initially not known + if cache.IsKnownReyden(host, warehouseID) { + t.Errorf("warehouse should not be known initially") + } + + // After marking, should be known + cache.MarkReyden(host, warehouseID) + if !cache.IsKnownReyden(host, warehouseID) { + t.Errorf("warehouse should be known after marking") + } +} + +func TestCacheHostCaseInsensitive(t *testing.T) { + cache := NewCache() + warehouseID := "wh-123" + + cache.MarkReyden("Host.Example.COM", warehouseID) + + // Lookup with different case should still find it + if !cache.IsKnownReyden("host.example.com", warehouseID) { + t.Errorf("warehouse lookup should be case-insensitive") + } +} + +func TestCacheDistinctHostsDoNotCollide(t *testing.T) { + cache := NewCache() + warehouseID := "wh-123" + + cache.MarkReyden("host-a", warehouseID) + + // Same warehouse id on different host should not be treated as Reyden + if cache.IsKnownReyden("host-b", warehouseID) { + t.Errorf("different hosts should not collide") + } +} + +func TestCacheDistinctWarehousesDoNotCollide(t *testing.T) { + cache := NewCache() + host := "host.example.com" + + cache.MarkReyden(host, "wh-a") + + // Different warehouse on same host should not collide + if cache.IsKnownReyden(host, "wh-b") { + t.Errorf("different warehouses should not collide") + } +} + +func TestCacheExpiry(t *testing.T) { + cache := &Cache{expiry: make(map[[2]string]time.Time)} + host := "host.example.com" + warehouseID := "wh-123" + + // Mark with zero TTL (entry immediately expires) + now := time.Now() + key := makeKey(host, warehouseID) + cache.mu.Lock() + cache.expiry[key] = now.Add(-1 * time.Nanosecond) + cache.mu.Unlock() + + // Should be expired + if cache.IsKnownReyden(host, warehouseID) { + t.Errorf("warehouse should be expired") + } +} + +func TestCacheSweepsExpiredEntries(t *testing.T) { + cache := &Cache{expiry: make(map[[2]string]time.Time)} + host := "host.example.com" + + // Add an entry that's already expired + now := time.Now() + key1 := makeKey(host, "old-wh") + key2 := makeKey(host, "new-wh") + + cache.mu.Lock() + cache.expiry[key1] = now.Add(-1 * time.Second) + cache.mu.Unlock() + + if len(cache.expiry) != 1 { + t.Errorf("cache should have 1 entry") + } + + // Mark a new warehouse should trigger sweep + cache.MarkReyden(host, "new-wh") + + cache.mu.RLock() + defer cache.mu.RUnlock() + + // Old entry should be gone + if _, ok := cache.expiry[key1]; ok { + t.Errorf("expired entry should be swept") + } + // New entry should be present + if _, ok := cache.expiry[key2]; !ok { + t.Errorf("new entry should be present after sweep") + } +} + +func TestGlobalCache(t *testing.T) { + ClearCache() + + host := "example.com" + warehouseID := "wh-global" + + if IsKnownReyden(host, warehouseID) { + t.Errorf("warehouse should not be known initially") + } + + MarkReyden(host, warehouseID) + + if !IsKnownReyden(host, warehouseID) { + t.Errorf("warehouse should be known after marking") + } + + ClearCache() + + if IsKnownReyden(host, warehouseID) { + t.Errorf("warehouse should not be known after clear") + } +} + +func TestCacheEmptyWarehouseID(t *testing.T) { + cache := NewCache() + + // Should not crash or add empty entries + cache.MarkReyden("host", "") + + if cache.IsKnownReyden("host", "") { + t.Errorf("empty warehouse id should not be cached") + } +} From 054bb55268016ce775e71595fe06906e34ceac00 Mon Sep 17 00:00:00 2001 From: Rahul Singhal Date: Wed, 9 Sep 2026 22:07:36 +0000 Subject: [PATCH 2/6] fix: skip driver telemetry for a recovered kernel connection skipDriverTelemetry read cfg.UseKernel, which stays false on the Reyden auto-recovery path (the fallback opens the kernel without mutating cfg). A recovered-kernel connection therefore kept the Go driver's telemetry active, duplicating the kernel's own telemetry. Derive the skip decision from the backend that actually opened instead. Mirrors the analogous fix in the Python driver. Co-authored-by: Isaac Signed-off-by: Rahul Singhal --- connector.go | 12 +++++++++++- connector_reyden_test.go | 11 +++++++++++ 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/connector.go b/connector.go index 382af388..cc199865 100644 --- a/connector.go +++ b/connector.go @@ -49,6 +49,16 @@ func skipDriverTelemetry(cfg *config.Config) bool { return cfg.UseKernel } +// shouldSkipDriverTelemetry reports whether driver-side telemetry should be +// skipped for the active backend. The kernel backend owns telemetry, so the +// driver skips its own to avoid duplication. This is derived from the backend +// that actually opened (not the config) so a Reyden auto-recovery onto the +// kernel — which does not set cfg.UseKernel — is still attributed correctly. +func shouldSkipDriverTelemetry(be backend.Backend) bool { + _, isThrift := be.(*thrift.Backend) + return !isThrift +} + // federatedTokenAuthenticator preserves the base provider for the kernel. type federatedTokenAuthenticator struct { auth.Authenticator @@ -88,7 +98,7 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) { // Skip driver telemetry on the kernel path. The kernel owns query execution // below the driver backend, so keeping the Go telemetry interceptor active // would duplicate kernel telemetry for the same connection/statements. - skipTelemetry := skipDriverTelemetry(c.cfg) + skipTelemetry := shouldSkipDriverTelemetry(be) if skipTelemetry { log.Debug().Msg("telemetry skipped: kernel backend owns telemetry") } diff --git a/connector_reyden_test.go b/connector_reyden_test.go index ec2303b1..b392925a 100644 --- a/connector_reyden_test.go +++ b/connector_reyden_test.go @@ -8,6 +8,7 @@ import ( dbsqlerr "github.com/databricks/databricks-sql-go/errors" "github.com/databricks/databricks-sql-go/internal/backend" + "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" @@ -346,6 +347,16 @@ func TestReydenDoubleFailureChaining(t *testing.T) { }) } +func TestReydenSkipDriverTelemetryFollowsActualBackend(t *testing.T) { + // Driver telemetry must be skipped whenever the ACTIVE backend is the + // kernel — including after a Reyden recovery, where cfg.UseKernel stays + // false. The decision is derived from the backend, not the config. + assert.True(t, shouldSkipDriverTelemetry(&fakeKernelBackend{}), + "kernel backend owns telemetry; driver telemetry should be skipped") + assert.False(t, shouldSkipDriverTelemetry(&thrift.Backend{}), + "thrift backend: driver telemetry should stay active") +} + func TestReydenDefaultBuildKernelNotCompiled(t *testing.T) { t.Run("Default build's newKernelBackend returns 'not compiled' error", func(t *testing.T) { // In the default pure-Go build, newKernelBackend returns an error indicating From 2d01313202e954e6efdd54f1568cd2c2807c4f93 Mon Sep 17 00:00:00 2001 From: Rahul Singhal Date: Wed, 9 Sep 2026 22:24:14 +0000 Subject: [PATCH 3/6] refactor: scope KP001 Reyden detection to OpenSession only CheckStatus is the shared status checker for every Thrift RPC, so mapping KP001 to the recoverable marker there gave it a wider blast radius than the recovery logic (which only wraps session open): a stray KP001 on any other RPC would have surfaced as ErrReydenThriftUnsupported with no handler. Keep CheckStatus generic and add an OpenSession-scoped CheckOpenSessionStatus that does the KP001 mapping; only the OpenSession wrapper uses it. Every other RPC now surfaces a KP001 as a plain error, unchanged from before. Co-authored-by: Isaac Signed-off-by: Rahul Singhal --- connector_reyden_test.go | 32 +++++++++++++++++++++----------- internal/client/client.go | 26 +++++++++++++++++++------- 2 files changed, 40 insertions(+), 18 deletions(-) diff --git a/connector_reyden_test.go b/connector_reyden_test.go index b392925a..735f401c 100644 --- a/connector_reyden_test.go +++ b/connector_reyden_test.go @@ -18,28 +18,39 @@ import ( "github.com/stretchr/testify/require" ) -// TestCheckStatusReydenDetection tests that CheckStatus detects SQLSTATE KP001 -// and returns a distinct error vs. other ERROR_STATUS codes. +// TestCheckStatusReydenDetection tests that KP001 maps to the Reyden marker +// ONLY on the OpenSession-scoped check, and stays a generic error everywhere +// else (the shared CheckStatus). func TestCheckStatusReydenDetection(t *testing.T) { - t.Run("KP001 on ERROR_STATUS returns ReydenThriftUnsupported marker", func(t *testing.T) { + reydenStatus := func() *cli_service.TOpenSessionResp { sqlState := "KP001" errMsg := "Lakehouse/RT is not supported for Thrift protocol" - reydenResp := &cli_service.TOpenSessionResp{ + return &cli_service.TOpenSessionResp{ Status: &cli_service.TStatus{ StatusCode: cli_service.TStatusCode_ERROR_STATUS, SqlState: &sqlState, ErrorMessage: &errMsg, }, } + } - err := client.CheckStatus(reydenResp) + t.Run("KP001 via CheckOpenSessionStatus returns the Reyden marker", func(t *testing.T) { + err := client.CheckOpenSessionStatus(reydenStatus()) require.Error(t, err) - // Verify it's the Reyden marker. assert.True(t, errors.Is(err, dbsqlerr.ErrReydenThriftUnsupported), - "error should satisfy errors.Is for ErrReydenThriftUnsupported") + "OpenSession KP001 should satisfy errors.Is for ErrReydenThriftUnsupported") + }) + + t.Run("KP001 via the shared CheckStatus is a generic error, NOT the marker", func(t *testing.T) { + // A KP001 on any non-OpenSession RPC must not become the marker — the + // recovery layer only handles it at session open. + err := client.CheckStatus(reydenStatus()) + require.Error(t, err) + assert.False(t, errors.Is(err, dbsqlerr.ErrReydenThriftUnsupported), + "a KP001 outside OpenSession must stay a generic error") }) - t.Run("non-KP001 ERROR_STATUS returns generic error, NOT Reyden marker", func(t *testing.T) { + t.Run("non-KP001 ERROR_STATUS is never the marker, even at OpenSession", func(t *testing.T) { sqlState := "42000" // Syntax error errMsg := "a syntax error" syntaxResp := &cli_service.TOpenSessionResp{ @@ -50,9 +61,8 @@ func TestCheckStatusReydenDetection(t *testing.T) { }, } - err := client.CheckStatus(syntaxResp) + err := client.CheckOpenSessionStatus(syntaxResp) require.Error(t, err) - // Verify it's NOT the Reyden marker. assert.False(t, errors.Is(err, dbsqlerr.ErrReydenThriftUnsupported), "generic error should NOT satisfy errors.Is for ErrReydenThriftUnsupported") }) @@ -64,7 +74,7 @@ func TestCheckStatusReydenDetection(t *testing.T) { }, } - err := client.CheckStatus(successResp) + err := client.CheckOpenSessionStatus(successResp) assert.NoError(t, err) }) } diff --git a/internal/client/client.go b/internal/client/client.go index b650d2e5..fec632cd 100644 --- a/internal/client/client.go +++ b/internal/client/client.go @@ -127,7 +127,8 @@ func (tsc *ThriftServiceClient) OpenSession(ctx context.Context, req *cli_servic recordResult(ctx, resp) - return resp, CheckStatus(resp) + // OpenSession is the only RPC that opts into Reyden KP001 detection. + return resp, CheckOpenSessionStatus(resp) } // CloseSession is a wrapper around the thrift operation CloseSession @@ -398,12 +399,6 @@ func CheckStatus(resp interface{}) error { if ok { status := rpcresp.GetStatus() if status.StatusCode == cli_service.TStatusCode_ERROR_STATUS { - // Detect Reyden warehouse rejection: SQLSTATE KP001 indicates the warehouse - // rejects the legacy Thrift protocol. Surface a distinct marker so the - // connection layer can transparently re-open on the kernel backend. - if status.GetSqlState() == "KP001" { - return dbsqlerrint.NewReydenThriftUnsupportedError(status.GetErrorMessage()) - } return errors.New(status.GetErrorMessage()) } if status.StatusCode == cli_service.TStatusCode_INVALID_HANDLE_STATUS { @@ -416,6 +411,23 @@ func CheckStatus(resp interface{}) error { return errors.New("thrift: invalid response") } +// CheckOpenSessionStatus is CheckStatus plus Reyden KP001 detection, used only +// for the OpenSession response. A Reyden / Real-Time warehouse rejects the +// legacy Thrift protocol with SQLSTATE KP001, and the connection layer's +// recovery only wraps session open — so the marker is scoped to this call +// rather than the shared CheckStatus, keeping a stray KP001 on any other RPC a +// plain error. +func CheckOpenSessionStatus(resp interface{}) error { + if rpcresp, ok := resp.(ThriftResponse); ok { + status := rpcresp.GetStatus() + if status.StatusCode == cli_service.TStatusCode_ERROR_STATUS && + status.GetSqlState() == "KP001" { + return dbsqlerrint.NewReydenThriftUnsupportedError(status.GetErrorMessage()) + } + } + return CheckStatus(resp) +} + // SprintGuid is a convenience function to format a byte array into GUID. func SprintGuid(bts []byte) string { if len(bts) == 16 { From 9f74e0f2daf93aeb172252ad79987b3d8d9332c5 Mon Sep 17 00:00:00 2001 From: Rahul Singhal Date: Thu, 10 Sep 2026 00:46:43 +0000 Subject: [PATCH 4/6] Address review: guardrail before pre-check, drop stale helper, pin wrapped-marker tests - connector: hoist the WithKernel*-without-WithUseKernel guardrail above the cache pre-check so the misconfiguration is rejected deterministically, regardless of process-global cache state (previously a warm cache let the pre-check open the kernel and silently bypass the guardrail). - connector: delete the now-unused skipDriverTelemetry(cfg); the skip decision is derived from the active backend via shouldSkipDriverTelemetry(be). Remove its redundant test (shouldSkipDriverTelemetry is already covered). - connector_reyden_test: inject the Reyden marker WRAPPED in NewRequestError, as thrift.Backend.OpenSession does in production, so the errors.Is unwrap chain the recovery relies on is pinned; add a test that the guardrail fires on a warm cache. Co-authored-by: Isaac Signed-off-by: Rahul Singhal --- connector.go | 30 +++++++++++++-------------- connector_kernel_u2m_test.go | 22 -------------------- connector_reyden_test.go | 39 +++++++++++++++++++++++++++++++----- 3 files changed, 48 insertions(+), 43 deletions(-) delete mode 100644 connector_kernel_u2m_test.go diff --git a/connector.go b/connector.go index cc199865..1fbf338d 100644 --- a/connector.go +++ b/connector.go @@ -45,10 +45,6 @@ type connector struct { thriftBackendFactory thriftBackendFactory // Seam for testing; nil in production } -func skipDriverTelemetry(cfg *config.Config) bool { - return cfg.UseKernel -} - // shouldSkipDriverTelemetry reports whether driver-side telemetry should be // skipped for the active backend. The kernel backend owns telemetry, so the // driver skips its own to avoid duplication. This is derived from the backend @@ -829,6 +825,18 @@ func WithTokenCache(enabled bool) ConnOption { // UseKernel nor other backend-selecting options). On success or failure, it // returns the backend, session latency, and error. func (c *connector) openSessionWithReydenFallback(ctx context.Context) (backend.Backend, int64, error) { + // Guardrail, checked up front — before the cache pre-check — so the outcome does not depend + // on process-global cache state. The experimental WithKernel* options have no Thrift-path + // equivalent, so a caller who sets one (a trusted-CA bundle, a hostname-verify skip, a proxy, + // a retry budget, or a CloudFetch chunk cap) and forgets WithUseKernel is rejected loudly + // rather than connecting as if it were never set. Every WithKernel* option allocates + // KernelExperimental, so this one gate covers them all; the message names the family rather + // than a stale subset that drifts as options are added. + if !c.cfg.UseKernel && c.cfg.KernelExperimental != nil { + return nil, 0, fmt.Errorf("databricks: a WithKernel* option %w; "+ + "add WithUseKernel(true) or remove it", dbsqlerr.ErrRequiresKernelBackend) + } + // Extract warehouse ID from HTTPPath for cache lookups. warehouseID := warehouse_cache.ExtractWarehouseID(c.cfg.HTTPPath) @@ -860,18 +868,8 @@ func (c *connector) openSessionWithReydenFallback(ctx context.Context) (backend. if c.cfg.UseKernel { be, err = c.getKernelBackend(ctx) } else { - // The experimental WithKernel* options have no Thrift-path equivalent — reject - // them loudly rather than silently ignore, so a caller who sets one (a - // trusted-CA bundle, a hostname-verify skip, a proxy, a retry budget, or a - // CloudFetch chunk cap) and forgets WithUseKernel learns the option had no - // effect instead of connecting as if it were never set. Every WithKernel* - // option allocates KernelExperimental, so this one gate covers them all; the - // message names the family rather than a stale subset that drifts as options - // are added. - if c.cfg.KernelExperimental != nil { - return nil, 0, fmt.Errorf("databricks: a WithKernel* option %w; "+ - "add WithUseKernel(true) or remove it", dbsqlerr.ErrRequiresKernelBackend) - } + // WithKernel*-without-WithUseKernel was already rejected up front; the default + // path is Thrift. be, err = c.getThriftBackend(ctx) } if err != nil { diff --git a/connector_kernel_u2m_test.go b/connector_kernel_u2m_test.go deleted file mode 100644 index 64713e3f..00000000 --- a/connector_kernel_u2m_test.go +++ /dev/null @@ -1,22 +0,0 @@ -package dbsql - -import ( - "testing" - - "github.com/databricks/databricks-sql-go/internal/config" - - "github.com/stretchr/testify/assert" -) - -func TestKernelSkipsDriverTelemetry(t *testing.T) { - assert.True( - t, - skipDriverTelemetry(&config.Config{UserConfig: config.UserConfig{UseKernel: true}}), - "kernel connections skip driver telemetry", - ) - assert.False( - t, - skipDriverTelemetry(&config.Config{UserConfig: config.UserConfig{UseKernel: false}}), - "thrift connections keep driver telemetry eligible", - ) -} diff --git a/connector_reyden_test.go b/connector_reyden_test.go index 735f401c..d052928f 100644 --- a/connector_reyden_test.go +++ b/connector_reyden_test.go @@ -176,10 +176,15 @@ func TestReydenReactiveRecovery(t *testing.T) { tt := NewTestReydenFallback() - // Fake Thrift backend that rejects OpenSession with the Reyden marker. + // Fake Thrift backend that rejects OpenSession with the Reyden marker, wrapped exactly + // as production does: thrift.Backend.OpenSession wraps every failure in NewRequestError, + // so this pins the errors.Is unwrap chain the recovery relies on. fakeThrift := &fakeThriftBackend{ - openSessionErr: dbsqlerrint.NewReydenThriftUnsupportedError( - "Lakehouse/RT is not supported for Thrift protocol"), + openSessionErr: dbsqlerrint.NewRequestError( + context.Background(), + "error connecting", + dbsqlerrint.NewReydenThriftUnsupportedError( + "Lakehouse/RT is not supported for Thrift protocol")), } kernelBackend := &fakeKernelBackend{} @@ -300,6 +305,25 @@ func TestReydenGuardrail(t *testing.T) { _, ok := be.(*fakeKernelBackend) assert.True(t, ok, "explicit UseKernel should return the kernel backend, got %T", be) }) + + t.Run("WithKernel* without WithUseKernel is rejected up front, even on a warm cache", func(t *testing.T) { + // The guardrail runs before the cache pre-check, so the same misconfiguration errors + // deterministically whether or not a sibling connection already cached the warehouse. + defer warehouse_cache.ClearCache() + + tt := NewTestReydenFallback() + // Prime the cache so the pre-check would otherwise fire and open the kernel directly. + warehouse_cache.MarkReyden(tt.host, tt.warehouseID) + + // Factories are nil: the guardrail must reject before any backend is built. + conn := tt.makeConnector(nil, nil) + conn.cfg.KernelExperimental = &config.KernelExperimentalConfig{} + + _, _, err := conn.openSessionWithReydenFallback(context.Background()) + require.Error(t, err) + assert.ErrorIs(t, err, dbsqlerr.ErrRequiresKernelBackend, + "a WithKernel* option without WithUseKernel must be rejected regardless of cache state") + }) } func TestReydenNonReydenError(t *testing.T) { @@ -335,8 +359,13 @@ func TestReydenDoubleFailureChaining(t *testing.T) { tt := NewTestReydenFallback() - thriftErr := dbsqlerrint.NewReydenThriftUnsupportedError( - "Lakehouse/RT is not supported for Thrift protocol") + // Wrap the marker as production does (NewRequestError), so the double-failure chain is + // exercised against the real error shape the connector sees. + thriftErr := dbsqlerrint.NewRequestError( + context.Background(), + "error connecting", + dbsqlerrint.NewReydenThriftUnsupportedError( + "Lakehouse/RT is not supported for Thrift protocol")) kernelErr := errors.New("kernel open failed") conn := tt.makeConnector( func(ctx context.Context, cfg *config.Config, client *http.Client) (backend.Backend, error) { From f75e10d424cc99fb354e6a9246156020b95d2978 Mon Sep 17 00:00:00 2001 From: Rahul Singhal Date: Thu, 10 Sep 2026 01:04:02 +0000 Subject: [PATCH 5/6] Fix CI: lint issues in warehouse_cache, gate default-build kernel test - warehouse_cache/cache.go: drop the redundant `match != nil` (staticcheck S1009: len(nil) is 0) and fix struct field gofmt alignment. Both failed golangci-lint. - Move TestReydenDefaultBuildKernelNotCompiled into a //go:build !databricks_kernel file. It asserts newKernelBackend reports "not compiled", which only holds in the default build; under -tags databricks_kernel the real backend is linked in, so the test was failing the "Test (kernel backend)" CI job. Co-authored-by: Isaac Signed-off-by: Rahul Singhal --- connector_reyden_defaultbuild_test.go | 31 +++++++++++++++++++++++++++ connector_reyden_test.go | 21 ------------------ internal/warehouse_cache/cache.go | 4 ++-- 3 files changed, 33 insertions(+), 23 deletions(-) create mode 100644 connector_reyden_defaultbuild_test.go diff --git a/connector_reyden_defaultbuild_test.go b/connector_reyden_defaultbuild_test.go new file mode 100644 index 00000000..ef9deedb --- /dev/null +++ b/connector_reyden_defaultbuild_test.go @@ -0,0 +1,31 @@ +//go:build !databricks_kernel + +package dbsql + +import ( + "context" + "errors" + "testing" + + dbsqlerr "github.com/databricks/databricks-sql-go/errors" + "github.com/databricks/databricks-sql-go/internal/config" + "github.com/stretchr/testify/assert" +) + +// TestReydenDefaultBuildKernelNotCompiled asserts that in the default pure-Go build +// newKernelBackend reports the kernel is not compiled in. It is gated to the default +// build because under -tags databricks_kernel newKernelBackend is the real (compiled) +// implementation and does not return this error. +func TestReydenDefaultBuildKernelNotCompiled(t *testing.T) { + t.Run("Default build's newKernelBackend returns 'not compiled' error", func(t *testing.T) { + cfg := config.WithDefaults() + cfg.UseKernel = true + + be, err := newKernelBackend(context.Background(), cfg) + + assert.Nil(t, be) + assert.Error(t, err) + assert.True(t, errors.Is(err, dbsqlerr.ErrKernelNotCompiled), + "default build should not have kernel backend compiled in") + }) +} diff --git a/connector_reyden_test.go b/connector_reyden_test.go index d052928f..618f7a05 100644 --- a/connector_reyden_test.go +++ b/connector_reyden_test.go @@ -395,24 +395,3 @@ func TestReydenSkipDriverTelemetryFollowsActualBackend(t *testing.T) { assert.False(t, shouldSkipDriverTelemetry(&thrift.Backend{}), "thrift backend: driver telemetry should stay active") } - -func TestReydenDefaultBuildKernelNotCompiled(t *testing.T) { - t.Run("Default build's newKernelBackend returns 'not compiled' error", func(t *testing.T) { - // In the default pure-Go build, newKernelBackend returns an error indicating - // the kernel backend is not compiled in. This naturally exercises the - // error-chaining path for double-failure scenarios. - cfg := config.WithDefaults() - cfg.UseKernel = true - - // Attempt to create a kernel backend in the default build. - // (This would fail with ErrKernelNotCompiled outside of a - // databricks_kernel+CGO_ENABLED=1 build.) - be, err := newKernelBackend(context.Background(), cfg) - - // In the default build, this should fail. - assert.Nil(t, be) - assert.Error(t, err) - assert.True(t, errors.Is(err, dbsqlerr.ErrKernelNotCompiled), - "default build should not have kernel backend compiled in") - }) -} diff --git a/internal/warehouse_cache/cache.go b/internal/warehouse_cache/cache.go index 904fab63..eaa7c731 100644 --- a/internal/warehouse_cache/cache.go +++ b/internal/warehouse_cache/cache.go @@ -41,7 +41,7 @@ func ExtractWarehouseID(httpPath string) string { return "" } match := warehousePathRE.FindStringSubmatch(httpPath) - if match != nil && len(match) > 1 { + if len(match) > 1 { return match[1] } return "" @@ -49,7 +49,7 @@ func ExtractWarehouseID(httpPath string) string { // Cache is a thread-safe cache of warehouses known to reject Thrift. type Cache struct { - mu sync.RWMutex + mu sync.RWMutex // (host_lowercased, warehouse_id) -> expiry deadline (monotonic time) expiry map[[2]string]time.Time } From 8bfe13ededf3226158b17cdf6dc8960029ea7edb Mon Sep 17 00:00:00 2001 From: Rahul Singhal Date: Thu, 10 Sep 2026 01:23:39 +0000 Subject: [PATCH 6/6] Address review: wrap pre-check kernel errors, tidy cache sweep - connector: on the cache pre-check path, wrap a kernel backend-create/OpenSession failure with context noting the warehouse was cached as Reyden and Thrift was skipped. Uses %w so the underlying error (including a default-build ErrKernelNotCompiled) stays reachable via errors.Is; adds a test asserting both. - warehouse_cache: range the sweep over (key, deadline) instead of re-indexing the map per key. Co-authored-by: Isaac Signed-off-by: Rahul Singhal --- connector.go | 12 ++++++++++-- connector_reyden_test.go | 22 ++++++++++++++++++++++ internal/warehouse_cache/cache.go | 4 ++-- 3 files changed, 34 insertions(+), 4 deletions(-) diff --git a/connector.go b/connector.go index 1fbf338d..e703a4fb 100644 --- a/connector.go +++ b/connector.go @@ -846,14 +846,22 @@ func (c *connector) openSessionWithReydenFallback(ctx context.Context) (backend. logger.Debug().Msgf( "warehouse %s on %s is known to require kernel backend; skipping Thrift", warehouseID, c.cfg.Host) + // Wrap failures with context: the pre-check trusted a cached "Reyden" marker + // and deliberately skipped Thrift, so a bare kernel error (including a default-build + // ErrKernelNotCompiled from a sibling connection's marking) would otherwise hide why + // Thrift was never attempted. %w keeps the underlying error for errors.Is. var be backend.Backend var err error if be, err = c.getKernelBackend(ctx); err != nil { - return nil, 0, err + return nil, 0, fmt.Errorf( + "databricks: warehouse %s is cached as Reyden so Thrift was skipped, but the "+ + "kernel backend could not be created: %w", warehouseID, err) } sessionStart := time.Now() if err := be.OpenSession(ctx); err != nil { - return nil, 0, err + return nil, 0, fmt.Errorf( + "databricks: warehouse %s is cached as Reyden so Thrift was skipped, but the "+ + "kernel OpenSession failed: %w", warehouseID, err) } return be, time.Since(sessionStart).Milliseconds(), nil } diff --git a/connector_reyden_test.go b/connector_reyden_test.go index 618f7a05..8f2aec6f 100644 --- a/connector_reyden_test.go +++ b/connector_reyden_test.go @@ -242,6 +242,28 @@ func TestReydenPreCheck(t *testing.T) { assert.False(t, thriftCalled, "pre-check must skip the Thrift OpenSession round-trip") assert.GreaterOrEqual(t, latency, int64(0), "latency should be non-negative") }) + + t.Run("Kernel failure on the pre-check path is wrapped with context but preserves the chain", func(t *testing.T) { + defer warehouse_cache.ClearCache() + + tt := NewTestReydenFallback() + warehouse_cache.MarkReyden(tt.host, tt.warehouseID) + + kernelErr := errors.New("kernel open failed") + conn := tt.makeConnector( + nil, // Thrift must not be attempted on the pre-check path. + func(ctx context.Context, cfg *config.Config) (backend.Backend, error) { + return &fakeKernelBackend{openSessionErr: kernelErr}, nil + }, + ) + + _, _, err := conn.openSessionWithReydenFallback(context.Background()) + require.Error(t, err) + // The underlying kernel error stays reachable for errors.Is, and the message explains + // that Thrift was skipped due to the cache. + assert.ErrorIs(t, err, kernelErr, "the underlying kernel error must be preserved") + assert.Contains(t, err.Error(), "cached as Reyden") + }) } func TestReydenCacheMarking(t *testing.T) { diff --git a/internal/warehouse_cache/cache.go b/internal/warehouse_cache/cache.go index eaa7c731..038b7ce9 100644 --- a/internal/warehouse_cache/cache.go +++ b/internal/warehouse_cache/cache.go @@ -80,8 +80,8 @@ func (c *Cache) MarkReyden(host, warehouseID string) { // rejection (rare), so purging every expired entry here is near-free // and bounds the cache to warehouses seen within the TTL window // rather than every warehouse ever seen. - for key := range c.expiry { - if c.expiry[key].Before(now) { + for key, deadline := range c.expiry { + if deadline.Before(now) { delete(c.expiry, key) } }