diff --git a/KERNEL_REV b/KERNEL_REV index 0fc1de14..74115902 100644 --- a/KERNEL_REV +++ b/KERNEL_REV @@ -1 +1 @@ -4d301455f7e70de9cb747e0f1143b8a3df8c5b48 +f29957733c8f6d79a7212bc8a50cf4e5ffa8dd69 diff --git a/connector.go b/connector.go index 3838cb9d..b9b42969 100644 --- a/connector.go +++ b/connector.go @@ -45,14 +45,16 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) { if c.cfg.UseKernel { be, err = newKernelBackend(ctx, c.cfg) } else { - // The experimental WithKernel* TLS options have no Thrift-path equivalent — - // reject them loudly rather than silently ignore, so a caller who sets a - // trusted-CA bundle / an independent hostname skip and forgets WithUseKernel - // learns the option had no effect instead of connecting with a - // weaker-than-intended (or unconfigured) TLS trust store. + // 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 "+ - "(WithKernelTrustedCerts / WithKernelSkipHostnameVerify) %w; "+ + 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) @@ -97,6 +99,15 @@ 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 @@ -603,3 +614,85 @@ func WithKernelSkipHostnameVerify() ConnOption { kernelExperimental(c).TLSSkipHostnameVerify = true } } + +// KernelProxy is the explicit-proxy configuration for WithKernelProxy. Its fields +// are named so a call site can't transpose the credentials — the four values are +// all strings, so a positional signature would let a Username/Password swap (or a +// misplaced BypassHosts) compile cleanly and fail only at runtime with wrong proxy +// credentials. +type KernelProxy struct { + // URL is the proxy URL (e.g. "http://proxy.internal:3128"). Required; empty is a + // no-op, leaving the environment-derived proxy (if any) in effect. + URL string + // Username / Password are optional out-of-band basic-auth credentials, supplied + // here rather than embedded in the URL userinfo. Empty means unset. + Username string + Password string + // BypassHosts is an optional comma-separated no-proxy host list. NO_PROXY is + // consumed during environment resolution and not forwarded to the kernel, so this + // is the only way to give the kernel a structured bypass list. Empty means unset. + BypassHosts string +} + +// WithKernelProxy configures an explicit HTTP proxy for the kernel backend, with +// optional out-of-band basic-auth credentials and a comma-separated bypass +// (no-proxy) host list. It overrides the HTTP(S)_PROXY / NO_PROXY environment the +// kernel path otherwise mirrors from the Thrift path. +// +// Use this instead of the proxy environment when you need the "advanced" fields +// the env-var path can't express: a structured bypass list (NO_PROXY is consumed +// during environment resolution, not forwarded to the kernel) or basic-auth +// credentials supplied out of band rather than embedded in the URL userinfo. +// KernelProxy.Username / Password / BypassHosts may be empty (passed to the kernel +// as NULL, i.e. unset). An empty KernelProxy.URL is a no-op — the environment-derived +// proxy, if any, stays in effect. A malformed URL is rejected at connect +// (errors.Is ErrInvalidKernelConfig). +// +// An explicit WithKernelProxy takes precedence over the environment: consulting +// both would be ambiguous, and an explicit proxy is a deliberate override. +// +// EXPERIMENTAL, kernel-only: the default (Thrift) backend rejects this at connect. +func WithKernelProxy(p KernelProxy) ConnOption { + return func(c *config.Config) { + ke := kernelExperimental(c) + ke.ProxyURL = p.URL + ke.ProxyUsername = p.Username + ke.ProxyPassword = p.Password + ke.ProxyBypassHosts = p.BypassHosts + } +} + +// WithKernelRetryOverallTimeout sets the cumulative retry budget across all +// attempts on the kernel backend — the total time the kernel may spend retrying a +// single logical request before giving up. This is the 4th retry knob, alongside +// the backoff bounds and max attempts carried by the backend-neutral WithRetries +// (RetryWaitMin / RetryWaitMax / RetryMax, which the kernel path also honors). +// +// It is a kernel-only option because the Thrift-path WithRetries surface has no +// overall-budget equivalent; it mirrors the pyo3/napi retry_overall_timeout knob. +// Zero (the default) keeps the kernel's built-in budget (900s). +// +// EXPERIMENTAL, kernel-only: the default (Thrift) backend rejects this at connect. +func WithKernelRetryOverallTimeout(d time.Duration) ConnOption { + return func(c *config.Config) { + kernelExperimental(c).RetryOverallTimeout = d + } +} + +// WithKernelMaxChunksInMemory bounds how many decompressed CloudFetch chunks the +// kernel holds in memory at once on the kernel backend — the knob that trades +// large-result throughput for peak RSS. Lower it (e.g. 4) to cap memory on wide, +// row-heavy result sets; raise it for more download parallelism at higher memory. +// A value <= 0 (the default) leaves the kernel's built-in default (16) in place. +// +// It is forwarded as the kernel's client-only "cloudfetch_max_chunks_in_memory" +// session conf, which the kernel applies to its result config and strips before +// the SEA wire — so it never reaches the server. +// +// EXPERIMENTAL, kernel-only: the default (Thrift) backend has no in-memory-chunk +// knob and rejects this at connect. +func WithKernelMaxChunksInMemory(n int) ConnOption { + return func(c *config.Config) { + kernelExperimental(c).MaxChunksInMemory = n + } +} diff --git a/doc.go b/doc.go index 2febbf55..24a67391 100644 --- a/doc.go +++ b/doc.go @@ -214,13 +214,14 @@ via the authType=oauthU2M DSN param); reading scalar, nested, and complex-typed results (CloudFetch is transparent); bound query parameters (positional and named); context cancellation during execute; the initial namespace (WithInitialNamespace, applied post-connect via USE CATALOG / USE SCHEMA); metric-view metadata -(WithEnableMetricViewMetadata); and the TLS, proxy, and session-conf (query tags, -statement timeout, time zone) options. Nothing is silently ignored: WithTimeout, a -retries-disabling WithRetries, token-provider / external / federated -authenticators, and custom M2M OAuth scopes (the kernel applies its own) are -rejected at connect; staging (PUT/GET/REMOVE on a Unity Catalog volume) is -rejected at execute. WithMaxRows and positive-limit WithRetries are -accepted but inert (the kernel manages fetching and retries below the C ABI). +(WithEnableMetricViewMetadata); the retry / backoff policy (WithRetries: +RetryWaitMin / RetryWaitMax / RetryMax, including the disable form, forwarded to the +kernel's HTTP retry config); and the TLS, proxy, and session-conf (query tags, +statement timeout, time zone) options. Nothing is silently ignored: WithTimeout, +token-provider / external / federated authenticators, and custom M2M OAuth scopes +(the kernel applies its own) are rejected at connect; staging (PUT/GET/REMOVE on a +Unity Catalog volume) is rejected at execute. WithMaxRows is accepted but inert (the +kernel manages fetching below the C ABI). OAuth U2M is interactive: on a cache miss, connecting launches the system browser and blocks until login completes or the kernel's ~120s callback timeout expires. Because @@ -233,23 +234,47 @@ backend exposes a U2M-scopes option, so this is a fixed difference between the t default sets, not a dropped setting; both authorize against the built-in public client. -Experimental kernel-only TLS options (rejected by the default backend; the +Experimental kernel-only options (rejected by the default backend; the WithKernel* prefix marks them experimental): - WithKernelTrustedCerts(pem) adds a PEM CA bundle on top of the system roots (for a re-signing proxy or on-prem CA). Required because the kernel's TLS stack does not read SSL_CERT_FILE. - WithKernelSkipHostnameVerify() skips only the hostname check while keeping chain validation (finer-grained than WithSkipTLSHostVerify). - -Setting either without WithUseKernel fails Connect with an error wrapping the + - WithKernelProxy(KernelProxy{URL, Username, Password, BypassHosts}) sets an explicit + HTTP proxy, overriding the HTTP(S)_PROXY / NO_PROXY environment the kernel path + otherwise mirrors. The fields are named (not positional) so the four same-typed + strings can't be transposed at the call site. Use it for the advanced fields the + env-var path can't express: a structured bypass (no-proxy) list, or basic-auth + credentials supplied out of band rather than embedded in the URL. An explicit proxy + takes precedence over the environment; empty credentials / bypass are passed to the + kernel as unset, and a malformed URL is rejected at connect. + - WithKernelRetryOverallTimeout(d) sets the cumulative retry budget across all + attempts — the 4th retry knob, alongside the backoff bounds and max attempts + carried by the backend-neutral WithRetries (also honored on the kernel path). It + is kernel-only because the Thrift WithRetries surface has no overall-budget + equivalent; it mirrors the pyo3/napi retry_overall_timeout knob. Zero keeps the + kernel's default budget (900s). + - WithKernelMaxChunksInMemory(n) bounds how many decompressed CloudFetch chunks the + kernel holds in memory at once — the knob that trades large-result download + throughput for peak RSS. Lower it (e.g. 4) to cap memory on wide, row-heavy result + sets; raise it for more parallelism at higher memory. It is forwarded as the + kernel's client-only cloudfetch_max_chunks_in_memory session conf, stripped before + the SEA wire so it never reaches the server. A value <= 0 keeps the kernel's + default (16). + +Setting any of these without WithUseKernel fails Connect with an error wrapping the 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. Result types render byte-for-byte identical to the +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 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 (WKT). The server query id is surfaced on the success path, so a +and GEOMETRY / GEOGRAPHY (WKT). The server query id is surfaced on the success path, so a QueryIdCallback (see below) fires with the real id and EXECUTE_STATEMENT telemetry carries it. @@ -257,6 +282,13 @@ On the read path, context cancellation is honored at result-batch boundaries, no mid-fetch: an in-flight CloudFetch batch runs to completion before the cancel takes effect. +OAuth token caching and HTTP client reuse are inherited from the kernel and need no +driver configuration: the kernel caches U2M tokens on disk +(~/.config/databricks-sql-kernel/oauth/, with refresh-token lifecycle) and M2M +tokens in-memory with background refresh, and reuses a single pooled HTTP client per +session across the control-plane, CloudFetch, and auth-refresh calls. There is no +driver-side cache or pool knob because these live below the C ABI. + # Programmatically Retrieving Connection and Query Id Use the driverctx package under driverctx/ctx.go to add callbacks to the query context to receive the connection id and query id. diff --git a/errors/errors.go b/errors/errors.go index 8eac7877..c90f3b1a 100644 --- a/errors/errors.go +++ b/errors/errors.go @@ -82,6 +82,15 @@ var ErrKernelNotCompiled error = errors.New("the SEA-via-kernel backend is not c // the case programmatically instead of matching on message text. var ErrRequiresKernelBackend error = errors.New("requires the SEA-via-kernel backend") +// value to be used with errors.Is() to determine that a kernel-backend option was +// itself malformed (e.g. a WithKernelProxy URL that does not parse) — as opposed to +// unsupported (ErrNotSupportedByKernel) or a transient connect failure. The kernel +// path validates such options in the Go layer before handing them to the kernel's C +// ABI, where the failure would otherwise surface as an opaque wrapped string a +// caller can't branch on. Lets a caller distinguish "fix your config" from a +// retryable error instead of matching on message text. +var ErrInvalidKernelConfig error = errors.New("invalid kernel backend configuration") + // Base interface for driver errors type DBError interface { // Descriptive message describing the error diff --git a/internal/arrowscan/coltype.go b/internal/arrowscan/coltype.go new file mode 100644 index 00000000..e795fcfc --- /dev/null +++ b/internal/arrowscan/coltype.go @@ -0,0 +1,145 @@ +package arrowscan + +import ( + "database/sql" + "math" + "reflect" + "time" + + "github.com/apache/arrow/go/v12/arrow" +) + +// ColumnTypeInfo is the per-column metadata database/sql surfaces through +// sql.ColumnType — the DatabaseTypeName, ScanType, and Length the optional +// driver.RowsColumnType* interfaces return. The kernel backend derives it from +// the result's Arrow schema; ColumnTypeInfoFor is the single mapping both the +// value scanner (ScanCellCached) and the type-metadata reporter agree on, so a +// column's reported type can never drift from what a row actually scans into. +// +// The mapping mirrors the Thrift backend (internal/rows/rows.go getScanType / +// ColumnTypeDatabaseTypeName / ColumnTypeLength) so a query reports byte-identical +// column metadata on either backend — the same "identical results across backends" +// contract the value renderers hold. The pure-Go guards are TestColumnTypeInfoFor, +// TestColumnTypeInfoScanTypeCoversScanner, and TestColumnTypeInfoMatchesThriftMapping +// (which cross-checks this mapping against the Thrift functions directly); the live +// TestKernelThriftColumnTypeParity confirms it against a real warehouse. +type ColumnTypeInfo struct { + // DatabaseTypeName is the Databricks type name (e.g. "BIGINT", "DECIMAL", + // "ARRAY"), matching what the Thrift path reports; "" for a type with no + // Databricks name. + DatabaseTypeName string + // ScanType is the Go type database/sql recommends scanning the column into, + // matching the Thrift path. nil (only for the NULL type) makes database/sql + // fall back to interface{}, exactly as the Thrift path's nil scan type does. + ScanType reflect.Type + // Length / HasLength report a variable-length column's length. As on the + // Thrift path, only variable-length types (string/binary/nested, and the + // server-stringified interval/geo types) report a length, and the reported + // value is math.MaxInt64 (unbounded); fixed-width types report (0, false). + Length int64 + HasLength bool +} + +// Scan types, matching the Thrift path's vars in internal/rows/rows.go so both +// backends recommend the identical Go destination type per column. +var ( + scanTypeBool = reflect.TypeOf(true) + scanTypeInt8 = reflect.TypeOf(int8(0)) + scanTypeInt16 = reflect.TypeOf(int16(0)) + scanTypeInt32 = reflect.TypeOf(int32(0)) + scanTypeInt64 = reflect.TypeOf(int64(0)) + scanTypeFloat32 = reflect.TypeOf(float32(0)) + scanTypeFloat64 = reflect.TypeOf(float64(0)) + scanTypeString = reflect.TypeOf("") + scanTypeDateTime = reflect.TypeOf(time.Time{}) + scanTypeRawBytes = reflect.TypeOf(sql.RawBytes{}) + scanTypeUnknown = reflect.TypeOf(new(any)) // *interface{}, as on the Thrift path +) + +// ColumnTypeInfoFor maps an Arrow column type to the metadata database/sql +// exposes, matching the Thrift backend for every Databricks type. The Arrow types +// listed here are exactly those ScanCellCached scans, so the type reported for a +// column and the value produced for its cells stay in lockstep. +// +// Notably: DECIMAL reports sql.RawBytes (the Thrift scan type for DECIMAL) even +// though the value is rendered as an exact string — matching Thrift, which also +// scans DECIMAL into RawBytes; and the interval / geo types report STRING because +// the Thrift server pre-formats them to strings and the kernel formats them +// Go-side to the same string, so both are indistinguishable to a caller. +func ColumnTypeInfoFor(dt arrow.DataType) ColumnTypeInfo { + switch dt.ID() { + case arrow.BOOL: + return ColumnTypeInfo{DatabaseTypeName: "BOOLEAN", ScanType: scanTypeBool} + case arrow.INT8: + return ColumnTypeInfo{DatabaseTypeName: "TINYINT", ScanType: scanTypeInt8} + case arrow.INT16: + return ColumnTypeInfo{DatabaseTypeName: "SMALLINT", ScanType: scanTypeInt16} + case arrow.INT32: + return ColumnTypeInfo{DatabaseTypeName: "INT", ScanType: scanTypeInt32} + case arrow.INT64: + return ColumnTypeInfo{DatabaseTypeName: "BIGINT", ScanType: scanTypeInt64} + case arrow.UINT8, arrow.UINT16, arrow.UINT32, arrow.UINT64: + // Databricks SQL has no unsigned types, so these do not occur in practice; + // this arm is defensive and stays in lockstep with ScanCellCached, which + // widens every unsigned integer to int64 (driver.Value has no uint64). Report + // BIGINT/int64 to match that scanned value rather than falling through to the + // generic *interface{} default. + return ColumnTypeInfo{DatabaseTypeName: "BIGINT", ScanType: scanTypeInt64} + case arrow.FLOAT32: + return ColumnTypeInfo{DatabaseTypeName: "FLOAT", ScanType: scanTypeFloat32} + case arrow.FLOAT64: + return ColumnTypeInfo{DatabaseTypeName: "DOUBLE", ScanType: scanTypeFloat64} + case arrow.STRING, arrow.LARGE_STRING: + // STRING covers VARCHAR/CHAR/VARIANT/GEOMETRY/GEOGRAPHY too — all arrive as + // Arrow Utf8, and the Thrift path collapses them to STRING as well. + return varLen("STRING", scanTypeString) + case arrow.BINARY: + return varLen("BINARY", scanTypeRawBytes) + case arrow.DATE32, arrow.DATE64: + return ColumnTypeInfo{DatabaseTypeName: "DATE", ScanType: scanTypeDateTime} + case arrow.TIMESTAMP: + return ColumnTypeInfo{DatabaseTypeName: "TIMESTAMP", ScanType: scanTypeDateTime} + case arrow.DECIMAL128: + // Thrift scans DECIMAL into sql.RawBytes; match it even though the kernel + // renders the value as an exact fixed-point string (both convert cleanly to + // a caller's *string/*[]byte, and reporting the same scan type keeps the + // backends indistinguishable). + return ColumnTypeInfo{DatabaseTypeName: "DECIMAL", ScanType: scanTypeRawBytes} + case arrow.LIST, arrow.LARGE_LIST, arrow.FIXED_SIZE_LIST: + return varLen("ARRAY", scanTypeRawBytes) + case arrow.MAP: + return varLen("MAP", scanTypeRawBytes) + case arrow.STRUCT: + return varLen("STRUCT", scanTypeRawBytes) + case arrow.DURATION: + // INTERVAL DAY TO SECOND. Parity target is what the Thrift backend REPORTS, + // which is config-dependent: in the prod default (native-interval Arrow off) + // the server pre-formats intervals to text and declares the Thrift column + // STRING_TYPE, so Thrift's GetDBTypeName yields "STRING" and MaxInt64 length + // — verified live against both backends. We therefore report STRING here even + // though the kernel receives a native arrow.DURATION (which it formats Go-side + // to the identical string). If a warehouse ever enables native-interval Thrift + // the server would instead declare INTERVAL_DAY_TIME and Thrift would report + // that; matching STRING is correct for the default path the parity test pins, + // and the scanned VALUE is identical either way — only this label would differ. + return varLen("STRING", scanTypeString) + case arrow.INTERVAL_MONTHS: + // INTERVAL YEAR TO MONTH — same server-config reasoning as arrow.DURATION. + return varLen("STRING", scanTypeString) + case arrow.NULL: + // The NULL type has no scan type on the Thrift path (nil → database/sql + // falls back to interface{}); mirror that. + return ColumnTypeInfo{DatabaseTypeName: "NULL", ScanType: nil} + default: + // A type ScanCellCached does not handle: report the Thrift default scan type + // (*interface{}) and no database name, rather than inventing one. + return ColumnTypeInfo{DatabaseTypeName: "", ScanType: scanTypeUnknown} + } +} + +// varLen builds a ColumnTypeInfo for a variable-length type, which reports an +// unbounded length (math.MaxInt64) just as the Thrift path does for +// string/binary/nested columns. +func varLen(name string, scan reflect.Type) ColumnTypeInfo { + return ColumnTypeInfo{DatabaseTypeName: name, ScanType: scan, Length: math.MaxInt64, HasLength: true} +} diff --git a/internal/arrowscan/coltype_test.go b/internal/arrowscan/coltype_test.go new file mode 100644 index 00000000..ac71ad0c --- /dev/null +++ b/internal/arrowscan/coltype_test.go @@ -0,0 +1,106 @@ +package arrowscan + +import ( + "database/sql" + "math" + "reflect" + "testing" + "time" + + "github.com/apache/arrow/go/v12/arrow" +) + +// TestColumnTypeInfoFor pins the Arrow → column-metadata mapping the kernel +// backend reports through sql.ColumnType, so it stays byte-identical to the +// Thrift path (internal/rows/rows.go). This is a pure-Go, no-warehouse guard: the +// gap it regresses (PECOBLR-3692) was invisible to the value-parity suites +// because they compare scanned VALUES, not Rows.ColumnType* metadata. The +// expected DatabaseTypeName / ScanType / Length values are the ground truth +// captured from the Thrift backend on a live warehouse. +func TestColumnTypeInfoFor(t *testing.T) { + str := reflect.TypeOf("") + raw := reflect.TypeOf(sql.RawBytes{}) + tm := reflect.TypeOf(time.Time{}) + + cases := []struct { + name string + dt arrow.DataType + wantDB string + wantST reflect.Type + wantLen int64 + wantOk bool + }{ + {"bool", arrow.FixedWidthTypes.Boolean, "BOOLEAN", reflect.TypeOf(true), 0, false}, + {"int8", arrow.PrimitiveTypes.Int8, "TINYINT", reflect.TypeOf(int8(0)), 0, false}, + {"int16", arrow.PrimitiveTypes.Int16, "SMALLINT", reflect.TypeOf(int16(0)), 0, false}, + {"int32", arrow.PrimitiveTypes.Int32, "INT", reflect.TypeOf(int32(0)), 0, false}, + {"int64", arrow.PrimitiveTypes.Int64, "BIGINT", reflect.TypeOf(int64(0)), 0, false}, + // Unsigned types don't occur in Databricks SQL, but ScanCellCached widens them + // to int64, so the metadata reports BIGINT/int64 to match (not the fallthrough). + {"uint8", arrow.PrimitiveTypes.Uint8, "BIGINT", reflect.TypeOf(int64(0)), 0, false}, + {"uint16", arrow.PrimitiveTypes.Uint16, "BIGINT", reflect.TypeOf(int64(0)), 0, false}, + {"uint32", arrow.PrimitiveTypes.Uint32, "BIGINT", reflect.TypeOf(int64(0)), 0, false}, + {"uint64", arrow.PrimitiveTypes.Uint64, "BIGINT", reflect.TypeOf(int64(0)), 0, false}, + {"float32", arrow.PrimitiveTypes.Float32, "FLOAT", reflect.TypeOf(float32(0)), 0, false}, + {"float64", arrow.PrimitiveTypes.Float64, "DOUBLE", reflect.TypeOf(float64(0)), 0, false}, + {"string", arrow.BinaryTypes.String, "STRING", str, math.MaxInt64, true}, + {"binary", arrow.BinaryTypes.Binary, "BINARY", raw, math.MaxInt64, true}, + {"date32", arrow.FixedWidthTypes.Date32, "DATE", tm, 0, false}, + {"timestamp", arrow.FixedWidthTypes.Timestamp_us, "TIMESTAMP", tm, 0, false}, + {"decimal", &arrow.Decimal128Type{Precision: 10, Scale: 2}, "DECIMAL", raw, 0, false}, + {"list", arrow.ListOf(arrow.PrimitiveTypes.Int64), "ARRAY", raw, math.MaxInt64, true}, + {"map", arrow.MapOf(arrow.BinaryTypes.String, arrow.PrimitiveTypes.Int64), "MAP", raw, math.MaxInt64, true}, + {"struct", arrow.StructOf(arrow.Field{Name: "x", Type: arrow.PrimitiveTypes.Int64}), "STRUCT", raw, math.MaxInt64, true}, + {"duration", &arrow.DurationType{Unit: arrow.Microsecond}, "STRING", str, math.MaxInt64, true}, + {"month_interval", arrow.FixedWidthTypes.MonthInterval, "STRING", str, math.MaxInt64, true}, + {"null", arrow.Null, "NULL", nil, 0, false}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + got := ColumnTypeInfoFor(c.dt) + if got.DatabaseTypeName != c.wantDB { + t.Errorf("DatabaseTypeName = %q, want %q", got.DatabaseTypeName, c.wantDB) + } + if got.ScanType != c.wantST { + t.Errorf("ScanType = %v, want %v", got.ScanType, c.wantST) + } + if got.Length != c.wantLen || got.HasLength != c.wantOk { + t.Errorf("Length = (%d,%v), want (%d,%v)", got.Length, got.HasLength, c.wantLen, c.wantOk) + } + }) + } +} + +// TestColumnTypeInfoScanTypeCoversScanner is the lockstep guard: every Arrow type +// the value scanner (ScanCellCached) handles must also have a non-fallback entry +// in ColumnTypeInfoFor, so a future scalar type added to the scanner without a +// matching type-metadata entry is caught here rather than silently reporting the +// generic *interface{} scan type at runtime. It checks the representative arrays +// the scanner switches on; the NULL type is intentionally excluded (its nil scan +// type is the correct, Thrift-matching value). +func TestColumnTypeInfoScanTypeCoversScanner(t *testing.T) { + unknown := reflect.TypeOf(new(any)) + scannerTypes := []arrow.DataType{ + arrow.FixedWidthTypes.Boolean, + arrow.PrimitiveTypes.Int8, arrow.PrimitiveTypes.Int16, + arrow.PrimitiveTypes.Int32, arrow.PrimitiveTypes.Int64, + arrow.PrimitiveTypes.Uint8, arrow.PrimitiveTypes.Uint16, + arrow.PrimitiveTypes.Uint32, arrow.PrimitiveTypes.Uint64, + arrow.PrimitiveTypes.Float32, arrow.PrimitiveTypes.Float64, + arrow.BinaryTypes.String, arrow.BinaryTypes.Binary, + arrow.FixedWidthTypes.Date32, arrow.FixedWidthTypes.Timestamp_us, + &arrow.Decimal128Type{Precision: 10, Scale: 2}, + arrow.ListOf(arrow.PrimitiveTypes.Int64), + arrow.MapOf(arrow.BinaryTypes.String, arrow.PrimitiveTypes.Int64), + arrow.StructOf(arrow.Field{Name: "x", Type: arrow.PrimitiveTypes.Int64}), + &arrow.DurationType{Unit: arrow.Microsecond}, + arrow.FixedWidthTypes.MonthInterval, + } + for _, dt := range scannerTypes { + info := ColumnTypeInfoFor(dt) + if info.ScanType == unknown || info.DatabaseTypeName == "" { + t.Errorf("%s: scanner-handled type fell through to the default mapping (db=%q scan=%v)", + dt, info.DatabaseTypeName, info.ScanType) + } + } +} diff --git a/internal/backend/kernel/backend.go b/internal/backend/kernel/backend.go index 361d8a2b..febf37e4 100644 --- a/internal/backend/kernel/backend.go +++ b/internal/backend/kernel/backend.go @@ -5,6 +5,23 @@ package kernel /* #include #include "databricks_kernel.h" + +// go_kernel_set_retry_config wraps kernel_session_config_set_retry_config. cgo's +// C parser silently drops the direct declaration of that symbol (a known cgo +// quirk: the declaration is valid C — it compiles under gcc and links from the +// archive — but cgo omits it from the generated bindings, so a direct +// C.kernel_session_config_set_retry_config call fails to build with "could not +// determine what … refers to"). A static inline shim forwarding to it IS parsed +// and links fine. The shim must live in the SAME file's preamble as its caller +// (applyRetry, below): a shim placed in another file's cgo preamble in this +// package is itself dropped the same way. Keeps the kernel header unchanged (it +// is valid C, used verbatim by the C-only ODBC consumer). +static inline KernelStatusCode go_kernel_set_retry_config( + KernelSessionConfig* config, uint64_t min_wait_ms, uint64_t max_wait_ms, + uint32_t max_retries, uint64_t overall_timeout_ms) { + return kernel_session_config_set_retry_config( + config, min_wait_ms, max_wait_ms, max_retries, overall_timeout_ms); +} */ import "C" @@ -140,18 +157,14 @@ func (k *KernelBackend) OpenSession(ctx context.Context) error { return err } - // Proxy: only when the environment configured one for this endpoint. NO_PROXY - // was already applied during resolution, so no bypass list is needed here; - // any credentials are carried in the URL userinfo (Go's proxy-env convention), - // so username/password are NULL. - if k.cfg.ProxyURL != "" { - url := newCStr(k.cfg.ProxyURL) - defer url.free() - if err := call(func() C.KernelStatusCode { - return C.kernel_session_config_set_proxy(cfg, url.c, nil, nil, nil) - }); err != nil { - return fmt.Errorf("kernel: set_proxy: %w", toConnError(err)) - } + // Proxy: env-derived or explicit (WithKernelProxy). See applyProxy. + if err := k.applyProxy(cfg); err != nil { + return err + } + + // Retry / backoff policy (WithRetries). See applyRetry. + if err := k.applyRetry(cfg); err != nil { + return err } // Session confs (STATEMENT_TIMEOUT, QUERY_TAGS, TIMEZONE, …) — the same map @@ -169,6 +182,13 @@ func (k *KernelBackend) OpenSession(ctx context.Context) error { } } + // Log the resolved retry policy + CloudFetch chunk cap before connecting, so if a + // customer reports a hung connect or a large-result OOM, on-call can see from the + // debug log what was actually applied (these are otherwise silent — forwarded to + // the kernel with no observable trace). Kept at Debug and cheap to format. + klogCtx(ctx, "OpenSession resolved: retry=%s cloudfetchMaxChunksInMemory=%s", + describeRetry(k.cfg.Retry), k.cfg.SessionConf[kernelMaxChunksInMemoryConfKey]) + // kernel_session_open takes ownership of cfg here. Its documented C-ABI // contract (databricks_kernel.h: "CONSUMES config on both success and failure // — do not use or free config afterwards") is what makes the unconditional @@ -237,6 +257,74 @@ func (k *KernelBackend) applyKernelTLS(cfg *C.KernelSessionConfig) error { return nil } +// applyProxy forwards the HTTP proxy config to the session config. The URL is set +// when either the environment configured one for this endpoint or the caller +// supplied one via WithKernelProxy (resolveKernelProxy decides which). The +// optional basic-auth credentials and bypass list ride along; each is NULL when +// empty (the env path folds credentials into the URL and consumes NO_PROXY during +// resolution, so it leaves them empty — Go's proxy-env convention). A no-op when +// ProxyURL is empty (direct connection). The kernel applies the URL as an explicit +// override of its own env-var behavior. +func (k *KernelBackend) applyProxy(cfg *C.KernelSessionConfig) error { + if k.cfg.ProxyURL == "" { + return nil + } + url := newCStr(k.cfg.ProxyURL) + defer url.free() + user := newCStrOrNull(k.cfg.ProxyUsername) + defer user.free() + pass := newCStrOrNull(k.cfg.ProxyPassword) + defer pass.free() + bypass := newCStrOrNull(k.cfg.ProxyBypassHosts) + defer bypass.free() + if err := call(func() C.KernelStatusCode { + return C.kernel_session_config_set_proxy(cfg, url.c, user.c, pass.c, bypass.c) + }); err != nil { + return fmt.Errorf("kernel: set_proxy: %w", toConnError(err)) + } + return nil +} + +// applyRetry forwards the driver's HTTP retry / backoff policy to the session +// config. A no-op when Config.Retry is nil, so the kernel's own default policy +// (exponential backoff with jitter, 5 retries, 1s..60s, 900s budget) is preserved +// otherwise. MaxRetries == 0 disables retries; OverallTimeout == 0 keeps the +// kernel's default budget (the setter maps a 0 ms budget to "keep default"). +func (k *KernelBackend) applyRetry(cfg *C.KernelSessionConfig) error { + r := k.cfg.Retry + if r == nil { + return nil + } + if err := call(func() C.KernelStatusCode { + // Via the go_kernel_set_retry_config shim in cgo.go — cgo drops the direct + // declaration of the underlying symbol (see the shim's comment). + return C.go_kernel_set_retry_config(cfg, + C.uint64_t(r.MinWait.Milliseconds()), C.uint64_t(r.MaxWait.Milliseconds()), + C.uint32_t(r.MaxRetries), C.uint64_t(r.OverallTimeout.Milliseconds())) + }); err != nil { + return fmt.Errorf("kernel: set_retry_config: %w", toConnError(err)) + } + return nil +} + +// kernelMaxChunksInMemoryConfKey mirrors config.KernelMaxChunksInMemoryConfKey — +// the client-only session conf carrying WithKernelMaxChunksInMemory. Duplicated as +// a local const rather than importing internal/config, which this cgo backend +// otherwise has no dependency on; it is read here only to log the resolved cap. +const kernelMaxChunksInMemoryConfKey = "cloudfetch_max_chunks_in_memory" + +// describeRetry renders the resolved retry policy for the OpenSession debug log. +// "kernel-default" means Config.Retry is nil, so the kernel keeps its own policy +// (5 retries, 1s..60s, 900s budget); otherwise it shows the forwarded values, with +// maxRetries=0 being the disable form and overallTimeout=0 meaning "keep default". +func describeRetry(r *RetryConfig) string { + if r == nil { + return "kernel-default" + } + return fmt.Sprintf("maxRetries=%d minWait=%s maxWait=%s overallTimeout=%s", + r.MaxRetries, r.MinWait, r.MaxWait, r.OverallTimeout) +} + // setAuth applies the resolved auth form to the session config via exactly one // kernel_session_config_set_auth_* call. PAT and M2M are plain value setters; U2M // records the client id / redirect port / scopes and the kernel owns the browser @@ -314,6 +402,36 @@ func trySetKernelTLS(cfg Config) error { return k.applyKernelTLS(c) } +// trySetProxy allocates a throwaway session config, applies the proxy config from +// cfg to it, and frees it — the analogous test seam to trySetKernelTLS, so a +// tagged test can exercise the real kernel_session_config_set_proxy cgo setter +// (URL + optional NULL-for-empty username / password / bypass) end to end. Not +// used in production. +func trySetProxy(cfg Config) error { + var c *C.KernelSessionConfig + if err := call(func() C.KernelStatusCode { return C.kernel_session_config_new(&c) }); err != nil { + return fmt.Errorf("config_new: %w", err) + } + defer C.kernel_session_config_free(c) + k := &KernelBackend{cfg: cfg} + return k.applyProxy(c) +} + +// trySetRetry allocates a throwaway session config, applies the retry config from +// cfg to it, and frees it — the analogous test seam to trySetProxy, so a tagged +// test can exercise the real kernel_session_config_set_retry_config cgo setter +// (the 4 knobs, plus the InvalidArgument rejections for a degenerate range) end to +// end. Not used in production. +func trySetRetry(cfg Config) error { + var c *C.KernelSessionConfig + if err := call(func() C.KernelStatusCode { return C.kernel_session_config_new(&c) }); err != nil { + return fmt.Errorf("config_new: %w", err) + } + defer C.kernel_session_config_free(c) + k := &KernelBackend{cfg: cfg} + return k.applyRetry(c) +} + // applyInitialNamespace runs USE CATALOG / USE SCHEMA to select the configured // initial namespace, since the kernel C ABI exposes no catalog/schema setter. // Identifiers are backtick-quoted (quoteIdent) so arbitrary names are safe. A diff --git a/internal/backend/kernel/config.go b/internal/backend/kernel/config.go index f0bc8e28..5727407c 100644 --- a/internal/backend/kernel/config.go +++ b/internal/backend/kernel/config.go @@ -35,11 +35,28 @@ type Config struct { TLSTrustedCertsPEM []byte TLSSkipHostnameVerify bool - // ProxyURL configures an HTTP proxy, already resolved for this endpoint from - // the same HTTP(S)_PROXY / NO_PROXY environment the Thrift path uses (NO_PROXY - // is applied during resolution). Empty leaves the kernel on a direct - // connection. - ProxyURL string + // ProxyURL configures an HTTP proxy. It is either resolved for this endpoint + // from the same HTTP(S)_PROXY / NO_PROXY environment the Thrift path uses + // (NO_PROXY applied during resolution), or set explicitly via WithKernelProxy. + // Empty leaves the kernel on a direct connection. + // + // ProxyUsername / ProxyPassword are out-of-band basic-auth credentials (an + // alternative to embedding them in ProxyURL's userinfo). ProxyBypassHosts is + // a comma-separated no-proxy list honored kernel-side. All three are only + // meaningful with an explicit WithKernelProxy — the env-var path folds + // credentials into the URL and consumes NO_PROXY during resolution, so it + // leaves these empty. Empty fields are passed as NULL (kernel default). + ProxyURL string + ProxyUsername string + ProxyPassword string + ProxyBypassHosts string + + // Retry carries the driver's WithRetries backoff/attempt policy (RetryWaitMin + // / RetryWaitMax / RetryMax) plus the kernel-only overall retry budget. nil + // leaves the kernel on its own default retry policy; non-nil forwards to + // set_retry_config so the caller's policy is authoritative. A pointer so + // "unset" is distinct from an explicit zero-retry (disable) request. + Retry *RetryConfig // Location is the session time zone used to render DATE / TIMESTAMP values, // matching the Thrift path which returns them in this location. nil means UTC. @@ -52,3 +69,19 @@ type Config struct { Catalog string Schema string } + +// RetryConfig is the driver's HTTP retry policy forwarded to the kernel: the +// backoff-wait bounds, the maximum number of retries after the initial attempt +// (MaxRetries == 0 disables retries), and the cumulative retry budget across all +// attempts (OverallTimeout; zero keeps the kernel's default 900s budget). The +// connector fills MinWait/MaxWait/MaxRetries from WithRetries and OverallTimeout +// from WithKernelRetryOverallTimeout; the kernel's own retry policy applies when +// Config.Retry is nil. Maps to kernel_session_config_set_retry_config. +type RetryConfig struct { + MinWait time.Duration + MaxWait time.Duration + MaxRetries uint32 + // OverallTimeout is the cumulative retry budget; zero => keep the kernel + // default (900s). Mirrors the pyo3/napi retry_overall_timeout knob. + OverallTimeout time.Duration +} diff --git a/internal/backend/kernel/kernel_test.go b/internal/backend/kernel/kernel_test.go index 4182477c..9a175c7d 100644 --- a/internal/backend/kernel/kernel_test.go +++ b/internal/backend/kernel/kernel_test.go @@ -10,6 +10,7 @@ import ( "errors" "os" "testing" + "time" "github.com/databricks/databricks-sql-go/driverctx" dbsqlerr "github.com/databricks/databricks-sql-go/errors" @@ -76,6 +77,63 @@ func TestSetKernelTLS(t *testing.T) { } } +// TestSetProxy exercises the real kernel_session_config_set_proxy cgo setter via +// the trySetProxy seam across every field combination — proving the arg +// marshalling and the NULL-for-empty handling of the optional username / password +// / bypass_hosts args (newCStrOrNull). The "no proxy" case is a no-op (ProxyURL +// empty). A failure here means the field→setter wiring or the C signature drifted. +func TestSetProxy(t *testing.T) { + cases := []struct { + name string + cfg Config + }{ + {"url only", Config{ProxyURL: "http://proxy:3128"}}, + {"url + credentials", Config{ProxyURL: "http://proxy:3128", ProxyUsername: "u", ProxyPassword: "p"}}, + {"url + bypass", Config{ProxyURL: "http://proxy:3128", ProxyBypassHosts: "localhost,*.internal"}}, //nolint:gosec // G101: test literals, not real credentials + {"all fields", Config{ProxyURL: "http://proxy:3128", ProxyUsername: "u", ProxyPassword: "p", ProxyBypassHosts: "localhost,*.internal"}}, //nolint:gosec // G101: test literals, not real credentials + {"none (no-op)", Config{}}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + if err := trySetProxy(c.cfg); err != nil { + t.Errorf("applyProxy(%s) = %v, want nil", c.name, err) + } + }) + } +} + +// TestSetRetry exercises the real kernel_session_config_set_retry_config cgo setter +// via the trySetRetry seam: a valid range succeeds (incl. the disable form, +// MaxRetries=0, and a non-zero overall budget), and a degenerate range (min=0 or +// max= len(r.colTypes) { + return nil + } + return r.colTypes[index].ScanType +} + +// ColumnTypeDatabaseTypeName returns the Databricks type name for a column (e.g. +// "BIGINT", "DECIMAL", "ARRAY"), matching the Thrift path. An out-of-range index +// returns "". +func (r *kernelRows) ColumnTypeDatabaseTypeName(index int) string { + if index < 0 || index >= len(r.colTypes) { + return "" + } + return r.colTypes[index].DatabaseTypeName +} + +// ColumnTypeNullable reports whether a column is nullable. The kernel result +// schema does not carry a reliable per-column nullability flag, so — exactly like +// the Thrift path — this always returns ok=false (nullability unknown). +func (r *kernelRows) ColumnTypeNullable(index int) (nullable, ok bool) { + return false, false +} + +// ColumnTypeLength returns a variable-length column's length (math.MaxInt64, +// unbounded) for string/binary/nested/interval types and (0, false) for +// fixed-width types, matching the Thrift path. An out-of-range index returns +// (0, false). +func (r *kernelRows) ColumnTypeLength(index int) (length int64, ok bool) { + if index < 0 || index >= len(r.colTypes) { + return 0, false + } + ct := r.colTypes[index] + return ct.Length, ct.HasLength +} + // Close releases the current batch, the kernel result stream, and (query-path // ownership) the server operation. Idempotent. func (r *kernelRows) Close() error { diff --git a/internal/config/config.go b/internal/config/config.go index f689f414..3b705290 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -76,6 +76,37 @@ type KernelExperimentalConfig struct { // of the blanket InsecureSkipVerify (which relaxes both chain and hostname). // Maps to kernel_session_config_set_tls_skip_hostname_verification. TLSSkipHostnameVerify bool + + // ProxyURL / ProxyUsername / ProxyPassword / ProxyBypassHosts configure an + // explicit HTTP proxy on the kernel path (WithKernelProxy), overriding the + // HTTP(S)_PROXY / NO_PROXY environment the driver otherwise mirrors. The + // credentials and bypass list are the "advanced" fields the env-var path + // can't express (a structured no-proxy list, out-of-band basic auth). Empty + // ProxyURL leaves the environment-derived proxy in effect. Maps to + // kernel_session_config_set_proxy(url, username, password, bypass_hosts). + ProxyURL string + ProxyUsername string + ProxyPassword string + ProxyBypassHosts string + + // RetryOverallTimeout is the cumulative retry budget across all attempts on + // the kernel path (WithKernelRetryOverallTimeout). Zero = keep the kernel + // default (900s). WithRetries only carries the per-attempt backoff bounds + + // max attempts (mirroring the Thrift RetryWaitMin/Max/RetryMax surface); the + // overall budget is a kernel-only knob the Thrift path has no equivalent for, + // so it lives here and maps to the 4th arg of + // kernel_session_config_set_retry_config (matching the pyo3/napi + // retry_overall_timeout knob). + RetryOverallTimeout time.Duration + + // MaxChunksInMemory bounds how many decompressed CloudFetch chunks the kernel + // holds in memory at once (WithKernelMaxChunksInMemory) — the knob that trades + // throughput for peak RSS on large result sets. Zero = keep the kernel default + // (16). A positive value is forwarded as the client-only + // "cloudfetch_max_chunks_in_memory" session conf, which the kernel folds into + // its result config and strips before the SEA wire; it is kernel-only because + // the Thrift path has no such in-memory-chunk knob. + MaxChunksInMemory int } // DeepCopy returns a deep copy of the experimental config, or nil for a nil @@ -85,7 +116,15 @@ func (k *KernelExperimentalConfig) DeepCopy() *KernelExperimentalConfig { if k == nil { return nil } - cp := &KernelExperimentalConfig{TLSSkipHostnameVerify: k.TLSSkipHostnameVerify} + cp := &KernelExperimentalConfig{ + TLSSkipHostnameVerify: k.TLSSkipHostnameVerify, + ProxyURL: k.ProxyURL, + ProxyUsername: k.ProxyUsername, + ProxyPassword: k.ProxyPassword, + ProxyBypassHosts: k.ProxyBypassHosts, + RetryOverallTimeout: k.RetryOverallTimeout, + MaxChunksInMemory: k.MaxChunksInMemory, + } if k.TLSTrustedCertsPEM != nil { cp.TLSTrustedCertsPEM = append([]byte(nil), k.TLSTrustedCertsPEM...) } @@ -100,6 +139,15 @@ func (k *KernelExperimentalConfig) DeepCopy() *KernelExperimentalConfig { // the literal (see EffectiveSessionParams). const MetricViewMetadataConfKey = "spark.sql.thriftserver.metadata.metricview.enabled" +// KernelMaxChunksInMemoryConfKey is the CLIENT-only session conf the kernel reads +// to bound how many decompressed CloudFetch chunks it holds in memory +// (WithKernelMaxChunksInMemory). Unlike MetricViewMetadataConfKey this is NOT a +// server SET parameter and is NOT added by EffectiveSessionParams: the kernel +// backend injects it into its own SessionConf only, and the kernel strips it +// before the SEA wire (it is absent from the kernel's server allowlist). It must +// match the key the kernel's apply_client_result_overrides looks for. +const KernelMaxChunksInMemoryConfKey = "cloudfetch_max_chunks_in_memory" + // EffectiveSessionParams returns the session confs to send to the server: the // user-supplied SessionParams plus any conf derived from a higher-level option // (currently metric-view metadata). Both backends call this, so the derivation is diff --git a/internal/decimalfmt/decimalfmt.go b/internal/decimalfmt/decimalfmt.go index 7d21ebf5..90b2b43c 100644 --- a/internal/decimalfmt/decimalfmt.go +++ b/internal/decimalfmt/decimalfmt.go @@ -6,42 +6,105 @@ package decimalfmt import ( "math/big" - "strings" "github.com/apache/arrow/go/v12/arrow/decimal128" ) +// decMaxLen bounds the longest exact string a real Databricks DECIMAL produces, +// so ExactString's stack scratch never has to grow onto the heap. A Databricks +// DECIMAL is at most precision 38, i.e. 38 significant digits; the widest render +// is a small magnitude at a large scale — "0." + (scale) fractional digits — plus +// a sign. 48 covers scale up to 45 with the sign and "0." prefix; anything longer +// (only reachable via a synthetic decimal128 scale the server never sends) still +// renders correctly, just with one heap growth in Append. +const decMaxLen = 48 + // ExactString renders an Arrow decimal128 as an exact fixed-point string, -// applying scale by string placement rather than float conversion. This -// preserves precision a float64 would lose beyond ~17 significant digits +// applying scale by digit placement rather than float conversion. This preserves +// precision a float64 would lose beyond ~17 significant digits // (databricks-sql-go#274). A negative scale is not produced by the server and is // treated as scale 0 rather than panicking. +// +// It allocates only the returned string (the digit rendering uses a stack +// scratch, no math/big for magnitudes that fit in a uint64 — i.e. every +// DECIMAL(≤18) and most DECIMAL(≤38) values). See Append for the zero-copy form. func ExactString(n decimal128.Num, scale int32) string { - unscaled := n.BigInt() // exact signed unscaled integer - neg := unscaled.Sign() < 0 - digits := new(big.Int).Abs(unscaled).String() + var buf [decMaxLen]byte + return string(Append(buf[:0], n, scale)) +} - var b strings.Builder +// Append renders n at the given scale into dst and returns the extended slice, +// exactly as ExactString would but without forcing a string allocation — for +// callers that render many decimals into a reused buffer (e.g. a batch-scoped +// arena). The output is byte-for-byte identical to ExactString. +// +// The magnitude is rendered from the unscaled integer with no heap allocation +// when it fits in a uint64 (the common case: |unscaled| < 2^64 covers every +// DECIMAL(≤18) and every DECIMAL(≤38) value whose magnitude stays under 2^64). +// A true 128-bit magnitude (DECIMAL(20+) with a very large value) falls back to +// math/big, matching the old renderer's output exactly. +func Append(dst []byte, n decimal128.Num, scale int32) []byte { + neg := n.Sign() < 0 + abs := n if neg { - b.WriteByte('-') + // Two's-complement negation; for a positive magnitude the high word is 0 + // iff the magnitude fits in a uint64. + abs = n.Negate() + } + + if abs.HighBits() == 0 { + // Fast path: |unscaled| < 2^64. Render the digits right-to-left into a + // stack scratch (uint64 is at most 20 decimal digits), no math/big. + var tmp [20]byte + mag := abs.LowBits() + i := len(tmp) + if mag == 0 { + i-- + tmp[i] = '0' + } else { + for mag > 0 { + i-- + tmp[i] = byte('0' + mag%10) + mag /= 10 + } + } + return appendPlaced(dst, tmp[i:], neg, scale) } + // Slow path: true 128-bit magnitude. Rare (only DECIMAL values ≥ 2^64), so a + // math/big allocation here is acceptable; the digit placement is identical. + mag := new(big.Int).Abs(n.BigInt()) + var tmp [40]byte // 2^127 has 39 decimal digits + return appendPlaced(dst, mag.Append(tmp[:0], 10), neg, scale) +} + +// appendPlaced writes the sign and the magnitude digits into dst with the decimal +// point inserted `scale` places from the right. digits is the magnitude's decimal +// ASCII with no leading zeros ("0" for a zero magnitude). This encodes the exact +// placement rules both backends must agree on; the fast and slow paths share it so +// they can never diverge. +func appendPlaced(dst, digits []byte, neg bool, scale int32) []byte { + if neg { + dst = append(dst, '-') + } if scale <= 0 { - b.WriteString(digits) - return b.String() + return append(dst, digits...) } s := int(scale) - if len(digits) <= s { - // Pad with leading zeros so there are exactly `scale` fractional digits - // and a single leading integer zero, e.g. 5 with scale 3 -> "0.005". - b.WriteString("0.") - b.WriteString(strings.Repeat("0", s-len(digits))) - b.WriteString(digits) - } else { - b.WriteString(digits[:len(digits)-s]) - b.WriteByte('.') - b.WriteString(digits[len(digits)-s:]) + nd := len(digits) + if nd <= s { + // Fewer digits than the scale: a single leading integer zero, then enough + // fractional zeros to reach exactly `scale` places, e.g. 5 at scale 3 -> + // "0.005". + dst = append(dst, '0', '.') + for k := 0; k < s-nd; k++ { + dst = append(dst, '0') + } + return append(dst, digits...) } - return b.String() + // Split the digits `scale` places from the right around the point. + dst = append(dst, digits[:nd-s]...) + dst = append(dst, '.') + return append(dst, digits[nd-s:]...) } diff --git a/internal/decimalfmt/decimalfmt_test.go b/internal/decimalfmt/decimalfmt_test.go index 7765c7c6..07773e31 100644 --- a/internal/decimalfmt/decimalfmt_test.go +++ b/internal/decimalfmt/decimalfmt_test.go @@ -1,7 +1,9 @@ package decimalfmt import ( + "math" "math/big" + "strings" "testing" "github.com/apache/arrow/go/v12/arrow/decimal128" @@ -34,10 +36,171 @@ func TestExactString(t *testing.T) { // A value beyond float64's exact integer range must survive intact — the whole // point of formatting from the 128-bit unscaled integer rather than a float. func TestExactStringHighPrecision(t *testing.T) { - // 12345678901234567890 (20 digits) — past float64's 2^53 exact-integer limit. + // 12345678901234567890 (20 digits) — past float64's 2^53 exact-integer limit, + // and past uint64's max (18446744073709551615), so it also exercises the + // 128-bit slow path. big20, _ := new(big.Int).SetString("12345678901234567890", 10) got := ExactString(decimal128.FromBigInt(big20), 4) if want := "1234567890123456.7890"; got != want { t.Errorf("got %q want %q", got, want) } } + +// exactStringOracle is the pre-optimization implementation, kept verbatim as the +// correctness oracle for the alloc-free rewrite. ExactString is on the SHARED +// Thrift + kernel render path, so its output must stay byte-for-byte identical to +// this reference for every input; TestExactStringOracleParity fuzzes the two +// against each other. Do NOT "simplify" this to call ExactString — it exists +// precisely to be an independent second implementation. +func exactStringOracle(n decimal128.Num, scale int32) string { + unscaled := n.BigInt() + neg := unscaled.Sign() < 0 + digits := new(big.Int).Abs(unscaled).String() + + var b strings.Builder + if neg { + b.WriteByte('-') + } + if scale <= 0 { + b.WriteString(digits) + return b.String() + } + s := int(scale) + if len(digits) <= s { + b.WriteString("0.") + b.WriteString(strings.Repeat("0", s-len(digits))) + b.WriteString(digits) + } else { + b.WriteString(digits[:len(digits)-s]) + b.WriteByte('.') + b.WriteString(digits[len(digits)-s:]) + } + return b.String() +} + +// TestExactStringOracleParity is the byte-parity gate for the shared render path: +// it asserts the alloc-free ExactString reproduces the old implementation +// (exactStringOracle) exactly across a wide grid — both signs, the full scale +// range, the uint64/128-bit boundary, and true 128-bit magnitudes including the +// DECIMAL(38) extreme. A single divergence here would be a silent DECIMAL +// corruption in BOTH backends. +func TestExactStringOracleParity(t *testing.T) { + // Magnitudes chosen to straddle every branch: zero, small, the uint64 max + // boundary (2^64-1 and 2^64), the largest 38-digit decimal, and the in-range + // 128-bit extreme. decimal128.FromBigInt panics above 2^127-1, so that is the + // largest fixture (a bare -2^127 would panic). + mags := []string{ + "0", + "1", + "5", + "9", + "12345", + "99999999999999999", // 17 nines (fits int64) + "18446744073709551615", // 2^64 - 1 (uint64 max, last fast-path magnitude) + "18446744073709551616", // 2^64 (first 128-bit slow-path magnitude) + "99999999999999999999999999999999999999", // DECIMAL(38) all-nines max + "170141183460469231731687303715884105727", // 2^127 - 1 (in-range 128-bit extreme) + } + scales := []int32{-1, 0, 1, 2, 3, 6, 17, 18, 37, 38} + + for _, m := range mags { + bi, ok := new(big.Int).SetString(m, 10) + if !ok { + t.Fatalf("bad fixture %q", m) + } + for _, sign := range []int{1, -1} { + v := new(big.Int).Set(bi) + if sign < 0 { + if v.Sign() == 0 { + continue // -0 == 0, already covered + } + v.Neg(v) + } + num := decimal128.FromBigInt(v) + for _, sc := range scales { + got := ExactString(num, sc) + want := exactStringOracle(num, sc) + if got != want { + t.Errorf("mag=%s sign=%d scale=%d: got %q want %q", m, sign, sc, got, want) + } + } + } + } +} + +// TestAppendMatchesExactString proves the reused-buffer Append form produces the +// identical bytes to ExactString (and that appending into a non-empty buffer only +// extends it), so a caller can render many decimals into one scratch/arena buffer +// with no behavioral difference. +func TestAppendMatchesExactString(t *testing.T) { + nums := []decimal128.Num{ + decimal128.FromI64(0), + decimal128.FromI64(5), + decimal128.FromI64(-12345), + decimal128.FromI64(math.MaxInt64), + decimal128.FromU64(math.MaxUint64), + } + big128, _ := new(big.Int).SetString("123456789012345678901234567890", 10) + nums = append(nums, decimal128.FromBigInt(big128), decimal128.FromBigInt(big128.Neg(big128))) + + for _, n := range nums { + for _, sc := range []int32{0, 2, 3, 18, 30} { + want := ExactString(n, sc) + // Into an empty buffer. + if got := string(Append(nil, n, sc)); got != want { + t.Errorf("Append(nil) n=%v scale=%d: got %q want %q", n, sc, got, want) + } + // Into a pre-filled buffer: only the suffix must match, and the prefix + // must be preserved. + pre := []byte("PFX|") + out := string(Append(pre, n, sc)) + if !strings.HasPrefix(out, "PFX|") || strings.TrimPrefix(out, "PFX|") != want { + t.Errorf("Append(pre) n=%v scale=%d: got %q want prefix+%q", n, sc, out, want) + } + } + } +} + +// BenchmarkExactString / BenchmarkExactStringSink / BenchmarkAppendReused document +// the alloc profile the fix targets. Run: +// +// go test -run '^$' -bench 'ExactString|AppendReused' -benchmem ./internal/decimalfmt/ +// +// Read the two ExactString benchmarks together: the discard form (result unused) +// lets escape analysis prove the string never leaves the frame and elides its heap +// allocation, reporting 0 B/0 allocs — which flatters but does not reflect +// production, where every caller keeps the result. The Sink form assigns to a +// package-level var so the string escapes, reporting the real 1-alloc-per-cell +// profile the render path actually pays. Both are still a large win over the old +// renderer (which allocated the big.Int scratch on top of the returned string). +// AppendReused is the genuinely zero-alloc form, available to callers that render +// into a reused buffer (no production caller does today). +func BenchmarkExactString(b *testing.B) { + n := decimal128.FromI64(1999) // DECIMAL(6,2) 19.99 — the common small case + b.ReportAllocs() + for i := 0; i < b.N; i++ { + _ = ExactString(n, 2) + } +} + +// benchSink keeps the ExactString result escaping so BenchmarkExactStringSink +// measures the real production alloc profile (see the note above). +var benchSink string + +func BenchmarkExactStringSink(b *testing.B) { + n := decimal128.FromI64(1999) + b.ReportAllocs() + for i := 0; i < b.N; i++ { + benchSink = ExactString(n, 2) + } +} + +func BenchmarkAppendReused(b *testing.B) { + n := decimal128.FromI64(1999) + buf := make([]byte, 0, decMaxLen) + b.ReportAllocs() + for i := 0; i < b.N; i++ { + buf = Append(buf[:0], n, 2) + } + _ = buf +} diff --git a/internal/rows/arrowbased/arrowscan_parity_test.go b/internal/rows/arrowbased/arrowscan_parity_test.go index 82faaa13..625086cf 100644 --- a/internal/rows/arrowbased/arrowscan_parity_test.go +++ b/internal/rows/arrowbased/arrowscan_parity_test.go @@ -449,12 +449,15 @@ func TestArrowbasedKernelTimestampTZParity(t *testing.T) { } } -// TestArrowbasedKernelVariantGeometryParity documents that VARIANT and GEOMETRY -// need no special rendering: over Arrow both arrive as plain STRING columns (a -// top-level VARIANT is its JSON text "{\"a\":1}", GEOMETRY is its WKT "POINT(1 2)"), -// so the existing string arm on both backends already produces identical output. -// Verified live on both backends. (GEOGRAPHY is intentionally not covered — it is -// not enabled on the benchmark warehouse and no consumer has asked for it.) +// TestArrowbasedKernelVariantGeometryParity documents that VARIANT, GEOMETRY, and +// GEOGRAPHY need no special rendering: over Arrow all three arrive as plain STRING +// columns (a top-level VARIANT is its JSON text "{\"a\":1}", GEOMETRY/GEOGRAPHY are +// their WKT "POINT(1 2)"), so the existing string arm on both backends already +// produces identical output. The kernel maps GEOMETRY and GEOGRAPHY to the same +// Arrow Utf8 shape (distinguished only by the databricks.type_name field-metadata +// hint, which does not change the scanned value — see the kernel post-processor, +// where both hit the same WKT arm), so both render identically to the Thrift path. +// Verified live on both backends. func TestArrowbasedKernelVariantGeometryParity(t *testing.T) { pool := memory.NewGoAllocator() loc := time.UTC @@ -464,6 +467,9 @@ func TestArrowbasedKernelVariantGeometryParity(t *testing.T) { {"variant_object", `{"a":1,"b":[2,3]}`}, {"variant_scalar", "42"}, {"geometry_wkt", "POINT(1 2)"}, + // GEOGRAPHY is the sibling geospatial type; like GEOMETRY it comes back as + // WKT text in a Utf8 column, so the string arm renders it identically. + {"geography_wkt", "POINT(1 2)"}, } for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { diff --git a/internal/rows/coltype_thrift_parity_test.go b/internal/rows/coltype_thrift_parity_test.go new file mode 100644 index 00000000..be178a64 --- /dev/null +++ b/internal/rows/coltype_thrift_parity_test.go @@ -0,0 +1,105 @@ +package rows + +import ( + "testing" + + "github.com/apache/arrow/go/v12/arrow" + + "github.com/databricks/databricks-sql-go/internal/arrowscan" + "github.com/databricks/databricks-sql-go/internal/cli_service" + "github.com/databricks/databricks-sql-go/internal/rows/rowscanner" +) + +// TestColumnTypeInfoMatchesThriftMapping is the pure-Go drift guard the kernel +// backend's column-metadata mapping (internal/arrowscan.ColumnTypeInfoFor) was +// missing. That mapping restates the Thrift path's per-type decisions — getScanType +// (this package), rowscanner.GetDBTypeName, and ColumnTypeLength — but pinned its own +// hardcoded expectations, so a change to the Thrift mapping (e.g. DECIMAL's scan +// type) would leave arrowscan silently divergent while CI stayed green; parity was +// then enforced only by the warehouse-gated, build-tagged TestKernelThriftColumnTypeParity +// (nightly, not PR CI). +// +// This test cross-checks the two mappings DIRECTLY, with no cgo and no warehouse, so +// drift fails a default CGO_ENABLED=0 PR run. The two switch on different type systems +// (Arrow IDs vs Thrift TTypeIds), so it can't share a map — instead it enumerates the +// shared Databricks types as {Arrow type, equivalent Thrift TColumnDesc} pairs and +// asserts ColumnTypeInfoFor(arrow) reports exactly what the Thrift functions produce +// for the paired column. +// +// It lives in package rows (white-box) to reach the unexported getScanType; +// arrowscan is a pure leaf import (no cycle). +func TestColumnTypeInfoMatchesThriftMapping(t *testing.T) { + // thriftCol builds the minimal TColumnDesc the Thrift mapping functions read. + thriftCol := func(id cli_service.TTypeId) *cli_service.TColumnDesc { + return &cli_service.TColumnDesc{ + TypeDesc: &cli_service.TTypeDesc{ + Types: []*cli_service.TTypeEntry{ + {PrimitiveEntry: &cli_service.TPrimitiveTypeEntry{Type: id}}, + }, + }, + } + } + + // Each shared Databricks type as it arrives on each backend: the Arrow type the + // kernel receives, and the Thrift TTypeId the server declares for the SAME logical + // column. ColumnTypeInfoFor(arrow) must match the Thrift functions on thriftID. + cases := []struct { + name string + arrow arrow.DataType + thriftID cli_service.TTypeId + }{ + {"boolean", arrow.FixedWidthTypes.Boolean, cli_service.TTypeId_BOOLEAN_TYPE}, + {"tinyint", arrow.PrimitiveTypes.Int8, cli_service.TTypeId_TINYINT_TYPE}, + {"smallint", arrow.PrimitiveTypes.Int16, cli_service.TTypeId_SMALLINT_TYPE}, + {"int", arrow.PrimitiveTypes.Int32, cli_service.TTypeId_INT_TYPE}, + {"bigint", arrow.PrimitiveTypes.Int64, cli_service.TTypeId_BIGINT_TYPE}, + {"float", arrow.PrimitiveTypes.Float32, cli_service.TTypeId_FLOAT_TYPE}, + {"double", arrow.PrimitiveTypes.Float64, cli_service.TTypeId_DOUBLE_TYPE}, + {"string", arrow.BinaryTypes.String, cli_service.TTypeId_STRING_TYPE}, + {"binary", arrow.BinaryTypes.Binary, cli_service.TTypeId_BINARY_TYPE}, + {"date", arrow.FixedWidthTypes.Date32, cli_service.TTypeId_DATE_TYPE}, + {"timestamp", arrow.FixedWidthTypes.Timestamp_us, cli_service.TTypeId_TIMESTAMP_TYPE}, + {"decimal", &arrow.Decimal128Type{Precision: 10, Scale: 2}, cli_service.TTypeId_DECIMAL_TYPE}, + {"array", arrow.ListOf(arrow.PrimitiveTypes.Int64), cli_service.TTypeId_ARRAY_TYPE}, + {"map", arrow.MapOf(arrow.BinaryTypes.String, arrow.PrimitiveTypes.Int64), cli_service.TTypeId_MAP_TYPE}, + {"struct", arrow.StructOf(arrow.Field{Name: "x", Type: arrow.PrimitiveTypes.Int64}), cli_service.TTypeId_STRUCT_TYPE}, + // Interval types: in the prod default (native-interval Arrow off) the server + // pre-formats intervals to text and declares the column STRING_TYPE, so the + // kernel — which receives native arrow.DURATION / arrow.INTERVAL_MONTHS and + // formats them Go-side to the same string — must report the STRING metadata the + // Thrift path reports for STRING_TYPE. Pairing them with STRING_TYPE here pins + // exactly that (see the arrowscan DURATION / INTERVAL_MONTHS arms). + {"interval_day_time", arrow.FixedWidthTypes.Duration_us, cli_service.TTypeId_STRING_TYPE}, + {"interval_year_month", arrow.FixedWidthTypes.MonthInterval, cli_service.TTypeId_STRING_TYPE}, + {"null", arrow.Null, cli_service.TTypeId_NULL_TYPE}, + } + + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + col := thriftCol(c.thriftID) + info := arrowscan.ColumnTypeInfoFor(c.arrow) + + // ScanType: ColumnTypeInfoFor must recommend the same Go type as the Thrift + // getScanType. This is the DECIMAL-drift guard the finding named. + if want := getScanType(col); info.ScanType != want { + t.Errorf("ScanType = %v, Thrift getScanType = %v", info.ScanType, want) + } + + // DatabaseTypeName: the Thrift name is the enum name minus the _TYPE suffix + // (e.g. BIGINT, DECIMAL, ARRAY), which ColumnTypeInfoFor reports verbatim. + if want := rowscanner.GetDBTypeName(col); info.DatabaseTypeName != want { + t.Errorf("DatabaseTypeName = %q, Thrift GetDBTypeName = %q", info.DatabaseTypeName, want) + } + + // Length: only variable-length types report an (unbounded) length. Drive the + // expectation from the SAME classifier ColumnTypeLength uses + // (columnTypeLengthForID), not a hand-copied switch, so a change to the + // Thrift-side length rule fails this test instead of the two copies drifting + // together and staying green. + wantLen, wantOk := columnTypeLengthForID(c.thriftID) + if info.Length != wantLen || info.HasLength != wantOk { + t.Errorf("Length = (%d,%v), Thrift = (%d,%v)", info.Length, info.HasLength, wantLen, wantOk) + } + }) + } +} diff --git a/internal/rows/rows.go b/internal/rows/rows.go index cbbefa33..a87e0ceb 100644 --- a/internal/rows/rows.go +++ b/internal/rows/rows.go @@ -381,10 +381,21 @@ func (r *rows) ColumnTypeLength(index int) (length int64, ok bool) { if err != nil { return 0, false } + return columnTypeLengthForID(rowscanner.GetDBTypeID(columnInfo)) +} - typeName := rowscanner.GetDBTypeID(columnInfo) - - switch typeName { +// columnTypeLengthForID is the variable-length classification behind +// ColumnTypeLength, factored out of the method (which needs a live *rows and a +// server round-trip to reach it) so the exact rule can be reused verbatim — in +// particular by the pure-Go coltype parity guard (TestColumnTypeInfoMatchesThriftMapping), +// which asserts the kernel path's arrowscan mapping against this same classifier +// rather than a hand-copied replica that could drift with it. +// +// Variable-length types (string/varchar/binary and the nested array/map/struct) +// report an unbounded length (math.MaxInt64); every fixed-width type reports +// (0, false). +func columnTypeLengthForID(id cli_service.TTypeId) (length int64, ok bool) { + switch id { case cli_service.TTypeId_STRING_TYPE, cli_service.TTypeId_VARCHAR_TYPE, cli_service.TTypeId_BINARY_TYPE, diff --git a/kernel_backend.go b/kernel_backend.go index e956058f..35800d64 100644 --- a/kernel_backend.go +++ b/kernel_backend.go @@ -30,10 +30,14 @@ func newKernelBackend(_ context.Context, cfg *config.Config) (backend.Backend, e // buildKernelConfig in kernel_config.go, unit-tested under CGO_ENABLED=0. // SPOG org routing rides in HTTPPath's ?o= and is parsed kernel-side. kc := buildKernelConfig(cfg, kauth) - // Proxy: the Thrift path uses http.ProxyFromEnvironment; mirror it by reading - // the same HTTP(S)_PROXY / NO_PROXY environment for the kernel. Resolved here - // (not in buildKernelConfig) since proxyForEndpoint has its own unit coverage. - kc.ProxyURL = proxyForEndpoint(cfg) + // Proxy: an explicit WithKernelProxy (carried on KernelExperimental) takes + // precedence over the environment — an explicit proxy is a deliberate + // override, and consulting both would be ambiguous. Otherwise mirror the + // Thrift path, which uses http.ProxyFromEnvironment, by reading the same + // HTTP(S)_PROXY / NO_PROXY environment for the kernel. resolveKernelProxy is + // pure Go (untagged, in kernel_config.go) with its own unit coverage, so it + // (and proxyForEndpoint) resolve here rather than in buildKernelConfig. + resolveKernelProxy(cfg, &kc) return kernel.New(kc), nil } diff --git a/kernel_config.go b/kernel_config.go index ecf4cc84..e4bf996f 100644 --- a/kernel_config.go +++ b/kernel_config.go @@ -3,6 +3,9 @@ package dbsql import ( "errors" "fmt" + "net/url" + "strconv" + "time" "github.com/databricks/databricks-sql-go/auth/noop" "github.com/databricks/databricks-sql-go/auth/pat" @@ -78,13 +81,23 @@ func validateKernelConfig(cfg *config.Config) (kernel.Auth, error) { return kernel.Auth{}, fmt.Errorf("databricks: WithTimeout (server query timeout) is %w; "+ "omit it or use the default (Thrift) backend", dbsqlerr.ErrNotSupportedByKernel) } - // WithRetries(-1) explicitly disables retries, but the kernel retries - // internally below the C ABI with no user-facing toggle — so a disable request - // would be silently violated. Reject it. Positive/default RetryMax is fine: the - // kernel provides retries (just not user-tunable), documented in doc.go. - if cfg.RetryMax < 0 { - return kernel.Auth{}, fmt.Errorf("databricks: disabling retries via WithRetries is %w "+ - "(the kernel retries internally); omit it or use the default (Thrift) backend", dbsqlerr.ErrNotSupportedByKernel) + // WithRetries (RetryWaitMin / RetryWaitMax / RetryMax) is honored on the kernel + // path: newKernelBackend forwards it via kernelRetryConfig → + // kernel_session_config_set_retry_config, and a negative RetryMax (the disable + // form) maps to zero kernel retries. So it is neither rejected nor silently + // ignored here. + // WithKernelProxy URL: validate it in the Go layer so a malformed URL is a clear, + // errors.Is-able config error (ErrInvalidKernelConfig) here rather than an opaque + // "kernel: set_proxy: …" wrap from the C ABI at connect. url.Parse is lenient, so + // also require a scheme and host — the shape kernel_session_config_set_proxy needs. + if ke := cfg.KernelExperimental; ke != nil && ke.ProxyURL != "" { + if u, perr := url.Parse(ke.ProxyURL); perr != nil { + return kernel.Auth{}, fmt.Errorf("databricks: the WithKernelProxy URL %q is %w: %v", + ke.ProxyURL, dbsqlerr.ErrInvalidKernelConfig, perr) + } else if u.Scheme == "" || u.Host == "" { + return kernel.Auth{}, fmt.Errorf("databricks: the WithKernelProxy URL %q is %w "+ + "(want a scheme and host, e.g. http://proxy:3128)", ke.ProxyURL, dbsqlerr.ErrInvalidKernelConfig) + } } return kauth, nil } @@ -124,10 +137,144 @@ func buildKernelConfig(cfg *config.Config, kauth kernel.Auth) kernel.Config { if ke := cfg.KernelExperimental; ke != nil { kc.TLSTrustedCertsPEM = ke.TLSTrustedCertsPEM kc.TLSSkipHostnameVerify = ke.TLSSkipHostnameVerify + // Kernel-only CloudFetch in-memory-chunk knob (WithKernelMaxChunksInMemory). + // Injected into the kernel backend's own SessionConf ONLY (not via + // EffectiveSessionParams, which both backends share) as the client-only key + // the kernel reads and strips before the SEA wire — so it never leaks to the + // server or the Thrift path. Zero/negative keeps the kernel default (16). + if ke.MaxChunksInMemory > 0 { + // EffectiveSessionParams returned a fresh map, so mutating it is safe. + kc.SessionConf[config.KernelMaxChunksInMemoryConfKey] = strconv.Itoa(ke.MaxChunksInMemory) + } } + // Retry / backoff policy from WithRetries (+ the kernel-only overall budget). + // nil leaves the kernel's own default policy in place. + kc.Retry = kernelRetryConfig(cfg) return kc } +// 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 +// to zero. The kernel setter validates the range (it rejects min == 0 / max < min), +// so a valid one must be passed even when the attempts make the backoff moot; any +// positive min<=max works, and the kernel's own defaults (1s / 60s) are the natural +// choice. +const ( + kernelRetryPlaceholderWaitMin = 1 * time.Second + kernelRetryPlaceholderWaitMax = 60 * time.Second + + // kernelRetryMinWaitFloor is the smallest wait the kernel setter accepts once the + // driver forwards it in milliseconds (kernel_session_config_set_retry_config rejects + // min_wait_ms == 0). A positive wait below this floor would truncate to 0ms and fail + // the connect, so kernelRetryConfig clamps up to it. + kernelRetryMinWaitFloor = 1 * time.Millisecond +) + +// kernelRetryConfig resolves the driver's WithRetries policy (RetryWaitMin / +// RetryWaitMax / RetryMax) into the kernel retry descriptor so the caller's +// backoff/attempt policy is authoritative on the kernel path — matching the Thrift +// path, which applies exactly these values via go-retryablehttp, and the same +// "retries after the initial attempt" semantics the kernel setter uses. The +// connector's WithDefaults guarantees positive waits (1s / 30s) and RetryMax 4. +// +// A negative RetryMax is the WithRetries disable form (retryablehttp treats it as +// zero retries); it maps to MaxRetries == 0 and is honored EVEN when the waits are +// zero (WithRetries(-1, 0, 0) is the idiomatic disable), by substituting a valid +// placeholder range the setter accepts — the backoff is unused with no retries. +// +// A positive RetryMax with a degenerate wait range is honored the same way: the +// caller's attempt count is authoritative and the placeholder waits are substituted +// so the setter accepts the range. This case is reachable on the NORMAL option path, +// not just from a hand-built Config: WithDefaults() runs before options, so +// WithRetries(n, 0, 0) — valid per its own godoc, which promises sane wait defaults — +// overwrites the waits back to zero and lands here with the caller's RetryMax. Without +// this the caller's RetryMax would be silently dropped to the kernel's default policy. +// +// Returns nil — leaving the kernel's own default policy in place — only when a +// degenerate range carries NEITHER a caller attempt count (RetryMax > 0) NOR an +// explicit overall budget (WithKernelRetryOverallTimeout): the zero-value signature +// of a Config assembled without WithDefaults (unusual outside tests), where there is +// nothing to preserve. Substituting placeholders there would force zero retries onto +// a hand-built config; the kernel default is the safer choice. When only the overall +// budget is set, it is still forwarded (placeholder waits + zero retries) rather than +// dropped. +// +// Sub-millisecond waits are clamped up to a 1ms floor before they are returned: +// applyRetry forwards RetryConfig waits via time.Duration.Milliseconds(), so a valid +// wait in (0, 1ms) would truncate to 0ms, which the kernel setter rejects +// (min_wait_ms must be > 0) — a connect failure the Thrift path (go-retryablehttp, +// which accepts any Duration) does not have. +func kernelRetryConfig(cfg *config.Config) *kernel.RetryConfig { + // Kernel-only overall retry budget (WithKernelRetryOverallTimeout), carried on + // KernelExperimental rather than WithRetries. Zero = keep the kernel default; + // only a positive value overrides it. + var overall time.Duration + if ke := cfg.KernelExperimental; ke != nil && ke.RetryOverallTimeout > 0 { + overall = ke.RetryOverallTimeout + } + + // Disable form: honor it regardless of the (often zero) waits. Substitute a + // valid placeholder range so the setter accepts it; with 0 retries it's unused. + if cfg.RetryMax < 0 { + return &kernel.RetryConfig{ + MinWait: kernelRetryPlaceholderWaitMin, + MaxWait: kernelRetryPlaceholderWaitMax, + MaxRetries: 0, + OverallTimeout: overall, + } + } + + // Degenerate wait range (WithRetries(n, 0, 0), or a Config assembled without + // WithDefaults). If the caller asked for a positive attempt count OR an explicit + // overall budget, honor it with the placeholder waits — dropping it to the kernel + // default would silently ignore the caller's RetryMax / RetryOverallTimeout. Only + // when there is nothing to preserve (no attempt count and no overall budget) do we + // fall back to the kernel's default policy. + minWait, maxWait := cfg.RetryWaitMin, cfg.RetryWaitMax + if minWait <= 0 || maxWait < minWait { + if cfg.RetryMax <= 0 && overall <= 0 { + return nil + } + minWait, maxWait = kernelRetryPlaceholderWaitMin, kernelRetryPlaceholderWaitMax + } + // Clamp waits up to a 1ms floor: applyRetry forwards them via + // time.Duration.Milliseconds(), so a valid sub-ms wait would truncate to 0ms and + // the setter would reject the connect. maxWait is floored too so it stays >= minWait. + if minWait < kernelRetryMinWaitFloor { + minWait = kernelRetryMinWaitFloor + } + if maxWait < minWait { + maxWait = minWait + } + return &kernel.RetryConfig{ + MinWait: minWait, + MaxWait: maxWait, + MaxRetries: uint32(cfg.RetryMax), //nolint:gosec // RetryMax >= 0 here (negative handled above) + OverallTimeout: overall, + } +} + +// resolveKernelProxy fills the kernel Config's proxy fields. An explicit +// WithKernelProxy (KernelExperimental.ProxyURL non-empty) wins verbatim, +// including its out-of-band credentials and bypass list; otherwise the +// endpoint's environment-derived proxy URL is used (with no credentials or +// bypass list — the env path folds credentials into the URL userinfo and +// consumes NO_PROXY during resolution, per Go's proxy-env convention). Kept +// untagged here (like buildKernelConfig) so the explicit-over-env precedence is +// asserted under CGO_ENABLED=0 (see TestResolveKernelProxy); newKernelBackend +// calls it after buildKernelConfig so it stays a thin assembler. +func resolveKernelProxy(cfg *config.Config, kc *kernel.Config) { + if ke := cfg.KernelExperimental; ke != nil && ke.ProxyURL != "" { + kc.ProxyURL = ke.ProxyURL + kc.ProxyUsername = ke.ProxyUsername + kc.ProxyPassword = ke.ProxyPassword + kc.ProxyBypassHosts = ke.ProxyBypassHosts + return + } + kc.ProxyURL = proxyForEndpoint(cfg) +} + // resolveKernelAuth picks the kernel auth form from the config. The kernel backend // drives the kernel's own OAuth flow from raw credentials rather than reusing the Go // authenticator's Authenticate method. It reads those credentials off diff --git a/kernel_config_test.go b/kernel_config_test.go index 4da8efcb..c152eeab 100644 --- a/kernel_config_test.go +++ b/kernel_config_test.go @@ -108,7 +108,6 @@ func TestValidateKernelConfig(t *testing.T) { mut func(*config.Config) }{ {"query timeout", func(c *config.Config) { c.QueryTimeout = 30 * time.Second }}, - {"disable retries", func(c *config.Config) { c.RetryMax = -1 }}, {"non-https protocol", func(c *config.Config) { c.Protocol = "http" }}, {"non-default port", func(c *config.Config) { c.Port = 8443 }}, {"custom transport", func(c *config.Config) { c.Transport = http.DefaultTransport }}, @@ -239,12 +238,52 @@ func TestValidateKernelConfig(t *testing.T) { } }) + t.Run("disable retries (WithRetries(-1)) accepted", func(t *testing.T) { + // The disable form (RetryMax < 0) maps to zero kernel retries via the + // retry-config setter, so it validates rather than erroring. + c := baseKernelConfig() + c.RetryMax = -1 + if _, err := validateKernelConfig(c); err != nil { + t.Errorf("WithRetries(-1) (disable) should validate, got %v", err) + } + }) + t.Run("default https/443 accepted", func(t *testing.T) { c := baseKernelConfig() // WithDefaults sets Protocol=https, Port=443 if _, err := validateKernelConfig(c); err != nil { t.Errorf("the default https/443 endpoint should validate, got %v", err) } }) + + t.Run("valid WithKernelProxy URL accepted", func(t *testing.T) { + c := baseKernelConfig() + WithKernelProxy(KernelProxy{URL: "http://proxy.internal:3128", Username: "u", Password: "p", BypassHosts: "*.internal"})(c) + if _, err := validateKernelConfig(c); err != nil { + t.Errorf("a well-formed proxy URL should validate, got %v", err) + } + }) + + t.Run("malformed WithKernelProxy URL rejected as ErrInvalidKernelConfig", func(t *testing.T) { + // A malformed URL must be caught in the Go layer with an errors.Is-able + // config error, not surface as an opaque "kernel: set_proxy: …" wrap at + // connect (or, worse, a URL missing a scheme/host that the C ABI can't use). + for _, tc := range []struct { + name, proxyURL string + }{ + {"control chars", "http://a\x7f:3128"}, // url.Parse returns an error + {"no scheme or host", "proxy:3128"}, // parses, but unusable shape + {"scheme only", "http://"}, // parses, empty host + } { + t.Run(tc.name, func(t *testing.T) { + c := baseKernelConfig() + WithKernelProxy(KernelProxy{URL: tc.proxyURL})(c) + _, err := validateKernelConfig(c) + if !errors.Is(err, dbsqlerr.ErrInvalidKernelConfig) { + t.Errorf("proxy URL %q should be rejected as ErrInvalidKernelConfig, got %v", tc.proxyURL, err) + } + }) + } + }) } // kernelConfigFieldDisposition records, for every UserConfig field, how the kernel @@ -275,16 +314,19 @@ var kernelConfigFieldDisposition = map[string]string{ // Rejected loudly by validateKernelConfig. "QueryTimeout": "rejected", // when > 0 (WithTimeout) - "RetryMax": "rejected", // when < 0 (disable retries) "Protocol": "rejected", // kernel is https-only; non-default rejected "Port": "rejected", // kernel connects on 443; non-default rejected "Transport": "rejected", // custom RoundTripper; kernel uses its own HTTP stack, so reject rather than drop + // Forwarded to the kernel's HTTP retry config via kernelRetryConfig → + // set_retry_config (WithRetries: backoff bounds + max attempts, incl. disable). + "RetryMax": "forwarded", + "RetryWaitMin": "forwarded", + "RetryWaitMax": "forwarded", + // Accepted but intentionally inert on the kernel path (documented in doc.go): // the kernel manages these internally, below the C ABI, with no user knob. "MaxRows": "inert", - "RetryWaitMin": "inert", - "RetryWaitMax": "inert", "UseLz4Compression": "inert", // kernel negotiates compression internally // Not applicable to the kernel path (Thrift/HTTP-transport or telemetry knobs @@ -400,4 +442,199 @@ func TestBuildKernelConfig(t *testing.T) { t.Errorf("auth not forwarded: got %+v, want %+v", kc.Auth, kauth) } }) + + t.Run("MaxChunksInMemory injected into kernel SessionConf", func(t *testing.T) { + c := baseKernelConfig() + c.KernelExperimental = &config.KernelExperimentalConfig{MaxChunksInMemory: 4} + kc := buildKernelConfig(c, kernel.Auth{Mode: kernel.AuthPAT, Token: "dapi-x"}) + if got := kc.SessionConf[config.KernelMaxChunksInMemoryConfKey]; got != "4" { + t.Errorf("SessionConf[%q] = %q, want %q", config.KernelMaxChunksInMemoryConfKey, got, "4") + } + }) + + t.Run("MaxChunksInMemory unset leaves the key absent (kernel default)", func(t *testing.T) { + c := baseKernelConfig() // KernelExperimental nil + kc := buildKernelConfig(c, kernel.Auth{Mode: kernel.AuthPAT, Token: "dapi-x"}) + if _, ok := kc.SessionConf[config.KernelMaxChunksInMemoryConfKey]; ok { + t.Errorf("SessionConf should not carry %q when unset (kernel keeps its default)", config.KernelMaxChunksInMemoryConfKey) + } + // A zero/negative explicit value is also a no-op. + c.KernelExperimental = &config.KernelExperimentalConfig{MaxChunksInMemory: 0} + kc = buildKernelConfig(c, kernel.Auth{Mode: kernel.AuthPAT, Token: "dapi-x"}) + if _, ok := kc.SessionConf[config.KernelMaxChunksInMemoryConfKey]; ok { + t.Errorf("SessionConf should not carry %q for a zero value", config.KernelMaxChunksInMemoryConfKey) + } + }) + + t.Run("MaxChunksInMemory is not leaked to the shared server params", func(t *testing.T) { + // The knob is kernel-only: it must NOT appear in EffectiveSessionParams + // (which the Thrift path also sends to the server). + c := baseKernelConfig() + c.KernelExperimental = &config.KernelExperimentalConfig{MaxChunksInMemory: 4} + if _, ok := c.EffectiveSessionParams()[config.KernelMaxChunksInMemoryConfKey]; ok { + t.Errorf("%q must not be in EffectiveSessionParams (kernel-only, not a server param)", config.KernelMaxChunksInMemoryConfKey) + } + }) + + t.Run("MaxChunksInMemory conf key matches the kernel's cross-repo contract", func(t *testing.T) { + // Unlike the sibling retry knob (a typed C setter whose signature drift is a + // link/compile error), this knob rides a stringly-typed session-conf key whose + // only consumer is the kernel's apply_client_result_overrides + // (CLIENT_CONF_CLOUDFETCH_MAX_CHUNKS in databricks-sql-kernel src/session.rs). + // There is no build-time coupling, so pin the exact literal here: an accidental + // edit to the Go constant fails in CGO_ENABLED=0 PR CI (not just the + // warehouse-gated nightly), and this test is the greppable anchor a kernel-side + // rename must update in lockstep. + const wantKey = "cloudfetch_max_chunks_in_memory" + if config.KernelMaxChunksInMemoryConfKey != wantKey { + t.Errorf("KernelMaxChunksInMemoryConfKey = %q, want %q — the kernel's "+ + "apply_client_result_overrides reads this exact key; a mismatch silently "+ + "no-ops the knob or leaks it to the server", + config.KernelMaxChunksInMemoryConfKey, wantKey) + } + }) +} + +// TestKernelRetryConfig covers the pure resolution of the driver's WithRetries +// policy into the kernel retry descriptor: the defaults forward the connector's +// positive backoff bounds + max attempts; the disable form maps to zero retries; +// WithRetries(n, 0, 0) still forwards the caller's attempt count (with placeholder +// waits) rather than silently dropping to the kernel default; a fully zero-value +// range (a Config without WithDefaults and no attempt count) returns nil so a stray +// zero can't fail the connect; and the kernel-only overall-timeout knob is read from +// KernelExperimental. Runs in the default CGO_ENABLED=0 build. +func TestKernelRetryConfig(t *testing.T) { + t.Run("defaults forward backoff + max attempts, no overall budget", func(t *testing.T) { + c := baseKernelConfig() // WithDefaults: RetryMax=4, RetryWaitMin=1s, RetryWaitMax=30s + r := kernelRetryConfig(c) + if r == nil { + t.Fatal("kernelRetryConfig returned nil for the default (positive) range") + } + if r.MinWait != time.Second || r.MaxWait != 30*time.Second || r.MaxRetries != 4 { + t.Errorf("resolved retry = %+v, want {1s, 30s, 4}", r) + } + if r.OverallTimeout != 0 { + t.Errorf("OverallTimeout = %v, want 0 (keep kernel default) when unset", r.OverallTimeout) + } + }) + + t.Run("disable form maps to zero retries", func(t *testing.T) { + c := baseKernelConfig() + c.RetryMax = -1 // WithRetries(-1) disable + r := kernelRetryConfig(c) + if r == nil { + t.Fatal("disable form should still forward a config (0 retries), got nil") + } + if r.MaxRetries != 0 { + t.Errorf("MaxRetries = %d, want 0 (disable maps to zero kernel retries)", r.MaxRetries) + } + }) + + t.Run("idiomatic disable WithRetries(-1,0,0) honored despite zero waits", func(t *testing.T) { + // The idiomatic disable zeroes the waits too. The resolver must still honor + // the disable (MaxRetries=0) and substitute a valid placeholder range so the + // kernel setter accepts it (it rejects min==0), rather than returning nil + // (which would leave the kernel's DEFAULT retry policy in place — not disabled). + c := baseKernelConfig() + c.RetryMax = -1 + c.RetryWaitMin = 0 + c.RetryWaitMax = 0 + r := kernelRetryConfig(c) + if r == nil { + t.Fatal("WithRetries(-1,0,0) must forward a disable config, got nil (kernel default would apply — not disabled)") + } + if r.MaxRetries != 0 { + t.Errorf("MaxRetries = %d, want 0", r.MaxRetries) + } + if r.MinWait <= 0 || r.MaxWait < r.MinWait { + t.Errorf("placeholder waits = {%v, %v}, want a valid range the kernel setter accepts", r.MinWait, r.MaxWait) + } + }) + + t.Run("WithRetries(n,0,0) honors the attempt count despite zero waits", func(t *testing.T) { + // The regression this guards: WithDefaults() runs before options, so + // WithRetries(10, 0, 0) — valid per its godoc, which promises sane wait + // defaults — overwrites the waits to zero. The resolver must still forward the + // caller's RetryMax (with placeholder waits), not return nil (which would drop + // to the kernel's default policy and silently ignore the requested attempts). + c := baseKernelConfig() + WithRetries(10, 0, 0)(c) + r := kernelRetryConfig(c) + if r == nil { + t.Fatal("WithRetries(10,0,0) must forward the caller's RetryMax, got nil (kernel default would apply)") + } + if r.MaxRetries != 10 { + t.Errorf("MaxRetries = %d, want 10 (caller's attempt count honored)", r.MaxRetries) + } + if r.MinWait <= 0 || r.MaxWait < r.MinWait { + t.Errorf("placeholder waits = {%v, %v}, want a valid range the kernel setter accepts", r.MinWait, r.MaxWait) + } + }) + + t.Run("overall timeout forwarded from KernelExperimental", func(t *testing.T) { + c := baseKernelConfig() + WithKernelRetryOverallTimeout(5 * time.Minute)(c) + r := kernelRetryConfig(c) + if r == nil || r.OverallTimeout != 5*time.Minute { + t.Errorf("OverallTimeout not forwarded: %+v", r) + } + }) + + t.Run("zero-value range with no attempt count returns nil (keep kernel default)", func(t *testing.T) { + // A Config assembled without WithDefaults and no RetryMax: zero waits are a + // nonsense range the kernel setter would reject and there is no caller attempt + // count to preserve, so resolve to nil rather than fail connect. + c := &config.Config{UserConfig: config.UserConfig{RetryWaitMin: 0, RetryWaitMax: 0}} + if r := kernelRetryConfig(c); r != nil { + t.Errorf("zero-value range with RetryMax 0 should return nil, got %+v", r) + } + }) + + t.Run("overall budget forwarded even when retries are zeroed", func(t *testing.T) { + // The regression this guards: WithRetries(0,0,0) leaves RetryMax==0 and a + // degenerate wait range, but an explicit WithKernelRetryOverallTimeout must + // still reach the kernel. Returning nil here (the old behavior) discarded the + // caller's overall budget and left the kernel's 900s default — the exact + // "silently ignored option" failure the kernel gate was built to prevent. + c := baseKernelConfig() + WithRetries(0, 0, 0)(c) + WithKernelRetryOverallTimeout(5 * time.Minute)(c) + r := kernelRetryConfig(c) + if r == nil { + t.Fatal("overall budget with zeroed retries must forward a config, got nil (5m budget dropped, kernel default would apply)") + } + if r.OverallTimeout != 5*time.Minute { + t.Errorf("OverallTimeout = %v, want 5m (forwarded despite zero retries)", r.OverallTimeout) + } + if r.MaxRetries != 0 { + t.Errorf("MaxRetries = %d, want 0", r.MaxRetries) + } + if r.MinWait <= 0 || r.MaxWait < r.MinWait { + t.Errorf("placeholder waits = {%v, %v}, want a valid range the kernel setter accepts", r.MinWait, r.MaxWait) + } + }) + + t.Run("sub-millisecond waits are clamped up to a 1ms floor", func(t *testing.T) { + // The regression this guards: kernelRetryConfig validates waits in Duration + // space, but applyRetry forwards them via time.Duration.Milliseconds(). A valid + // wait in (0, 1ms) passes the guard yet truncates to min_wait_ms == 0, which the + // kernel setter rejects (InvalidArgument) — a connect failure the Thrift path, + // which accepts any Duration, does not have. The resolver must clamp up so the + // forwarded millisecond value stays > 0. + c := baseKernelConfig() + WithRetries(5, 999*time.Microsecond, 30*time.Second)(c) + r := kernelRetryConfig(c) + if r == nil { + t.Fatal("sub-ms MinWait should still resolve a config, got nil") + } + if r.MinWait.Milliseconds() < 1 { + t.Errorf("MinWait = %v (%d ms), want clamped to >= 1ms so the kernel setter accepts it", r.MinWait, r.MinWait.Milliseconds()) + } + if r.MaxWait < r.MinWait { + t.Errorf("MaxWait = %v < MinWait = %v after clamp", r.MaxWait, r.MinWait) + } + if r.MaxRetries != 5 { + t.Errorf("MaxRetries = %d, want 5 (attempt count preserved)", r.MaxRetries) + } + }) } diff --git a/kernel_e2e_test.go b/kernel_e2e_test.go index cc92d091..1abaea10 100644 --- a/kernel_e2e_test.go +++ b/kernel_e2e_test.go @@ -180,6 +180,12 @@ func TestKernelE2EDataTypes(t *testing.T) { db := kernelTestDB(t) defer db.Close() + // skipOnErrCases lists subtests to skip (rather than fail) when the server + // rejects the expression — for a type that may be gated on some warehouses + // (GEOGRAPHY is not enabled everywhere). The value comes back identically to + // GEOMETRY when it IS enabled, which is what the case documents. + skipOnErrCases := map[string]bool{"geography": true} + cases := []struct { name string expr string // the single SELECT expression @@ -204,13 +210,19 @@ func TestKernelE2EDataTypes(t *testing.T) { {"date", "CAST('2026-07-09' AS DATE)", time.Date(2026, time.July, 9, 0, 0, 0, 0, time.UTC)}, {"timestamp", "CAST('2026-07-09 12:34:56' AS TIMESTAMP)", time.Date(2026, time.July, 9, 12, 34, 56, 0, time.UTC)}, {"null", "CAST(NULL AS STRING)", nil}, - // Nested types render to a JSON string; VARIANT arrives nested, GEOMETRY - // as a WKT/WKB string. + // Nested types render to a JSON string; VARIANT arrives nested, + // GEOMETRY/GEOGRAPHY as a WKT/WKB string. {"array", "array(1, 2, 3)", "[1,2,3]"}, {"map", "map('k', 9)", `{"k":9}`}, {"struct", "named_struct('a', 1, 'b', 'x')", `{"a":1,"b":"x"}`}, {"variant", `parse_json('{"a":1,"b":[2,3]}')`, `{"a":1,"b":[2,3]}`}, {"geometry", "st_point(1, 2)", "POINT(1 2)"}, + // GEOGRAPHY is the sibling geospatial type: it comes back as WKT text in a + // Utf8 column, so it scans to the same string as GEOMETRY. The ::geography + // cast may be gated on some warehouses; the subtest skips (rather than + // fails) when the server rejects the expression, so the case documents the + // parity without breaking runs on a warehouse without it. + {"geography", "cast(st_point(1, 2) as geography)", "POINT(1 2)"}, } for _, c := range cases { @@ -218,6 +230,9 @@ func TestKernelE2EDataTypes(t *testing.T) { var got any err := db.QueryRowContext(context.Background(), "SELECT "+c.expr).Scan(&got) if err != nil { + if skipOnErrCases[c.name] { + t.Skipf("server rejected %q (type may be gated on this warehouse): %v", c.expr, err) + } t.Fatalf("scan %s: %v", c.expr, err) } if !dataTypeEqual(got, c.want) { @@ -312,6 +327,44 @@ func TestKernelE2ECloudFetch(t *testing.T) { } } +// TestKernelE2EMaxChunksInMemory drives a CloudFetch-sized result with the +// in-memory-chunk knob lowered (WithKernelMaxChunksInMemory), proving the knob +// flows through the C-ABI set_session_conf, is accepted by the kernel, and is +// stripped before the SEA wire — a wrong/unrecognized session conf would surface +// as a server error, and the kernel would reject an invalid client key. The +// result must be complete and correct regardless of the chunk bound (it only +// changes peak memory, not output). +func TestKernelE2EMaxChunksInMemory(t *testing.T) { + db := kernelTestDBWith(t, WithKernelMaxChunksInMemory(4)) + defer db.Close() + + const want = 1_000_000 + rows, err := db.QueryContext(context.Background(), "SELECT id FROM range(0, 1000000)") + if err != nil { + t.Fatalf("query: %v", err) + } + defer rows.Close() + + var count, last int64 + for rows.Next() { + var id int64 + if err := rows.Scan(&id); err != nil { + t.Fatalf("scan at row %d: %v", count, err) + } + count++ + last = id + } + if err := rows.Err(); err != nil { + t.Fatalf("iteration: %v", err) + } + if count != want { + t.Errorf("row count = %d, want %d (result must be complete regardless of chunk bound)", count, want) + } + if last != want-1 { + t.Errorf("last id = %d, want %d", last, want-1) + } +} + // TestKernelE2EStatementID proves the kernel backend surfaces the server query id // on the success path: a registered QueryIdCallback fires with a non-empty id // (driven by kernelOp.StatementID() → kernel_executed_statement_query_id). This is diff --git a/kernel_experimental_test.go b/kernel_experimental_test.go index 1c3534d2..b54aeed4 100644 --- a/kernel_experimental_test.go +++ b/kernel_experimental_test.go @@ -4,7 +4,9 @@ import ( "context" "errors" "reflect" + "strings" "testing" + "time" "github.com/databricks/databricks-sql-go/internal/config" @@ -27,6 +29,12 @@ import ( var kernelExperimentalFieldDisposition = map[string]string{ "TLSTrustedCertsPEM": "forwarded", // set_tls_trusted_certs "TLSSkipHostnameVerify": "forwarded", // set_tls_skip_hostname_verification + "ProxyURL": "forwarded", // set_proxy (url) + "ProxyUsername": "forwarded", // set_proxy (username) + "ProxyPassword": "forwarded", // set_proxy (password) + "ProxyBypassHosts": "forwarded", // set_proxy (bypass_hosts) + "RetryOverallTimeout": "forwarded", // set_retry_config (overall_timeout_ms, 4th knob) + "MaxChunksInMemory": "forwarded", // set_session_conf (cloudfetch_max_chunks_in_memory, client-only) } func TestKernelExperimentalFieldsClassified(t *testing.T) { @@ -66,6 +74,16 @@ func TestWithKernelTLSOptionsSetExperimental(t *testing.T) { {"skip hostname", WithKernelSkipHostnameVerify(), func(k *config.KernelExperimentalConfig) bool { return k.TLSSkipHostnameVerify }}, + {"proxy", WithKernelProxy(KernelProxy{URL: "http://proxy:3128", Username: "u", Password: "p", BypassHosts: "*.internal"}), func(k *config.KernelExperimentalConfig) bool { + return k.ProxyURL == "http://proxy:3128" && k.ProxyUsername == "u" && + k.ProxyPassword == "p" && k.ProxyBypassHosts == "*.internal" + }}, + {"retry overall timeout", WithKernelRetryOverallTimeout(5 * time.Minute), func(k *config.KernelExperimentalConfig) bool { + return k.RetryOverallTimeout == 5*time.Minute + }}, + {"max chunks in memory", WithKernelMaxChunksInMemory(4), func(k *config.KernelExperimentalConfig) bool { + return k.MaxChunksInMemory == 4 + }}, } for _, c := range cases { t.Run(c.name, func(t *testing.T) { @@ -95,6 +113,9 @@ func TestWithKernelOptionsRejectedOnThriftPath(t *testing.T) { }{ {"trusted certs", WithKernelTrustedCerts([]byte("ca"))}, {"skip hostname", WithKernelSkipHostnameVerify()}, + {"proxy", WithKernelProxy(KernelProxy{URL: "http://proxy:3128"})}, + {"retry overall timeout", WithKernelRetryOverallTimeout(5 * time.Minute)}, + {"max chunks in memory", WithKernelMaxChunksInMemory(4)}, } for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { @@ -109,11 +130,22 @@ func TestWithKernelOptionsRejectedOnThriftPath(t *testing.T) { t.Fatalf("NewConnector: %v", err) } // No WithUseKernel — the Thrift path must reject the kernel-only option. - if _, err = c.Connect(context.Background()); err == nil { + _, err = c.Connect(context.Background()) + if err == nil { t.Fatal("Connect should reject a WithKernel* option on the Thrift path, got nil") - } else if !errors.Is(err, dbsqlerr.ErrRequiresKernelBackend) { + } + if !errors.Is(err, dbsqlerr.ErrRequiresKernelBackend) { t.Errorf("error should wrap ErrRequiresKernelBackend; got: %v", err) } + // The message must name the WithKernel* family generically, not a stale + // subset (it used to hardcode the two TLS options, misdirecting a caller who + // set proxy / retry / chunk options). Guard against that regression. + if msg := err.Error(); !strings.Contains(msg, "WithKernel* option") { + t.Errorf("rejection message should name the WithKernel* family; got: %q", msg) + } + if msg := err.Error(); strings.Contains(msg, "WithKernelTrustedCerts") || strings.Contains(msg, "WithKernelSkipHostnameVerify") { + t.Errorf("rejection message should not name a stale option subset; got: %q", msg) + } }) } } @@ -145,11 +177,27 @@ func TestKernelExperimentalDeepCopy(t *testing.T) { orig := &config.KernelExperimentalConfig{ TLSTrustedCertsPEM: []byte("ca-bundle"), TLSSkipHostnameVerify: true, + ProxyURL: "http://proxy:3128", + ProxyUsername: "u", + ProxyPassword: "p", + ProxyBypassHosts: "*.internal", + RetryOverallTimeout: 5 * time.Minute, + MaxChunksInMemory: 4, } cp := orig.DeepCopy() if cp == nil || string(cp.TLSTrustedCertsPEM) != "ca-bundle" || !cp.TLSSkipHostnameVerify { t.Fatalf("DeepCopy lost data: %+v", cp) } + if cp.ProxyURL != "http://proxy:3128" || cp.ProxyUsername != "u" || + cp.ProxyPassword != "p" || cp.ProxyBypassHosts != "*.internal" { + t.Errorf("DeepCopy lost proxy fields: %+v", cp) + } + if cp.RetryOverallTimeout != 5*time.Minute { + t.Errorf("DeepCopy lost RetryOverallTimeout: %v", cp.RetryOverallTimeout) + } + if cp.MaxChunksInMemory != 4 { + t.Errorf("DeepCopy lost MaxChunksInMemory: %v", cp.MaxChunksInMemory) + } cp.TLSTrustedCertsPEM[0] = 'X' if orig.TLSTrustedCertsPEM[0] == 'X' { t.Error("DeepCopy aliased the CA byte slice; a copy mutation reached the original") diff --git a/kernel_newitems_e2e_test.go b/kernel_newitems_e2e_test.go new file mode 100644 index 00000000..da9f3368 --- /dev/null +++ b/kernel_newitems_e2e_test.go @@ -0,0 +1,200 @@ +//go:build cgo && databricks_kernel + +package dbsql + +import ( + "bufio" + "context" + "encoding/base64" + "net" + "net/http" + "sync" + "testing" + "time" +) + +// This file holds the live end-to-end tests for the "new-items" kernel features +// (advanced proxy, Geography, configurable backoff, conn/error telemetry). They +// run against the same DATABRICKS_PECOTESTING_* warehouse as the rest of the +// kernel E2E suite (via kernelTestDBWith / pecoTestingCreds) and self-skip when +// those credentials are unset. + +// TestKernelE2EProxyRejectsBadURL proves the advanced-proxy setter (WithKernelProxy) +// is actually applied on the kernel path. No proxy is provisioned for the staging +// warehouse, so a real round-trip can't be asserted; instead we route through an +// unreachable proxy and require the connect to FAIL. An ignored proxy setter would +// let the connection succeed directly, so a failure here is the proof the URL +// reached the kernel's HTTP stack. A bounded context keeps a hung dial from +// stalling the test. +func TestKernelE2EProxyRejectsBadURL(t *testing.T) { + // A routable-but-dead proxy address: connections are refused/time out rather + // than silently bypassed. 127.0.0.1:1 has no listener. + db := kernelTestDBWith(t, WithKernelProxy(KernelProxy{URL: "http://127.0.0.1:1"})) + defer db.Close() + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + var got int64 + err := db.QueryRowContext(ctx, "SELECT 1").Scan(&got) + if err == nil { + t.Fatal("expected the query to fail when routed through an unreachable proxy " + + "(an ignored proxy setter would connect directly), got success") + } + t.Logf("connect through unreachable proxy failed as expected: %v", err) +} + +// TestKernelE2EProxyForwardsCredentials proves WithKernelProxy's basic-auth +// credentials reach the kernel in the RIGHT slots — the gap TestSetProxy can't cover, +// since it only asserts the setter returns OK and the kernel's C config is opaque +// (no readback), so a username↔password swap or a dropped bypass would pass every +// unit test yet break proxy basic-auth in production. +// +// A local CONNECT proxy captures the Proxy-Authorization header the kernel's HTTP +// stack sends, then refuses to tunnel. The connect is EXPECTED to fail (nothing is +// tunneled), but the captured header is the proof: we assert it decodes to exactly +// "user:pass" (a slot swap would decode to "pass:user"). Observing the credential on +// the wire needs no real warehouse — only that the kernel forwarded it correctly. +func TestKernelE2EProxyForwardsCredentials(t *testing.T) { + const wantUser, wantPass = "proxyuser", "proxypass" + + var mu sync.Mutex + var gotAuth string + sawConnect := make(chan struct{}, 1) + + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("listen: %v", err) + } + defer func() { _ = ln.Close() }() + + go func() { + for { + conn, err := ln.Accept() + if err != nil { + return // listener closed + } + go func(c net.Conn) { + defer func() { _ = c.Close() }() + br := bufio.NewReader(c) + // Handle both proxy-auth timings on one connection: a client that sends + // Proxy-Authorization preemptively, and one that waits for a 407 + // challenge before resending. Loop so a same-connection retry after our + // 407 is captured; clients that open a fresh connection are covered by + // the accept loop. + for { + req, err := http.ReadRequest(br) + if err != nil { + return + } + if auth := req.Header.Get("Proxy-Authorization"); auth != "" { + mu.Lock() + gotAuth = auth + mu.Unlock() + select { + case sawConnect <- struct{}{}: + default: + } + // Captured — refuse to tunnel so the connect fails fast. + _, _ = c.Write([]byte("HTTP/1.1 502 Bad Gateway\r\n\r\n")) + return + } + // No credentials yet — challenge, then read the retry on this conn. + _, _ = c.Write([]byte("HTTP/1.1 407 Proxy Authentication Required\r\n" + + "Proxy-Authenticate: Basic realm=\"test\"\r\n" + + "Content-Length: 0\r\n\r\n")) + } + }(conn) + } + }() + + proxyURL := "http://" + ln.Addr().String() + db := kernelTestDBWith(t, WithKernelProxy(KernelProxy{URL: proxyURL, Username: wantUser, Password: wantPass})) + defer db.Close() + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + // The query is expected to fail (the proxy never tunnels); we only need the + // kernel to have attempted CONNECT through the proxy so the header is captured. + var got int64 + _ = db.QueryRowContext(ctx, "SELECT 1").Scan(&got) + + select { + case <-sawConnect: + case <-time.After(30 * time.Second): + t.Fatal("kernel never issued a CONNECT through the configured proxy") + } + + mu.Lock() + auth := gotAuth + mu.Unlock() + if auth == "" { + t.Fatal("proxy saw no Proxy-Authorization header — credentials were dropped") + } + const prefix = "Basic " + if len(auth) <= len(prefix) || auth[:len(prefix)] != prefix { + t.Fatalf("Proxy-Authorization = %q, want a Basic credential", auth) + } + decoded, err := base64.StdEncoding.DecodeString(auth[len(prefix):]) + if err != nil { + t.Fatalf("decode Proxy-Authorization: %v", err) + } + if want := wantUser + ":" + wantPass; string(decoded) != want { + t.Errorf("proxy credentials = %q, want %q (a username/password slot swap would show %q)", + decoded, want, wantPass+":"+wantUser) + } +} + +// TestKernelE2ERetryConfig proves a tuned WithRetries policy (backoff bounds + max +// attempts) is accepted at connect and the connection still works. It is a +// connect-succeeds smoke test — a happy-path query exercises no retry, so it can't +// observe whether the policy is actually applied vs ignored; the pure-Go +// TestKernelRetryConfig is what pins the WithRetries -> kernel.RetryConfig mapping. +func TestKernelE2ERetryConfig(t *testing.T) { + db := kernelTestDBWith(t, WithRetries(6, 500*time.Millisecond, 20*time.Second)) + defer db.Close() + + var got int64 + if err := db.QueryRowContext(context.Background(), "SELECT 1").Scan(&got); err != nil { + t.Fatalf("query with tuned retries: %v", err) + } + if got != 1 { + t.Errorf("SELECT 1 = %d, want 1", got) + } +} + +// TestKernelE2ERetryDisabled proves the disable form (WithRetries(-1)) is accepted +// on the kernel path — previously it was rejected at connect. It maps to zero +// kernel retries, and a normal query still succeeds. +func TestKernelE2ERetryDisabled(t *testing.T) { + db := kernelTestDBWith(t, WithRetries(-1, 0, 0)) + defer db.Close() + + var got int64 + if err := db.QueryRowContext(context.Background(), "SELECT 1").Scan(&got); err != nil { + t.Fatalf("query with retries disabled: %v", err) + } + if got != 1 { + t.Errorf("SELECT 1 = %d, want 1", got) + } +} + +// TestKernelE2ERetryOverallTimeout proves the kernel-only overall-budget knob +// (WithKernelRetryOverallTimeout, the 4th retry control) is accepted at connect and +// the connection works with it set. Like TestKernelE2ERetryConfig this is a +// connect-succeeds smoke test — the happy path triggers no retry, so it can't +// distinguish budget-applied from budget-ignored; TestKernelRetryConfig pins the +// forwarding in pure Go. +func TestKernelE2ERetryOverallTimeout(t *testing.T) { + db := kernelTestDBWith(t, WithKernelRetryOverallTimeout(5*time.Minute)) + defer db.Close() + + var got int64 + if err := db.QueryRowContext(context.Background(), "SELECT 1").Scan(&got); err != nil { + t.Fatalf("query with overall retry budget: %v", err) + } + if got != 1 { + t.Errorf("SELECT 1 = %d, want 1", got) + } +} diff --git a/kernel_parity_test.go b/kernel_parity_test.go index 5744c3ce..4987e152 100644 --- a/kernel_parity_test.go +++ b/kernel_parity_test.go @@ -5,6 +5,7 @@ package dbsql import ( "context" "database/sql" + "fmt" "testing" ) @@ -92,6 +93,164 @@ func TestKernelE2EInterval(t *testing.T) { } } +// TestKernelThriftColumnTypeParity is the live guard for PECOBLR-3692: the kernel +// backend must report the SAME sql.ColumnType metadata (DatabaseTypeName, +// ScanType, Nullable, Length) as the Thrift backend for every type. The +// value-parity suites above compare scanned VALUES and so are blind to this — the +// gap was a Rows.ColumnType* omission (kernelRows implemented only driver.Rows), +// which surfaces as "" / interface{} metadata, not a wrong value. Named +// TestKernelThrift* / TestKernel* so the nightly kernel -run picks it up. +func TestKernelThriftColumnTypeParity(t *testing.T) { + // One column per Databricks type the driver scans, mirroring the mixed-type + // query used to capture the ground truth. VARCHAR/CHAR/VARIANT/GEOMETRY all + // collapse to STRING on both backends (they arrive as Arrow Utf8), and both + // interval types are server/Go-stringified to STRING — included so the parity + // covers those collapses too. + const query = "SELECT " + + "CAST(1 AS TINYINT) a_tinyint, CAST(1 AS SMALLINT) a_smallint, " + + "CAST(1 AS INT) a_int, CAST(1 AS BIGINT) a_bigint, " + + "CAST(1 AS FLOAT) a_float, CAST(1 AS DOUBLE) a_double, " + + "CAST(1 AS BOOLEAN) a_bool, CAST('x' AS STRING) a_string, " + + "CAST('x' AS VARCHAR(10)) a_varchar, CAST('x' AS CHAR(3)) a_char, " + + "CAST(1.5 AS DECIMAL(10,2)) a_decimal, CAST('2020-01-01' AS DATE) a_date, " + + "CAST('2020-01-01 00:00:00' AS TIMESTAMP) a_ts, CAST('abc' AS BINARY) a_binary, " + + "array(1,2,3) a_array, map('k',1) a_map, named_struct('x',1) a_struct, " + + `parse_json('{"a":1}') a_variant, INTERVAL '1' DAY a_iv_dt, ` + + "INTERVAL '1' MONTH a_iv_ym, st_point(1,2) a_geom" + + kernelDB := kernelTestDB(t) + defer kernelDB.Close() + thriftDB := thriftTestDB(t) + defer thriftDB.Close() + + kernelCT := columnTypeStrings(t, kernelDB, query) + thriftCT := columnTypeStrings(t, thriftDB, query) + + if len(kernelCT) != len(thriftCT) { + t.Fatalf("column count differs: kernel=%d thrift=%d", len(kernelCT), len(thriftCT)) + } + for i := range kernelCT { + if kernelCT[i] != thriftCT[i] { + t.Errorf("col %d metadata differs:\n kernel=%s\n thrift=%s", i, kernelCT[i], thriftCT[i]) + } + } +} + +// columnTypeStrings renders each column's sql.ColumnType metadata to a stable +// string (name, DatabaseTypeName, ScanType, Nullable, Length), so the kernel and +// Thrift backends can be compared field-for-field. A nil ScanType renders as +// "" so the pre-fix interface{} fallback would differ visibly. +func columnTypeStrings(t *testing.T, db *sql.DB, query string) []string { + t.Helper() + rows, err := db.QueryContext(context.Background(), query) + if err != nil { + t.Fatalf("query: %v", err) + } + defer rows.Close() + cts, err := rows.ColumnTypes() + if err != nil { + t.Fatalf("ColumnTypes: %v", err) + } + out := make([]string, len(cts)) + for i, ct := range cts { + scan := "" + if ct.ScanType() != nil { + scan = ct.ScanType().String() + } + nl, nlok := ct.Nullable() + ln, lnok := ct.Length() + out[i] = fmt.Sprintf("name=%s db=%s scan=%s nullable=%v/%v len=%d/%v", + ct.Name(), ct.DatabaseTypeName(), scan, nl, nlok, ln, lnok) + } + return out +} + +// TestKernelThriftDecimalScaleParity exercises the shared exact-decimal renderer +// (decimalfmt, PECOBLR-3691) at scale against a live warehouse: it generates many +// decimal rows so the kernel result spans multiple Arrow batches, and asserts +// every rendered value matches the Thrift backend byte-for-byte — the guarantee +// that the alloc-free ExactString rewrite changed only cost, not output, over a +// real multi-batch stream and including magnitudes past float64 precision. The +// pure-Go TestExactStringOracleParity pins the renderer against its own +// pre-rewrite implementation; this pins the integrated kernel scan path against +// Thrift. Named TestKernelThrift* for the nightly -run. +func TestKernelThriftDecimalScaleParity(t *testing.T) { + // 50k rows × several DECIMALs, with a value beyond float64's exact range + // (DECIMAL(38,4)) so a lossy path would diverge. range() drives enough rows to + // cross batch boundaries on the kernel stream. + const rowCount = 50000 + query := "SELECT " + + "CAST(id AS DECIMAL(10,2)) d1, " + + "CAST(id * -1.25 AS DECIMAL(20,4)) d2, " + + "CAST(id AS DECIMAL(38,4)) + 123456789012345678901234.5678 d3, " + + "CAST(id % 7 AS DECIMAL(5,3)) d4 " + + "FROM range(" + fmt.Sprint(rowCount) + ") ORDER BY id" + + kernelDB := kernelTestDB(t) + defer kernelDB.Close() + thriftDB := thriftTestDB(t) + defer thriftDB.Close() + + kernelRows := scanAllRowsAsStrings(t, kernelDB, query) + thriftRows := scanAllRowsAsStrings(t, thriftDB, query) + + if len(kernelRows) != len(thriftRows) { + t.Fatalf("row count differs: kernel=%d thrift=%d", len(kernelRows), len(thriftRows)) + } + if len(kernelRows) != rowCount { + t.Fatalf("expected %d rows, got %d", rowCount, len(kernelRows)) + } + for i := range kernelRows { + if kernelRows[i] != thriftRows[i] { + t.Fatalf("row %d differs:\n kernel=%s\n thrift=%s", i, kernelRows[i], thriftRows[i]) + } + } + t.Logf("verified %d decimal rows byte-identical across backends (arena spans multiple batches)", len(kernelRows)) +} + +// scanAllRowsAsStrings scans every row into a "|"-joined string of column wire +// forms (via sql.RawBytes), so two backends can be compared row-by-row +// independent of Go-type coercion. +func scanAllRowsAsStrings(t *testing.T, db *sql.DB, query string) []string { + t.Helper() + rows, err := db.QueryContext(context.Background(), query) + if err != nil { + t.Fatalf("query: %v", err) + } + defer rows.Close() + cols, err := rows.Columns() + if err != nil { + t.Fatalf("columns: %v", err) + } + var out []string + raw := make([]sql.RawBytes, len(cols)) + dest := make([]any, len(cols)) + for i := range raw { + dest[i] = &raw[i] + } + for rows.Next() { + if err := rows.Scan(dest...); err != nil { + t.Fatalf("scan: %v", err) + } + line := "" + for i, b := range raw { + if i > 0 { + line += "|" + } + if b == nil { + line += "" + } else { + line += string(b) + } + } + out = append(out, line) + } + if err := rows.Err(); err != nil { + t.Fatalf("rows err: %v", err) + } + return out +} + // paramCase is one parameterized-query parity case: the same SQL + args run on // both backends must yield byte-identical output. type paramCase struct { diff --git a/kernel_proxy_test.go b/kernel_proxy_test.go index 0da7c72c..3bfed871 100644 --- a/kernel_proxy_test.go +++ b/kernel_proxy_test.go @@ -6,6 +6,7 @@ import ( "net/url" "testing" + "github.com/databricks/databricks-sql-go/internal/backend/kernel" "github.com/databricks/databricks-sql-go/internal/config" ) @@ -91,3 +92,45 @@ func TestProxyForEndpoint(t *testing.T) { // The production wrapper wires http.ProxyFromEnvironment and must not panic. _ = proxyForEndpoint(validCfg()) } + +// resolveKernelProxy: an explicit WithKernelProxy (KernelExperimental.ProxyURL +// set) wins verbatim — url + credentials + bypass list — over the environment; +// with no explicit proxy it falls back to the endpoint's env-derived URL and +// leaves credentials / bypass empty. Pure Go, default CGO_ENABLED=0 build. +func TestResolveKernelProxy(t *testing.T) { + t.Run("explicit WithKernelProxy wins verbatim over env", func(t *testing.T) { + c := config.WithDefaults() + c.Host = "my-workspace.databricks.com" + c.Port = 443 + c.WarehouseID = "abc" + c.KernelExperimental = &config.KernelExperimentalConfig{ //nolint:gosec // G101: test literals (ProxyPassword), not real credentials + ProxyURL: "http://explicit-proxy:8080", + ProxyUsername: "user", + ProxyPassword: "pass", + ProxyBypassHosts: "localhost,*.internal", + } + var kc kernel.Config + resolveKernelProxy(c, &kc) + if kc.ProxyURL != "http://explicit-proxy:8080" { + t.Errorf("ProxyURL = %q, want the explicit WithKernelProxy URL", kc.ProxyURL) + } + if kc.ProxyUsername != "user" || kc.ProxyPassword != "pass" || kc.ProxyBypassHosts != "localhost,*.internal" { + t.Errorf("explicit proxy credentials/bypass not forwarded: %+v", kc) + } + }) + + t.Run("no explicit proxy falls back to env (no creds/bypass)", func(t *testing.T) { + c := config.WithDefaults() + c.Host = "my-workspace.databricks.com" + c.Port = 443 + c.WarehouseID = "abc" + // No KernelExperimental proxy → env resolution. With no proxy env set in + // the test process this resolves to direct (""), and the credential / + // bypass fields must stay empty (the env path can't carry them). + var kc kernel.Config + resolveKernelProxy(c, &kc) + if kc.ProxyUsername != "" || kc.ProxyPassword != "" || kc.ProxyBypassHosts != "" { + t.Errorf("env path must not populate credentials/bypass: %+v", kc) + } + }) +} diff --git a/kernel_telemetry.go b/kernel_telemetry.go new file mode 100644 index 00000000..2639f51d --- /dev/null +++ b/kernel_telemetry.go @@ -0,0 +1,111 @@ +package dbsql + +import ( + "github.com/databricks/databricks-sql-go/internal/backend/kernel" + "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: +// the connection-config telemetry payload is assembled from pure Go config (no +// kernel C symbol), so keeping it untagged lets its test run under CGO_ENABLED=0. +// The connector (also untagged) emits it on the kernel path right after +// CREATE_SESSION. + +// kernelConnectionTelemetry builds the connection-configuration telemetry payload +// for a kernel-backed connection: the resolved connection parameters (http path, +// proxy usage, arrow, query tags, metric-view metadata) and the auth mechanism. +// It populates the DriverConnectionParameters shape the proto defines (which the +// Thrift path never populates on this driver either); the kernel path is the first +// to actually emit it. Pure Go so it is unit-testable in the default build. +// +// It reports only what the kernel path genuinely applies: EnableArrow is always +// true (the kernel returns Arrow results), UseProxy reflects whether a proxy was +// resolved (env-var or WithKernelProxy), and it does not claim direct-results or a +// socket timeout the kernel C ABI has no knob for. +func kernelConnectionTelemetry(cfg *config.Config) *telemetry.DriverConnectionParameters { + params := &telemetry.DriverConnectionParameters{ + HTTPPath: cfg.HTTPPath, + // The `mode` telemetry field is a closed server-side enum + // (DatabricksClientType = SEA | THRIFT | TYPE_UNSPECIFIED); it has no + // "kernel" member. The kernel backend speaks the Statement Execution API, + // so it reports "SEA" — the same value the JDBC and Python SEA paths land. + // Emitting "kernel" here would drop the field to NULL on ingestion (the + // enum rejects unknown members). + Mode: "SEA", + EnableArrow: true, + HostInfo: &telemetry.HostDetails{ + HostURL: cfg.Host, + Port: int32(cfg.Port), //nolint:gosec // port is a small positive int + }, + UseProxy: kernelUsesProxy(cfg), + EnableMetricViewMeta: cfg.EnableMetricViewMetadata, + } + if qt := cfg.SessionParams["QUERY_TAGS"]; qt != "" { + params.QueryTags = qt + } + mech, flow := kernelAuthMech(cfg) + params.AuthMech = mech + params.AuthFlow = flow + + // Resolved retry policy + CloudFetch memory cap, so a hung-connect or + // large-result-OOM report can be diagnosed from telemetry. Only positive values + // are emitted (omitempty): a non-positive RetryMax is the disable/default form, + // and a zero timeout / chunk cap means "keep the kernel default". Mirrors what + // buildKernelConfig / kernelRetryConfig forward to the kernel. + if cfg.RetryMax > 0 { + params.RetryMaxAttempts = int32(cfg.RetryMax) //nolint:gosec // small positive attempt count + } + if ke := cfg.KernelExperimental; ke != nil { + if ke.RetryOverallTimeout > 0 { + params.RetryOverallTimeoutMs = ke.RetryOverallTimeout.Milliseconds() + } + if ke.MaxChunksInMemory > 0 { + params.MaxChunksInMemory = int32(ke.MaxChunksInMemory) //nolint:gosec // small positive chunk count + } + } + return params +} + +// kernelUsesProxy reports whether the kernel connection will route through a +// proxy: either an explicit WithKernelProxy or an environment-resolved one for +// this endpoint. Mirrors resolveKernelProxy's proxy-source precedence. +func kernelUsesProxy(cfg *config.Config) bool { + if ke := cfg.KernelExperimental; ke != nil && ke.ProxyURL != "" { + return true + } + return proxyForEndpoint(cfg) != "" +} + +// kernelAuthMech maps the resolved kernel auth form to the telemetry auth_mech and +// auth_flow fields. Both are closed server-side enums, NOT the driver's own auth +// names — emitting a non-enum string (e.g. "Pat", "OauthM2M") drops the field to +// NULL on ingestion. The mapping mirrors the Python driver's +// TelemetryHelper.get_auth_mechanism / get_auth_flow so all drivers land the same +// values: +// - auth_mech (AuthMech enum): PAT | OAUTH +// - auth_flow (AuthFlow enum): CLIENT_CREDENTIALS (M2M) | BROWSER_BASED_AUTHENTICATION (U2M) +// +// PAT has no auth_flow (it is not an OAuth flow), so auth_flow is "" there. An +// unresolvable auth (rejected at connect anyway) reports empty (unspecified). +func kernelAuthMech(cfg *config.Config) (mech, flow string) { + const ( + authMechPAT = "PAT" + authMechOAuth = "OAUTH" + + authFlowClientCreds = "CLIENT_CREDENTIALS" //nolint:gosec // G101: telemetry auth_flow enum value, not a credential + authFlowBrowser = "BROWSER_BASED_AUTHENTICATION" + ) + ka, err := resolveKernelAuth(cfg) + if err != nil { + return "", "" + } + switch ka.Mode { + case kernel.AuthM2M: + return authMechOAuth, authFlowClientCreds + case kernel.AuthU2M: + return authMechOAuth, authFlowBrowser + default: + return authMechPAT, "" + } +} diff --git a/kernel_telemetry_test.go b/kernel_telemetry_test.go new file mode 100644 index 00000000..3bf3152d --- /dev/null +++ b/kernel_telemetry_test.go @@ -0,0 +1,145 @@ +package dbsql + +import ( + "testing" + "time" + + "github.com/databricks/databricks-sql-go/auth/pat" + "github.com/databricks/databricks-sql-go/internal/config" +) + +// These tests run in the default CGO_ENABLED=0 build (kernel_telemetry.go is +// untagged): the connection-config payload is assembled from pure Go config. + +// kernelConnectionTelemetry must emit the closed-enum values the ingestion schema +// accepts — mode="SEA" (not "kernel"), auth_mech ∈ {PAT, OAUTH}, auth_flow ∈ +// {CLIENT_CREDENTIALS, BROWSER_BASED_AUTHENTICATION} — and reflect the resolved +// connection config (arrow, http path, query tags, metric-view, proxy usage). +func TestKernelConnectionTelemetry(t *testing.T) { + t.Run("PAT connection reports SEA + PAT with no auth_flow", func(t *testing.T) { + cfg := config.WithDefaults() + cfg.Host = "example.cloud.databricks.com" + cfg.Port = 443 + cfg.HTTPPath = "/sql/1.0/warehouses/abc" + cfg.AccessToken = "dapi-x" + cfg.Authenticator = &pat.PATAuth{AccessToken: "dapi-x"} + + p := kernelConnectionTelemetry(cfg) + if p.Mode != "SEA" { + t.Errorf("Mode = %q, want SEA (the kernel speaks the Statement Execution API; 'kernel' would NULL the field)", p.Mode) + } + if p.AuthMech != "PAT" { + t.Errorf("AuthMech = %q, want PAT", p.AuthMech) + } + if p.AuthFlow != "" { + t.Errorf("AuthFlow = %q, want empty for PAT (not an OAuth flow)", p.AuthFlow) + } + if !p.EnableArrow { + t.Error("EnableArrow = false, want true (the kernel returns Arrow results)") + } + if p.HTTPPath != "/sql/1.0/warehouses/abc" { + t.Errorf("HTTPPath = %q, want the configured path", p.HTTPPath) + } + if p.HostInfo == nil || p.HostInfo.HostURL != "example.cloud.databricks.com" || p.HostInfo.Port != 443 { + t.Errorf("HostInfo = %+v, want host/port populated", p.HostInfo) + } + }) + + t.Run("query tags + metric-view are reflected", func(t *testing.T) { + cfg := config.WithDefaults() + cfg.AccessToken = "dapi-x" + cfg.SessionParams = map[string]string{"QUERY_TAGS": "team=peco"} + cfg.EnableMetricViewMetadata = true + + p := kernelConnectionTelemetry(cfg) + if p.QueryTags != "team=peco" { + t.Errorf("QueryTags = %q, want team=peco", p.QueryTags) + } + if !p.EnableMetricViewMeta { + t.Error("EnableMetricViewMeta = false, want true") + } + }) + + t.Run("explicit WithKernelProxy marks UseProxy", func(t *testing.T) { + cfg := config.WithDefaults() + cfg.AccessToken = "dapi-x" + WithKernelProxy(KernelProxy{URL: "http://proxy:3128"})(cfg) + + if p := kernelConnectionTelemetry(cfg); !p.UseProxy { + t.Error("UseProxy = false, want true when WithKernelProxy is set") + } + }) + + t.Run("resolved retry policy + chunk cap are reflected", func(t *testing.T) { + cfg := config.WithDefaults() + cfg.AccessToken = "dapi-x" + WithRetries(7, 0, 0)(cfg) + WithKernelRetryOverallTimeout(90 * time.Second)(cfg) + WithKernelMaxChunksInMemory(4)(cfg) + + p := kernelConnectionTelemetry(cfg) + if p.RetryMaxAttempts != 7 { + t.Errorf("RetryMaxAttempts = %d, want 7", p.RetryMaxAttempts) + } + if p.RetryOverallTimeoutMs != 90_000 { + t.Errorf("RetryOverallTimeoutMs = %d, want 90000", p.RetryOverallTimeoutMs) + } + if p.MaxChunksInMemory != 4 { + t.Errorf("MaxChunksInMemory = %d, want 4", p.MaxChunksInMemory) + } + }) + + t.Run("retry/chunk fields stay zero (omitted) at defaults", func(t *testing.T) { + // WithDefaults sets RetryMax=4, but the disable/default form and the unset + // kernel-only knobs must not fabricate telemetry: only an explicit positive + // RetryMax is emitted, so document that the default RetryMax IS surfaced while + // the kernel-only knobs stay zero when unset. + cfg := config.WithDefaults() + cfg.AccessToken = "dapi-x" + p := kernelConnectionTelemetry(cfg) + if int(p.RetryMaxAttempts) != cfg.RetryMax { + t.Errorf("RetryMaxAttempts = %d, want %d (the resolved RetryMax)", p.RetryMaxAttempts, cfg.RetryMax) + } + if p.RetryOverallTimeoutMs != 0 || p.MaxChunksInMemory != 0 { + t.Errorf("kernel-only knobs should be zero when unset: overall=%d chunks=%d", + p.RetryOverallTimeoutMs, p.MaxChunksInMemory) + } + }) +} + +// kernelAuthMech maps each resolvable auth form to the closed AuthMech/AuthFlow +// enums. M2M and U2M are both OAUTH but differ in flow; PAT is PAT with no flow. +func TestKernelAuthMech(t *testing.T) { + t.Run("PAT", func(t *testing.T) { + cfg := config.WithDefaults() + cfg.AccessToken = "dapi-x" + mech, flow := kernelAuthMech(cfg) + if mech != "PAT" || flow != "" { + t.Errorf("PAT -> (%q, %q), want (PAT, \"\")", mech, flow) + } + }) + + // M2M and U2M are the values that silently NULL server-side if the (mech, flow) + // enums are wrong, so assert both arms explicitly — not just PAT. resolveKernelAuth + // selects the arm by asserting the M2M/U2M provider interfaces on cfg.Authenticator, + // which the fake authenticators satisfy (see kernel_config_test.go). + t.Run("M2M -> OAUTH / CLIENT_CREDENTIALS", func(t *testing.T) { + cfg := config.WithDefaults() + cfg.AccessToken = "" + cfg.Authenticator = fakeM2MAuth{id: "cid", secret: "sec", scopes: []string{"all-apis"}} + mech, flow := kernelAuthMech(cfg) + if mech != "OAUTH" || flow != "CLIENT_CREDENTIALS" { + t.Errorf("M2M -> (%q, %q), want (OAUTH, CLIENT_CREDENTIALS)", mech, flow) + } + }) + + t.Run("U2M -> OAUTH / BROWSER_BASED_AUTHENTICATION", func(t *testing.T) { + cfg := config.WithDefaults() + cfg.AccessToken = "" + cfg.Authenticator = fakeU2MAuth{id: "databricks-sql-connector"} + mech, flow := kernelAuthMech(cfg) + if mech != "OAUTH" || flow != "BROWSER_BASED_AUTHENTICATION" { + t.Errorf("U2M -> (%q, %q), want (OAUTH, BROWSER_BASED_AUTHENTICATION)", mech, flow) + } + }) +} diff --git a/telemetry/aggregator.go b/telemetry/aggregator.go index 76565129..7e661b38 100644 --- a/telemetry/aggregator.go +++ b/telemetry/aggregator.go @@ -160,16 +160,29 @@ func (agg *metricsAggregator) recordMetric(ctx context.Context, metric *telemetr } case "error": - // Check if terminal error + // Defensive: no exported API currently emits a standalone "error" metric — + // errors ride along on "operation" (errorType set) and "statement" metrics + // instead — so this arm is unreachable from production code today and its + // routing is exercised only by direct recordMetric unit tests. It is kept so + // that if a producer is ever added the routing is already correct rather than + // silently dropping. NOTE for a future producer: each non-terminal + // connection-scoped error here triggers an immediate un-batched flushUnlocked + // (one export per error) — batch at the producer if that volume matters. if metric.errorType != "" && isTerminalError(&simpleError{msg: metric.errorType}) { - // Flush terminal errors immediately + // Terminal error: flush immediately. agg.batch = append(agg.batch, metric) agg.flushUnlocked(ctx) + } else if stmt, exists := agg.statements[metric.statementID]; exists { + // Non-terminal error with a live statement: buffer onto it, flushed when + // the statement completes. + stmt.errors = append(stmt.errors, metric.errorType) } else { - // Buffer non-terminal errors with statement - if stmt, exists := agg.statements[metric.statementID]; exists { - stmt.errors = append(stmt.errors, metric.errorType) - } + // Non-terminal error with no statement to attach to — a connection-scoped + // error (empty statementID), or one whose statement was already flushed. + // Emit it standalone rather than dropping it, so a connection-level error + // would still land if a producer existed. + agg.batch = append(agg.batch, metric) + agg.flushUnlocked(ctx) } } } diff --git a/telemetry/conn_error_event_test.go b/telemetry/conn_error_event_test.go new file mode 100644 index 00000000..7ede3415 --- /dev/null +++ b/telemetry/conn_error_event_test.go @@ -0,0 +1,203 @@ +package telemetry + +import ( + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "sync" + "testing" + "time" +) + +// decodeEvents parses a captured TelemetryRequest body into its TelemetryEvents. +func decodeEvents(t *testing.T, body []byte) []*TelemetryEvent { + t.Helper() + var req TelemetryRequest + if err := json.Unmarshal(body, &req); err != nil { + t.Fatalf("unmarshal request: %v", err) + } + var evs []*TelemetryEvent + for _, pl := range req.ProtoLogs { + var fl TelemetryFrontendLog + if err := json.Unmarshal([]byte(pl), &fl); err != nil { + t.Fatalf("unmarshal protoLog: %v", err) + } + if fl.Entry != nil && fl.Entry.SQLDriverLog != nil { + evs = append(evs, fl.Entry.SQLDriverLog) + } + } + return evs +} + +// createTelemetryRequest must serialize connParams into DriverConnectionParameters +// (and mirror AuthMech onto the top-level AuthType) for a "connection" metric, and +// must leave those fields nil/empty for a plain statement metric — so existing +// metrics serialize byte-identically to before this change. +func TestCreateTelemetryRequest_ConnParams(t *testing.T) { + connMetric := &telemetryMetric{ + metricType: "connection", + timestamp: time.Now(), + sessionID: "sess-1", + connParams: &DriverConnectionParameters{ + Mode: "SEA", + AuthMech: "PAT", + EnableArrow: true, + HTTPPath: "/sql/1.0/warehouses/abc", + }, + } + stmtMetric := &telemetryMetric{ + metricType: "statement", + timestamp: time.Now(), + sessionID: "sess-1", + statementID: "stmt-1", + latencyMs: 5, + } + + req, err := createTelemetryRequest([]*telemetryMetric{connMetric, stmtMetric}, "v-test") + if err != nil { + t.Fatalf("createTelemetryRequest: %v", err) + } + if len(req.ProtoLogs) != 2 { + t.Fatalf("want 2 protoLogs, got %d", len(req.ProtoLogs)) + } + + // The connection metric carries the params + AuthType mirror. + var connLog TelemetryFrontendLog + if err := json.Unmarshal([]byte(req.ProtoLogs[0]), &connLog); err != nil { + t.Fatal(err) + } + ev := connLog.Entry.SQLDriverLog + if ev.DriverConnectionParameters == nil { + t.Fatal("connection metric: DriverConnectionParameters must be populated") + } + if ev.DriverConnectionParameters.Mode != "SEA" || ev.DriverConnectionParameters.AuthMech != "PAT" { + t.Errorf("conn params = %+v, want Mode=SEA AuthMech=PAT", ev.DriverConnectionParameters) + } + if ev.AuthType != "PAT" { + t.Errorf("AuthType = %q, want PAT (mirrored from AuthMech)", ev.AuthType) + } + + // The statement metric must NOT carry conn params or auth type — byte-identical + // to a pre-change statement metric. Guards the "additive, behavior-preserving" + // contract for the default (Thrift) path's telemetry. + var stmtLog TelemetryFrontendLog + if err := json.Unmarshal([]byte(req.ProtoLogs[1]), &stmtLog); err != nil { + t.Fatal(err) + } + sev := stmtLog.Entry.SQLDriverLog + if sev.DriverConnectionParameters != nil { + t.Errorf("statement metric leaked DriverConnectionParameters: %+v", sev.DriverConnectionParameters) + } + if sev.AuthType != "" { + t.Errorf("statement metric leaked AuthType = %q, want empty", sev.AuthType) + } +} + +// Full loop: RecordConnectionConfig flushes a connection event immediately (the +// connection may close before the next batch), and it lands at the HTTP endpoint +// carrying the conn params. +func TestRecordConnectionConfig_EndToEnd(t *testing.T) { + cfg := DefaultConfig() + var mu sync.Mutex + var bodies [][]byte + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + b, _ := io.ReadAll(r.Body) + mu.Lock() + bodies = append(bodies, b) + mu.Unlock() + w.WriteHeader(http.StatusOK) + })) + defer server.Close() + + exporter := newTelemetryExporter(server.URL, "v-test", "ua", &http.Client{Timeout: 5 * time.Second}, cfg) + aggregator := newMetricsAggregator(exporter, cfg) + interceptor := newInterceptor(aggregator, true) + + interceptor.RecordConnectionConfig(context.Background(), "sess-1", &DriverConnectionParameters{ + Mode: "SEA", AuthMech: "PAT", EnableArrow: true, + }) + + // The connection metric flushes immediately; give the async export a moment. + waitForBodies(t, &mu, &bodies, 1) + + found := false + mu.Lock() + for _, b := range bodies { + for _, ev := range decodeEvents(t, b) { + if ev.DriverConnectionParameters != nil && ev.DriverConnectionParameters.Mode == "SEA" { + found = true + } + } + } + mu.Unlock() + if !found { + t.Fatal("connection-config event with Mode=SEA never reached the endpoint") + } +} + +// TestConnectionScopedError_NotDropped pins the routing of the aggregator's +// defensive "error" arm: a non-terminal error with no buffered statement to attach +// to (empty statementID) is emitted standalone and flushes immediately, rather than +// being dropped. No exported API emits a standalone "error" metric today, so this +// drives recordMetric directly — it guards the routing for a future producer, not a +// currently-reachable production path. +func TestConnectionScopedError_NotDropped(t *testing.T) { + cfg := DefaultConfig() + var mu sync.Mutex + var bodies [][]byte + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + b, _ := io.ReadAll(r.Body) + mu.Lock() + bodies = append(bodies, b) + mu.Unlock() + w.WriteHeader(http.StatusOK) + })) + defer server.Close() + + exporter := newTelemetryExporter(server.URL, "v-test", "ua", &http.Client{Timeout: 5 * time.Second}, cfg) + aggregator := newMetricsAggregator(exporter, cfg) + + // A non-terminal error with no statement to attach to (connection-scoped). + aggregator.recordMetric(context.Background(), &telemetryMetric{ + metricType: "error", + timestamp: time.Now(), + sessionID: "sess-1", + errorType: "sqlstate_42P01", // a SQL error, non-terminal + }) + + waitForBodies(t, &mu, &bodies, 1) + + found := false + mu.Lock() + for _, b := range bodies { + for _, ev := range decodeEvents(t, b) { + if ev.ErrorInfo != nil && ev.ErrorInfo.ErrorName == "sqlstate_42P01" { + found = true + } + } + } + mu.Unlock() + if !found { + t.Fatal("connection-scoped error was silently dropped — expected it to flush and land") + } +} + +// waitForBodies polls until at least n request bodies have been captured or times out. +func waitForBodies(t *testing.T, mu *sync.Mutex, bodies *[][]byte, n int) { + t.Helper() + deadline := time.Now().Add(3 * time.Second) + for { + mu.Lock() + got := len(*bodies) + mu.Unlock() + if got >= n { + return + } + if time.Now().After(deadline) { + t.Fatalf("timed out waiting for %d request body(ies), got %d", n, got) + } + time.Sleep(10 * time.Millisecond) + } +} diff --git a/telemetry/exporter.go b/telemetry/exporter.go index 3e16793a..d956f2ae 100644 --- a/telemetry/exporter.go +++ b/telemetry/exporter.go @@ -42,6 +42,12 @@ type telemetryMetric struct { latencyMs int64 errorType string tags map[string]interface{} + // connParams carries connection-configuration telemetry (a "connection" + // metric). nil for every statement/operation/error metric, so a metric that + // does not set it serializes byte-identically to before (see + // createTelemetryRequest). Populated by Interceptor.RecordConnectionConfig on + // the kernel connect path. + connParams *DriverConnectionParameters } // ensureHTTPScheme adds https:// prefix to host if no scheme is present. diff --git a/telemetry/interceptor.go b/telemetry/interceptor.go index a2abfa92..9e6552ae 100644 --- a/telemetry/interceptor.go +++ b/telemetry/interceptor.go @@ -172,24 +172,29 @@ func (i *Interceptor) AddTag(ctx context.Context, key string, value interface{}) } } -// recordConnection records a connection event. -// -//nolint:unused // Will be used in Phase 8+ -func (i *Interceptor) recordConnection(ctx context.Context, tags map[string]interface{}) { - if !i.enabled { +// RecordConnectionConfig records a connection-configuration event: the connection +// parameters (mode, auth mechanism/flow, proxy usage, arrow, query tags, …) the +// backend connected with. It fires once per connection, at connect, carrying only +// the session id (no statement id — this is connection-scoped, not per-statement). +// A "connection" metric flushes immediately (the connection may close before the +// next batch). Panic-firewalled and gated on enabled, like the other Record* +// methods. Exported for use by the driver package. +func (i *Interceptor) RecordConnectionConfig(ctx context.Context, sessionID string, params *DriverConnectionParameters) { + if !i.enabled || params == nil { return } defer func() { if r := recover(); r != nil { - logger.Debug().Msgf("telemetry: recordConnection panic: %v", r) + logger.Debug().Msgf("telemetry: recordConnectionConfig panic: %v", r) } }() metric := &telemetryMetric{ metricType: "connection", timestamp: time.Now(), - tags: tags, + sessionID: sessionID, + connParams: params, } i.aggregator.recordMetric(ctx, metric) diff --git a/telemetry/request.go b/telemetry/request.go index 6390a363..26ffa0a3 100644 --- a/telemetry/request.go +++ b/telemetry/request.go @@ -91,6 +91,15 @@ type DriverConnectionParameters struct { QueryTags string `json:"query_tags,omitempty"` EnableMetricViewMeta bool `json:"enable_metric_view_metadata,omitempty"` SocketTimeout int64 `json:"socket_timeout,omitempty"` + // Retry policy and CloudFetch memory cap resolved for the connection. Populated + // on the kernel path (the Thrift path leaves them zero for now); they surface + // what a customer actually configured so a hung-connect or large-result-OOM + // report can be diagnosed from telemetry, not just a debug log. RetryMaxAttempts + // / RetryOverallTimeoutMs zero means the backend default is in effect; + // MaxChunksInMemory zero means the kernel default (16). + RetryMaxAttempts int32 `json:"retry_max_attempts,omitempty"` + RetryOverallTimeoutMs int64 `json:"retry_overall_timeout_ms,omitempty"` + MaxChunksInMemory int32 `json:"max_chunks_in_memory,omitempty"` } // SQLExecutionEvent maps to SqlExecutionEvent in the proto schema. @@ -218,6 +227,15 @@ func createTelemetryRequest(metrics []*telemetryMetric, driverVersion string) (* } } + // Add connection-configuration params if present (a "connection" metric). + // Only set when populated, so a statement/operation/error metric — which + // leaves connParams nil — serializes exactly as before. AuthType mirrors the + // AuthMech onto the top-level event field the schema also carries. + if metric.connParams != nil { + frontendLog.Entry.SQLDriverLog.DriverConnectionParameters = metric.connParams + frontendLog.Entry.SQLDriverLog.AuthType = metric.connParams.AuthMech + } + jsonBytes, err := json.Marshal(frontendLog) if err != nil { return nil, err