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..e703a4fb 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,17 +23,36 @@ 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 { - 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. @@ -46,40 +66,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, @@ -100,7 +94,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") } @@ -139,6 +133,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 +814,117 @@ 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) { + // 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) + + // 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) + // 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, 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, 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 + } + + // 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 { + // WithKernel*-without-WithUseKernel was already rejected up front; the default + // path is Thrift. + 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_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_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 new file mode 100644 index 00000000..8f2aec6f --- /dev/null +++ b/connector_reyden_test.go @@ -0,0 +1,419 @@ +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/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" + 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 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) { + reydenStatus := func() *cli_service.TOpenSessionResp { + sqlState := "KP001" + errMsg := "Lakehouse/RT is not supported for Thrift protocol" + return &cli_service.TOpenSessionResp{ + Status: &cli_service.TStatus{ + StatusCode: cli_service.TStatusCode_ERROR_STATUS, + SqlState: &sqlState, + ErrorMessage: &errMsg, + }, + } + } + + t.Run("KP001 via CheckOpenSessionStatus returns the Reyden marker", func(t *testing.T) { + err := client.CheckOpenSessionStatus(reydenStatus()) + require.Error(t, err) + assert.True(t, errors.Is(err, dbsqlerr.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 is never the marker, even at OpenSession", 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.CheckOpenSessionStatus(syntaxResp) + require.Error(t, err) + 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.CheckOpenSessionStatus(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, 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.NewRequestError( + context.Background(), + "error connecting", + 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") + }) + + 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) { + 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) + }) + + 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) { + 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() + + // 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) { + 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 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") +} 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..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 @@ -410,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 { 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..038b7ce9 --- /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 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, deadline := range c.expiry { + if deadline.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") + } +}