From bceb2a140c500bd4c182cbe588d5920d57624233 Mon Sep 17 00:00:00 2001 From: Jay Xiao Date: Fri, 21 Aug 2026 21:36:29 +0000 Subject: [PATCH] Forward telemetry config to kernel backend --- CONNECTION_PARAMETERS.md | 22 ++++++----- README.md | 15 ++++---- connector.go | 33 ++++------------ connector_kernel_u2m_test.go | 34 ++++++----------- doc.go | 9 ++--- internal/backend/kernel/config.go | 36 ++++++++++++++++++ kernel_config.go | 38 +++++++++++++++++++ kernel_config_test.go | 62 +++++++++++++++++++++++++++++-- telemetry/system_info.go | 8 ++++ 9 files changed, 184 insertions(+), 73 deletions(-) diff --git a/CONNECTION_PARAMETERS.md b/CONNECTION_PARAMETERS.md index 2b13f8e4..d788b029 100644 --- a/CONNECTION_PARAMETERS.md +++ b/CONNECTION_PARAMETERS.md @@ -118,19 +118,21 @@ variables. ## Telemetry -Telemetry applies to **both** backends. When `enableTelemetry` is left unset (the -default), a server-side feature flag decides whether telemetry is active; setting it -explicitly overrides the flag. (Exception: on the kernel backend with OAuth **U2M**, -telemetry is skipped entirely to avoid a second interactive browser flow at connect.) +Go wrapper telemetry applies to the default Thrift backend. When `enableTelemetry` is +left unset (the default), a server-side feature flag decides whether wrapper telemetry is +active; setting it explicitly overrides the flag. On the kernel backend, the Go wrapper +skips its telemetry interceptor so it does not duplicate kernel-owned telemetry for the +same connection and statements, and forwards the kernel-owned telemetry knobs into the +kernel config. | DSN parameter | Thrift | Kernel | Default | Notes | |---|:---:|:---:|---|---| -| `enableTelemetry` | ✅ | ✅ | unset (server flag decides) | Force telemetry on/off, overriding the server feature flag. | -| `telemetry_batch_size` | ✅ | ✅ | `200` | Events per batch. | -| `telemetry_flush_interval` | ✅ | ✅ | `30s` | Flush interval. | +| `enableTelemetry` | ✅ | ✅ | unset (server flag decides wrapper telemetry; kernel telemetry defaults on) | Force Go wrapper telemetry on/off on the Thrift path, overriding the server feature flag. On the kernel path, forwarded to kernel-owned telemetry; unset forwards enabled. | +| `telemetry_batch_size` | ✅ | ✅ | `200` wrapper default; kernel default when unset | Events per batch. Forwarded to the kernel only when explicitly set. | +| `telemetry_flush_interval` | ✅ | ⚠️ | `30s` | Flush interval for Go wrapper telemetry. Parsed but not forwarded on the kernel path because the kernel owns its flush policy. | | `telemetry_retry_count` | ⚠️ | ⚠️ | — | **Deprecated and ignored** (retries are owned by the HTTP client + circuit breaker); logs a one-time warning. | | `telemetry_retry_delay` | ⚠️ | ⚠️ | — | **Deprecated and ignored** (see above). | -The kernel path additionally emits a connection-config telemetry event at connect (mode, -auth mechanism/flow, proxy, arrow, query tags, metric-view); the Thrift path's telemetry -is unchanged. See [`telemetry/DESIGN.md`](./telemetry/DESIGN.md). +The Go wrapper telemetry interceptor is skipped on the kernel path so it does not +duplicate kernel-owned telemetry for the same connection and statements. See +[`telemetry/DESIGN.md`](./telemetry/DESIGN.md). diff --git a/README.md b/README.md index 702cbc18..06bc8cd4 100644 --- a/README.md +++ b/README.md @@ -352,13 +352,14 @@ backends — the driver exposes no `GetCatalogs`/`GetSchemas`/`GetTables`/`GetCo ## Telemetry -The driver includes optional telemetry to help improve performance and reliability; it -applies to both backends. When `enableTelemetry` is left unset (the default), a -**server-side feature flag** decides whether telemetry is active — so it may be enabled -without an explicit opt-in. Setting `enableTelemetry` explicitly overrides the flag. -(One exception: on the kernel backend with OAuth **U2M**, telemetry is skipped entirely -— regardless of `enableTelemetry` — to avoid a second interactive browser flow at -connect.) +The driver includes optional telemetry to help improve performance and reliability. +Go wrapper telemetry applies to the default Thrift backend. When `enableTelemetry` is +left unset (the default), a **server-side feature flag** decides whether wrapper +telemetry is active — so it may be enabled without an explicit opt-in. Setting +`enableTelemetry` explicitly overrides the flag. On the kernel backend, the Go wrapper +skips its telemetry interceptor entirely so it does not duplicate kernel-owned +telemetry; `enableTelemetry` and `telemetry_batch_size` are forwarded into the kernel +telemetry config instead. ``` # force on (regardless of the server flag): diff --git a/connector.go b/connector.go index 040c3271..d55f67ab 100644 --- a/connector.go +++ b/connector.go @@ -31,11 +31,8 @@ type connector struct { client *http.Client } -// interactiveU2MAuthenticator is satisfied only by the browser-based U2M -// authenticator (U2MClientID is unique to it; PAT/M2M lack it). Matches the -// structural check the kernel backend uses to detect U2M. -type interactiveU2MAuthenticator interface { - U2MClientID() string +func skipDriverTelemetry(cfg *config.Config) bool { + return cfg.UseKernel } // Connect returns a connection to the Databricks database from a connection pool. @@ -93,17 +90,12 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) { telemetryClient = withSpogHeaders(c.client, spogHeaders) } - // Skip telemetry on the kernel U2M path: the kernel owns the interactive browser - // flow, so the telemetry/feature-flag call through the interactive authenticator - // would launch a second, redundant browser (and can block connect on its - // callback). Telemetry is best-effort, so it's dropped here rather than made to - // prompt. Unauthenticated telemetry (Python/Node parity) is tracked in PECOBLR-3839. - skipTelemetry := false - if c.cfg.UseKernel { - if _, isU2M := c.cfg.Authenticator.(interactiveU2MAuthenticator); isU2M { - skipTelemetry = true - log.Debug().Msg("telemetry skipped: kernel U2M owns the interactive auth flow") - } + // 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) + if skipTelemetry { + log.Debug().Msg("telemetry skipped: kernel backend owns telemetry") } // Initialize telemetry: client config overlay decides; if unset, feature flags decide @@ -121,15 +113,6 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) { if conn.telemetry != nil { log.Debug().Msg("telemetry initialized for connection") conn.telemetry.RecordOperation(ctx, conn.id, "", telemetry.OperationTypeCreateSession, sessionLatencyMs, nil) - // Connection-configuration telemetry on the kernel path only, so the - // default (Thrift) path's emitted telemetry stays byte-identical (the - // Thrift path has never populated DriverConnectionParameters). Emits mode / - // auth mech+flow / proxy / arrow / query-tags / metric-view for the - // just-opened session. Gated on the kernel backend, not just WithUseKernel, - // so it never fires when the kernel wasn't actually selected. - if _, ok := be.(*thrift.Backend); !ok { - conn.telemetry.RecordConnectionConfig(ctx, conn.id, kernelConnectionTelemetry(c.cfg)) - } } // ServerProtocolVersion is Thrift-specific (not on the neutral backend diff --git a/connector_kernel_u2m_test.go b/connector_kernel_u2m_test.go index abe11709..64713e3f 100644 --- a/connector_kernel_u2m_test.go +++ b/connector_kernel_u2m_test.go @@ -3,30 +3,20 @@ package dbsql import ( "testing" - "github.com/databricks/databricks-sql-go/auth/oauth/m2m" - "github.com/databricks/databricks-sql-go/auth/pat" + "github.com/databricks/databricks-sql-go/internal/config" "github.com/stretchr/testify/assert" ) -// u2mShaped stands in for the browser-based U2M authenticator, which exposes -// U2MClientID() (its real constructor does live OIDC discovery, unusable in a unit -// test). The kernel-U2M telemetry guard keys off exactly this method. -type u2mShaped struct{} - -func (u2mShaped) U2MClientID() string { return "databricks-sql-connector" } - -// TestInteractiveU2MAuthenticatorDetection pins the structural check the kernel-U2M -// telemetry guard relies on: only a U2M-shaped authenticator satisfies it, so -// PAT/M2M keep authenticated telemetry while U2M does not trigger a 2nd browser. -func TestInteractiveU2MAuthenticatorDetection(t *testing.T) { - _, isU2M := interface{}(u2mShaped{}).(interactiveU2MAuthenticator) - assert.True(t, isU2M, "a U2M-shaped authenticator must satisfy interactiveU2MAuthenticator") - - _, patIsU2M := interface{}(&pat.PATAuth{AccessToken: "dapi-x"}).(interactiveU2MAuthenticator) - assert.False(t, patIsU2M, "PAT must NOT satisfy interactiveU2MAuthenticator") - - m2mAuth := m2m.NewAuthenticator("cid", "secret", "dbc-1234.cloud.databricks.com") - _, m2mIsU2M := m2mAuth.(interactiveU2MAuthenticator) - assert.False(t, m2mIsU2M, "M2M must NOT satisfy interactiveU2MAuthenticator") +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/doc.go b/doc.go index 266deff0..cc934cb0 100644 --- a/doc.go +++ b/doc.go @@ -266,11 +266,10 @@ Setting any of these without WithUseKernel fails Connect with an error wrapping sentinel ErrRequiresKernelBackend, detectable with errors.Is. Features above the backend seam are inherited unchanged: the database/sql connection -pool, per-connection telemetry (CREATE_SESSION / EXECUTE_STATEMENT / DELETE_SESSION), -and the telemetry circuit breaker. The kernel path additionally emits a -connection-configuration telemetry event at connect (mode=SEA, auth mechanism/flow, -proxy usage, arrow, query tags, metric-view metadata); this is kernel-only, so the -default (Thrift) path's telemetry is unchanged. Result types render byte-for-byte identical to the +pool and connection lifecycle. The Go wrapper telemetry interceptor is skipped on the +kernel path so it does not duplicate kernel-owned telemetry for the same connection and +statements; `enableTelemetry` and `telemetry_batch_size` are forwarded to kernel-owned +telemetry config. Result types render byte-for-byte identical to the Thrift backend: scalars, DECIMAL (exact string), TIMESTAMP / TIMESTAMP_NTZ (shifted into the session time zone), INTERVAL, nested Array/Map/Struct and VARIANT (as JSON), and GEOMETRY / GEOGRAPHY (WKT). The server query id is surfaced on the success path, so a diff --git a/internal/backend/kernel/config.go b/internal/backend/kernel/config.go index 0ae2c7c0..4ff099f8 100644 --- a/internal/backend/kernel/config.go +++ b/internal/backend/kernel/config.go @@ -76,6 +76,15 @@ type Config struct { // DecimalAsFloat scans top-level DECIMAL columns to a lossy float64 instead of // the exact string (from WithKernelDecimalAsFloat). Kernel still sends Decimal128. DecimalAsFloat bool + + // Telemetry carries the kernel-owned telemetry collection settings. The Go + // wrapper telemetry interceptor is disabled on the kernel path; these fields + // are the configuration the kernel must use for its own telemetry runtime. + Telemetry *TelemetryConfig + + // DriverSystemConfiguration is the driver/runtime identity stamped onto + // kernel-owned telemetry. Nil lets the kernel use its built-in defaults. + DriverSystemConfiguration *DriverSystemConfiguration } // RetryConfig is the driver's HTTP retry policy forwarded to the kernel: the @@ -93,3 +102,30 @@ type RetryConfig struct { // default (900s). Mirrors the pyo3/napi retry_overall_timeout knob. OverallTimeout time.Duration } + +// TelemetryConfig is the kernel telemetry subset exposed by the Go driver. It +// mirrors the Python kernel kwargs from databricks-sql-python#925: +// telemetry_enabled follows the user-supplied enableTelemetry value, defaulting +// to true when unset; telemetry_batch_size is forwarded only when explicitly set. +type TelemetryConfig struct { + Enabled bool + BatchSize int +} + +// DriverSystemConfiguration mirrors the kernel's DriverSystemConfiguration +// fields without importing the Go telemetry package into this lower-level +// backend package. +type DriverSystemConfiguration struct { + DriverVersion string + RuntimeName string + RuntimeVersion string + RuntimeVendor string + OSName string + OSVersion string + OSArch string + DriverName string + ClientAppName string + LocaleName string + CharSetEncoding string + ProcessName string +} diff --git a/kernel_config.go b/kernel_config.go index 4a2e346e..098e893a 100644 --- a/kernel_config.go +++ b/kernel_config.go @@ -14,6 +14,7 @@ import ( "github.com/databricks/databricks-sql-go/internal/backend/kernel" "github.com/databricks/databricks-sql-go/internal/client" "github.com/databricks/databricks-sql-go/internal/config" + "github.com/databricks/databricks-sql-go/telemetry" ) // This file is intentionally NOT behind the `cgo && databricks_kernel` build tag. @@ -128,6 +129,11 @@ func buildKernelConfig(cfg *config.Config, kauth kernel.Auth) kernel.Config { // the same effective params the Thrift backend forwards, so they flow to the // server identically with no per-backend translation. SessionConf: cfg.EffectiveSessionParams(), + // Kernel-owned telemetry. The Go wrapper interceptor is skipped on the + // kernel path, but the kernel still needs the user's telemetry knobs and + // this binding's system identity for its own runtime. + Telemetry: kernelTelemetryConfig(cfg), + DriverSystemConfiguration: kernelDriverSystemConfiguration(cfg), } // TLS: the driver honors TLSConfig only for InsecureSkipVerify (see // internal/client), so map exactly that knob to the kernel. @@ -159,6 +165,38 @@ func buildKernelConfig(cfg *config.Config, kauth kernel.Auth) kernel.Config { return kc } +func kernelTelemetryConfig(cfg *config.Config) *kernel.TelemetryConfig { + enabled := true + if val, isSet := cfg.EnableTelemetry.Get(); isSet { + enabled = val + } + return &kernel.TelemetryConfig{ + Enabled: enabled, + BatchSize: cfg.TelemetryBatchSize, + } +} + +func kernelDriverSystemConfiguration(cfg *config.Config) *kernel.DriverSystemConfiguration { + system := telemetry.GetSystemConfiguration(cfg.DriverVersion) + if system == nil { + return nil + } + return &kernel.DriverSystemConfiguration{ + DriverVersion: system.DriverVersion, + RuntimeName: system.RuntimeName, + RuntimeVersion: system.RuntimeVersion, + RuntimeVendor: system.RuntimeVendor, + OSName: system.OSName, + OSVersion: system.OSVersion, + OSArch: system.OSArch, + DriverName: system.DriverName, + ClientAppName: system.ClientAppName, + LocaleName: system.LocaleName, + CharSetEncoding: system.CharSetEncoding, + ProcessName: system.ProcessName, + } +} + // kernelRetryPlaceholderWaits are the backoff bounds substituted when the caller // gave no valid wait range but a definite attempt count to honor — the disable form // (RetryMax < 0), or WithRetries(n, 0, 0) where WithDefaults' waits were overwritten diff --git a/kernel_config_test.go b/kernel_config_test.go index bf0c97c1..13e41f33 100644 --- a/kernel_config_test.go +++ b/kernel_config_test.go @@ -340,10 +340,12 @@ var kernelConfigFieldDisposition = map[string]string{ // Rides in the forwarded User-Agent header (set_custom_header). "UserAgentEntry": "forwarded", - // Not applicable to the kernel path (Thrift/HTTP-transport or telemetry knobs - // that don't reach the kernel binding). - "EnableTelemetry": "inert", - "TelemetryBatchSize": "inert", + // Forwarded to kernel-owned telemetry. The Go wrapper interceptor is disabled + // on the kernel path, but the kernel still needs these knobs for its own + // telemetry runtime. + "EnableTelemetry": "forwarded", + "TelemetryBatchSize": "forwarded", + // Wrapper-only scheduling knob; the kernel owns its own flush policy. "TelemetryFlushInterval": "inert", "UseArrowNativeDecimalDSN": "inert", // DSN carrier; kernel renders decimals exactly regardless @@ -438,6 +440,7 @@ func TestBuildKernelConfig(t *testing.T) { t.Run("core fields + auth forwarded", func(t *testing.T) { c := baseKernelConfig() + c.DriverVersion = "9.8.7-test" c.Catalog = "main" c.Schema = "sys" kauth := kernel.Auth{Mode: kernel.AuthPAT, Token: "dapi-x"} @@ -456,6 +459,57 @@ func TestBuildKernelConfig(t *testing.T) { if want := client.BuildUserAgent(c); kc.UserAgent == "" || kc.UserAgent != want { t.Errorf("UserAgent not forwarded: got %q, want %q", kc.UserAgent, want) } + if kc.DriverSystemConfiguration == nil { + t.Fatal("DriverSystemConfiguration not forwarded") + } + if kc.DriverSystemConfiguration.DriverName != "databricks-sql-go" { + t.Errorf("DriverSystemConfiguration.DriverName = %q, want databricks-sql-go", + kc.DriverSystemConfiguration.DriverName) + } + if kc.DriverSystemConfiguration.DriverVersion != "9.8.7-test" { + t.Errorf("DriverSystemConfiguration.DriverVersion = %q, want 9.8.7-test", + kc.DriverSystemConfiguration.DriverVersion) + } + if kc.DriverSystemConfiguration.RuntimeName != "go" { + t.Errorf("DriverSystemConfiguration.RuntimeName = %q, want go", + kc.DriverSystemConfiguration.RuntimeName) + } + }) + + t.Run("kernel telemetry config defaults enabled and omits unset batch size", func(t *testing.T) { + c := baseKernelConfig() + kc := buildKernelConfig(c, kernel.Auth{Mode: kernel.AuthPAT, Token: "dapi-x"}) + if kc.Telemetry == nil { + t.Fatal("Telemetry not forwarded") + } + if !kc.Telemetry.Enabled { + t.Error("Telemetry.Enabled = false, want true when enableTelemetry is unset") + } + if kc.Telemetry.BatchSize != 0 { + t.Errorf("Telemetry.BatchSize = %d, want 0 when telemetry_batch_size is unset", kc.Telemetry.BatchSize) + } + }) + + t.Run("kernel telemetry config follows explicit enableTelemetry and batch size", func(t *testing.T) { + c := baseKernelConfig() + c.EnableTelemetry = config.NewConfigValue(false) + c.TelemetryBatchSize = 17 + kc := buildKernelConfig(c, kernel.Auth{Mode: kernel.AuthPAT, Token: "dapi-x"}) + if kc.Telemetry == nil { + t.Fatal("Telemetry not forwarded") + } + if kc.Telemetry.Enabled { + t.Error("Telemetry.Enabled = true, want false from explicit enableTelemetry=false") + } + if kc.Telemetry.BatchSize != 17 { + t.Errorf("Telemetry.BatchSize = %d, want 17", kc.Telemetry.BatchSize) + } + + c.EnableTelemetry = config.NewConfigValue(true) + kc = buildKernelConfig(c, kernel.Auth{Mode: kernel.AuthPAT, Token: "dapi-x"}) + if !kc.Telemetry.Enabled { + t.Error("Telemetry.Enabled = false, want true from explicit enableTelemetry=true") + } }) t.Run("MaxChunksInMemory injected into kernel SessionConf", func(t *testing.T) { diff --git a/telemetry/system_info.go b/telemetry/system_info.go index 56b979b9..d0cc90b6 100644 --- a/telemetry/system_info.go +++ b/telemetry/system_info.go @@ -47,6 +47,14 @@ func getSystemConfiguration(driverVersion string) *DriverSystemConfiguration { } } +// GetSystemConfiguration returns the driver/runtime identity used in telemetry +// payloads. Kernel-backed connections pass the same identity into the kernel so +// kernel-owned telemetry is attributed to the Go driver rather than the kernel +// binding defaults. +func GetSystemConfiguration(driverVersion string) *DriverSystemConfiguration { + return getSystemConfiguration(driverVersion) +} + func getOSName() string { switch runtime.GOOS { case "darwin":