feat: auto-recover Reyden Thrift connections onto the kernel backend

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).

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

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