diff --git a/auth/oauth/m2m/m2m.go b/auth/oauth/m2m/m2m.go index 9e7c9c20..d08b6b89 100644 --- a/auth/oauth/m2m/m2m.go +++ b/auth/oauth/m2m/m2m.go @@ -8,7 +8,7 @@ import ( "github.com/databricks/databricks-sql-go/auth" "github.com/databricks/databricks-sql-go/auth/oauth" - "github.com/rs/zerolog/log" + "github.com/databricks/databricks-sql-go/logger" "golang.org/x/oauth2" "golang.org/x/oauth2/clientcredentials" ) @@ -72,12 +72,14 @@ func (c *authClient) Authenticate(r *http.Request) error { c.tokenSource = GetTokenSource(config) token, err := c.tokenSource.Token() - log.Debug().Msgf("databricks OAuth token fetched successfully") if err != nil { - log.Err(err).Msg("failed to get token") + logger.Err(err).Msg("failed to get token") return err } + // Log via the driver's configurable logger (defaults to Warn) rather than the + // global zerolog logger, so this line honors SetLogLevel like every other. + logger.Debug().Msgf("databricks OAuth token fetched successfully") token.SetAuthHeader(r) return nil diff --git a/auth/oauth/oauth.go b/auth/oauth/oauth.go index 0525e8cd..d0fd3f66 100644 --- a/auth/oauth/oauth.go +++ b/auth/oauth/oauth.go @@ -11,7 +11,7 @@ import ( "time" "github.com/coreos/go-oidc/v3/oidc" - "github.com/rs/zerolog/log" + "github.com/databricks/databricks-sql-go/logger" "golang.org/x/oauth2" ) @@ -99,7 +99,7 @@ func resolveOIDCIssuer(ctx context.Context, client *http.Client, hostName string cfgURL := fmt.Sprintf("https://%s/.well-known/databricks-config", hostName) meta, ok := fetchHostMetadata(ctx, client, cfgURL) if !ok || meta.OIDCEndpoint == "" { - log.Debug().Msgf("oauth: no usable databricks-config for %q; using bare-host OIDC issuer", hostName) + logger.Debug().Msgf("oauth: no usable databricks-config for %q; using bare-host OIDC issuer", hostName) return fallback } @@ -107,13 +107,13 @@ func resolveOIDCIssuer(ctx context.Context, client *http.Client, hostName string // placeholder would resolve to a malformed ".../accounts/" issuer. Fall back // rather than emit it (the function's documented contract). if strings.Contains(meta.OIDCEndpoint, accountIDPlaceholder) && meta.AccountID == "" { - log.Warn().Msgf("oauth: databricks-config for %q has an %s placeholder but empty account_id; using bare-host OIDC issuer", hostName, accountIDPlaceholder) + logger.Warn().Msgf("oauth: databricks-config for %q has an %s placeholder but empty account_id; using bare-host OIDC issuer", hostName, accountIDPlaceholder) return fallback } issuer := substituteAccountID(meta) if !isValidDatabricksIssuer(issuer) { - log.Warn().Msgf("oauth: databricks-config for %q advertised an unusable oidc_endpoint %q; using bare-host OIDC issuer", hostName, issuer) + logger.Warn().Msgf("oauth: databricks-config for %q advertised an unusable oidc_endpoint %q; using bare-host OIDC issuer", hostName, issuer) return fallback } return issuer diff --git a/auth/oauth/oauth_test.go b/auth/oauth/oauth_test.go index e9e64291..96d71a2f 100644 --- a/auth/oauth/oauth_test.go +++ b/auth/oauth/oauth_test.go @@ -4,6 +4,7 @@ import ( "context" "net/http" "net/http/httptest" + "reflect" "strings" "testing" ) @@ -62,6 +63,28 @@ func TestGetAzureDnsZone(t *testing.T) { } } +// TestGetScopes pins the literal cloud-specific scope sets (not GetScopes's own +// output) so a regression in the AWS/GCP-vs-Azure branching — the source of the +// kernel U2M access_denied incident — is caught rather than asserted tautologically. +func TestGetScopes(t *testing.T) { + cases := []struct { + name string + host string + want []string + }{ + {"AWS", "dbc-1234.cloud.databricks.com", []string{"offline_access", "sql"}}, + {"GCP", "x.gcp.databricks.com", []string{"offline_access", "sql"}}, + {"Azure", "adb-123.azuredatabricks.net", []string{"offline_access", "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d/user_impersonation"}}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if got := GetScopes(tc.host, nil); !reflect.DeepEqual(got, tc.want) { + t.Fatalf("GetScopes(%q, nil) = %v, want %v", tc.host, got, tc.want) + } + }) + } +} + // TestResolveOIDCIssuer drives resolveOIDCIssuer end-to-end against an httptest // server: the server stands in for the connection host, so a 200 with a given // databricks-config body exercises the real fetch + substitution + validation + diff --git a/auth/oauth/u2m/authenticator.go b/auth/oauth/u2m/authenticator.go index 1d5b46bf..834265c0 100644 --- a/auth/oauth/u2m/authenticator.go +++ b/auth/oauth/u2m/authenticator.go @@ -20,7 +20,7 @@ import ( "github.com/databricks/databricks-sql-go/auth" "github.com/databricks/databricks-sql-go/auth/oauth" - "github.com/rs/zerolog/log" + "github.com/databricks/databricks-sql-go/logger" "golang.org/x/oauth2" ) @@ -84,15 +84,9 @@ type u2mAuthenticator struct { // mode. It structurally satisfies the U2MCredentialsProvider interface the kernel // backend asserts (defined in internal/backend/kernel, off the public API). // -// Only the client id crosses to the kernel — NOT the scopes. The Thrift path -// requests offline_access + sql (AWS/GCP) or offline_access + / -// user_impersonation (Azure) via oauth.GetScopes, while the kernel's U2M flow -// applies its own default set (all-apis + offline_access). Neither backend exposes -// a user-facing U2M-scopes option, so this is a fixed-vs-fixed difference, not a -// dropped caller choice; against the built-in databricks-sql-connector public -// client (which all-apis targets) both authorize successfully. If a U2M-scopes -// option is ever added, forward it via kernel.Auth.Scopes (already wired through -// set_auth_u2m) so the two paths request the same set. +// Only the client id is read here; resolveKernelAuth pairs it with the same +// oauth.GetScopes set the Thrift path requests, so both backends authorize against +// the built-in databricks-sql-connector client identically (not the kernel default). func (c *u2mAuthenticator) U2MClientID() string { return c.clientID } // Auth will start the OAuth Authorization Flow to authenticate the cli client @@ -160,7 +154,7 @@ func (tsp *tokenSourceProvider) GetTokenSource() (oauth2.TokenSource, error) { loginURL := tsp.config.AuthCodeURL(state, challenge, challengeMethod) tsp.state = state - log.Info().Msgf("listening on %s://%s/", tsp.redirectURL.Scheme, tsp.redirectURL.Host) + logger.Info().Msgf("listening on %s://%s/", tsp.redirectURL.Scheme, tsp.redirectURL.Host) listener, err := net.Listen("tcp", tsp.redirectURL.Host) if err != nil { return nil, err @@ -225,28 +219,28 @@ func (tsp *tokenSourceProvider) ServeHTTP(w http.ResponseWriter, r *http.Request // Do some checking of the response here to show more relevant content if resp.err != "" { - log.Error().Msg(resp.err) + logger.Error().Msg(resp.err) w.WriteHeader(http.StatusBadRequest) _, err := w.Write([]byte(errorHTML("Identity Provider returned an error: " + resp.err))) //nolint:gosec // XSS not a concern for local OAuth callback if err != nil { - log.Error().Err(err).Msg("unable to write error response") + logger.Error().Err(err).Msg("unable to write error response") } return } if resp.state != tsp.state && r.URL.String() != "/favicon.ico" { msg := "Authentication state received did not match original request. Please try to login again." - log.Error().Msg(msg) + logger.Error().Msg(msg) w.WriteHeader(http.StatusBadRequest) _, err := w.Write([]byte(errorHTML(msg))) if err != nil { - log.Error().Err(err).Msg("unable to write error response") + logger.Error().Err(err).Msg("unable to write error response") } return } _, err := w.Write([]byte(infoHTML("CLI Login Success", "You may close this window anytime now and go back to terminal"))) if err != nil { - log.Error().Err(err).Msg("unable to write success response") + logger.Error().Err(err).Msg("unable to write success response") } } diff --git a/auth/tokenprovider/authenticator.go b/auth/tokenprovider/authenticator.go index b4ef57ff..9452811c 100644 --- a/auth/tokenprovider/authenticator.go +++ b/auth/tokenprovider/authenticator.go @@ -6,7 +6,7 @@ import ( "net/http" "github.com/databricks/databricks-sql-go/auth" - "github.com/rs/zerolog/log" + "github.com/databricks/databricks-sql-go/logger" ) // TokenProviderAuthenticator implements auth.Authenticator using a TokenProvider. @@ -47,7 +47,7 @@ func (a *TokenProviderAuthenticator) Authenticate(r *http.Request) error { } token.SetAuthHeader(r) - log.Debug().Msgf("token provider authenticator: authenticated using provider %s", a.provider.Name()) + logger.Debug().Msgf("token provider authenticator: authenticated using provider %s", a.provider.Name()) return nil } diff --git a/auth/tokenprovider/cached.go b/auth/tokenprovider/cached.go index fdb86c89..5c4de8b7 100644 --- a/auth/tokenprovider/cached.go +++ b/auth/tokenprovider/cached.go @@ -6,7 +6,7 @@ import ( "sync" "time" - "github.com/rs/zerolog/log" + "github.com/databricks/databricks-sql-go/logger" ) // CachedTokenProvider wraps another provider and caches tokens @@ -43,7 +43,7 @@ func (p *CachedTokenProvider) GetToken(ctx context.Context) (*Token, error) { // If cache is valid and not being refreshed, return a copy if cached != nil && !needsRefresh { - log.Debug().Msgf("cached token provider: using cached token for provider %s", p.provider.Name()) + logger.Debug().Msgf("cached token provider: using cached token for provider %s", p.provider.Name()) // Return a copy to avoid concurrent modification issues return copyToken(cached), nil } @@ -75,7 +75,7 @@ func (p *CachedTokenProvider) GetToken(ctx context.Context) (*Token, error) { p.mutex.Unlock() // Fetch new token WITHOUT holding the lock - log.Debug().Msgf("cached token provider: fetching new token from provider %s", p.provider.Name()) + logger.Debug().Msgf("cached token provider: fetching new token from provider %s", p.provider.Name()) token, err := p.provider.GetToken(ctx) // Update cache with result diff --git a/auth/tokenprovider/exchange.go b/auth/tokenprovider/exchange.go index 09c9268d..c0c5161e 100644 --- a/auth/tokenprovider/exchange.go +++ b/auth/tokenprovider/exchange.go @@ -10,8 +10,8 @@ import ( "strings" "time" + "github.com/databricks/databricks-sql-go/logger" "github.com/golang-jwt/jwt/v5" - "github.com/rs/zerolog/log" ) // FederationProvider wraps another token provider and automatically handles token exchange @@ -61,16 +61,16 @@ func (p *FederationProvider) GetToken(ctx context.Context) (*Token, error) { // Check if token is a JWT and needs exchange if p.needsTokenExchange(baseToken.AccessToken) { - log.Debug().Msgf("federation provider: attempting token exchange for %s", p.baseProvider.Name()) + logger.Debug().Msgf("federation provider: attempting token exchange for %s", p.baseProvider.Name()) // Try token exchange exchangedToken, err := p.tryTokenExchange(ctx, baseToken.AccessToken) if err != nil { - log.Warn().Err(err).Msg("federation provider: token exchange failed, using original token") + logger.Warn().Err(err).Msg("federation provider: token exchange failed, using original token") return baseToken, nil // Fall back to original token } - log.Debug().Msg("federation provider: token exchange successful") + logger.Debug().Msg("federation provider: token exchange successful") return exchangedToken, nil } @@ -87,7 +87,7 @@ func (p *FederationProvider) needsTokenExchange(tokenString string) bool { // 3. Token validation will be done by Databricks during exchange token, _, err := new(jwt.Parser).ParseUnverified(tokenString, jwt.MapClaims{}) if err != nil { - log.Debug().Err(err).Msg("federation provider: not a JWT token, skipping exchange") + logger.Debug().Err(err).Msg("federation provider: not a JWT token, skipping exchange") return false } @@ -114,7 +114,7 @@ func (p *FederationProvider) tryTokenExchange(ctx context.Context, subjectToken exchangeURL = "https://" + exchangeURL } else if strings.HasPrefix(exchangeURL, "http://") { // Warn if using insecure HTTP for token exchange - log.Warn().Msgf("federation provider: using insecure HTTP for token exchange: %s", exchangeURL) + logger.Warn().Msgf("federation provider: using insecure HTTP for token exchange: %s", exchangeURL) } if !strings.HasSuffix(exchangeURL, "/") { exchangeURL += "/" @@ -179,7 +179,7 @@ func (p *FederationProvider) tryTokenExchange(ctx context.Context, subjectToken return nil, fmt.Errorf("token exchange returned empty access token") } if tokenResp.TokenType == "" { - log.Debug().Msg("token exchange: token_type not specified, defaulting to Bearer") + logger.Debug().Msg("token exchange: token_type not specified, defaulting to Bearer") tokenResp.TokenType = "Bearer" } if tokenResp.ExpiresIn < 0 { @@ -211,7 +211,7 @@ func (p *FederationProvider) isSameHost(url1, url2 string) bool { u2, err2 := url.Parse(parsedURL2) if err1 != nil || err2 != nil { - log.Debug().Msgf("federation provider: failed to parse URLs for comparison: url1=%s err1=%v, url2=%s err2=%v", + logger.Debug().Msgf("federation provider: failed to parse URLs for comparison: url1=%s err1=%v, url2=%s err2=%v", url1, err1, parsedURL2, err2) return false } @@ -219,7 +219,7 @@ func (p *FederationProvider) isSameHost(url1, url2 string) bool { // Use Hostname() instead of Host to ignore port differences // This handles cases like "host.com:443" == "host.com" for HTTPS isSame := u1.Hostname() == u2.Hostname() - log.Debug().Msgf("federation provider: host comparison: %s vs %s = %v", u1.Hostname(), u2.Hostname(), isSame) + logger.Debug().Msgf("federation provider: host comparison: %s vs %s = %v", u1.Hostname(), u2.Hostname(), isSame) return isSame } diff --git a/internal/arrowscan/arrowscan.go b/internal/arrowscan/arrowscan.go index cee634e2..cb24b0a8 100644 --- a/internal/arrowscan/arrowscan.go +++ b/internal/arrowscan/arrowscan.go @@ -173,7 +173,15 @@ func ScanCellCached(col arrow.Array, row int, loc *time.Location, keys *StructKe // it. The cross-backend parity test pins both sides to µs (matching the wire // reality), so it can't exercise a non-µs value; TestScanCellTimestampUnits // covers this arm's unit-correctness directly instead. - return inLocation(c.Value(row).ToTime(dt.Unit), loc), nil + // + // Convert via the unit-specific time.Unix* constructors rather than arrow's + // ToTime: ToTime forms an int64-nanosecond intermediate (value*multiplier), + // which overflows for TIMESTAMPs outside ~1678–2262 and silently wraps a valid + // instant to a wrong one (e.g. TIMESTAMP '0001-01-01' → 1754). time.UnixMicro/ + // UnixMilli/Unix split into (sec, nsec) internally, so the full 0001–9999 range + // round-trips. The Thrift default path never hits this — it receives TIMESTAMP + // as a preformatted server string — so this divergence is kernel-only. + return inLocation(timestampToTime(int64(c.Value(row)), dt.Unit), loc), nil case *array.Decimal128: dt := col.DataType().(*arrow.Decimal128Type) return decimalfmt.ExactString(c.Value(row), dt.Scale), nil @@ -185,6 +193,13 @@ func ScanCellCached(col arrow.Array, row int, loc *time.Location, keys *StructKe // shared renderer to reuse, and this stays kernel-side). dt := col.DataType().(*arrow.DurationType) return formatDayTimeInterval(int64(c.Value(row)), dt.Unit), nil + case *array.DayTimeInterval: + // INTERVAL DAY TO SECOND can also arrive as arrow's DayTimeInterval ({Days, + // Milliseconds}) rather than a Duration. Fold it to total milliseconds and reuse + // the same renderer. int32 days × 86_400_000 stays well within int64, so the + // day→ms scaling cannot overflow. + v := c.Value(row) + return formatDayTimeInterval(int64(v.Days)*86_400_000+int64(v.Milliseconds), arrow.Millisecond), nil case *array.MonthInterval: // INTERVAL YEAR TO MONTH arrives as a month count; Thrift's server string is // "years-months". @@ -276,6 +291,24 @@ func abs64(x int64) int64 { return x } +// timestampToTime converts an arrow timestamp value in the given unit to a UTC +// time.Time without arrow's ToTime int64-nanosecond overflow (see the *array.Timestamp +// scan arm). Each constructor takes its argument in native units and splits into +// (seconds, nanoseconds) internally, so the whole 0001–9999 range is representable; +// only the nanosecond arm passes ns directly (already the finest unit, cannot overflow). +func timestampToTime(v int64, unit arrow.TimeUnit) time.Time { + switch unit { + case arrow.Second: + return time.Unix(v, 0).UTC() + case arrow.Millisecond: + return time.UnixMilli(v).UTC() + case arrow.Microsecond: + return time.UnixMicro(v).UTC() + default: // arrow.Nanosecond + return time.Unix(0, v).UTC() + } +} + // inLocation renders t in loc, matching the Thrift path's .In(location); a nil // loc leaves the value in UTC (arrow's ToTime default). func inLocation(t time.Time, loc *time.Location) time.Time { diff --git a/internal/arrowscan/arrowscan_test.go b/internal/arrowscan/arrowscan_test.go index 957a8f47..b8d31bc3 100644 --- a/internal/arrowscan/arrowscan_test.go +++ b/internal/arrowscan/arrowscan_test.go @@ -193,6 +193,34 @@ func TestScanCellInterval(t *testing.T) { } }) } + + // DAY-TIME intervals can arrive as arrow's DayTimeInterval ({Days, Milliseconds}) + // instead of a Duration; it must render to the identical "D HH:MM:SS.nnnnnnnnn". + dayTimeInterval := []struct { + name string + v arrow.DayTimeInterval + want string + }{ + {"one_day", arrow.DayTimeInterval{Days: 1, Milliseconds: 0}, "1 00:00:00.000000000"}, + {"day_hms_millis", arrow.DayTimeInterval{Days: 1, Milliseconds: 3661_500}, "1 01:01:01.500000000"}, + {"negative", arrow.DayTimeInterval{Days: -1, Milliseconds: -3661_500}, "-1 01:01:01.500000000"}, + } + for _, tc := range dayTimeInterval { + t.Run("daytimeinterval_"+tc.name, func(t *testing.T) { + b := array.NewDayTimeIntervalBuilder(pool) + defer b.Release() + b.Append(tc.v) + arr := b.NewArray() + defer arr.Release() + v, err := ScanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + if v.(string) != tc.want { + t.Errorf("got %q, want %q", v, tc.want) + } + }) + } } // ScanCell renders DATE / TIMESTAMP in the requested location, matching the @@ -274,6 +302,40 @@ func TestScanCellTimestampUnits(t *testing.T) { } } +// TestScanCellTimestampOutOfNanoRange pins the fix for the microsecond→nanosecond +// overflow: arrow's ToTime forms value*multiplier as an int64 ns count, which +// overflows for instants outside ~1678–2262 and silently wraps to a wrong time +// (TIMESTAMP '0001-01-01' scanned back as 1754, '9999-12-31' as 1816). The +// unit-split constructors used by ScanCell must round-trip the full 0001–9999 +// range that Databricks TIMESTAMP allows. +func TestScanCellTimestampOutOfNanoRange(t *testing.T) { + pool := memory.NewGoAllocator() + cases := []time.Time{ + time.Date(1, time.January, 1, 0, 0, 0, 0, time.UTC), // TIMESTAMP '0001-01-01 00:00:00' + time.Date(9999, time.December, 31, 23, 59, 59, 0, time.UTC), // TIMESTAMP '9999-12-31 23:59:59' + time.Date(1677, time.January, 1, 0, 0, 0, 0, time.UTC), // just below the int64-ns floor + time.Date(2263, time.January, 1, 0, 0, 0, 0, time.UTC), // just above the int64-ns ceiling + } + for _, want := range cases { + t.Run(want.Format("2006-01-02"), func(t *testing.T) { + // Databricks TIMESTAMP is always microseconds on the wire. + b := array.NewTimestampBuilder(pool, &arrow.TimestampType{Unit: arrow.Microsecond}) + defer b.Release() + b.Append(arrow.Timestamp(want.UnixMicro())) + arr := b.NewArray() + defer arr.Release() + + v, err := ScanCell(arr, 0, nil) + if err != nil { + t.Fatal(err) + } + if got := v.(time.Time); !got.Equal(want) { + t.Errorf("got %v, want %v (arrow ToTime would have wrapped this)", got.UTC(), want) + } + }) + } +} + // ScanCell renders nested types (list/struct/map) to a JSON string matching the // Thrift path. func TestScanCellNested(t *testing.T) { diff --git a/internal/arrowscan/coltype.go b/internal/arrowscan/coltype.go index e795fcfc..a0a5548e 100644 --- a/internal/arrowscan/coltype.go +++ b/internal/arrowscan/coltype.go @@ -10,32 +10,18 @@ import ( ) // 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. +// sql.ColumnType. The kernel derives it from the result's Arrow schema via +// ColumnTypeInfoFor — the mapping the value scanner and type reporter share, kept +// byte-identical to the Thrift backend (guarded by the coltype parity tests). 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 is the Databricks type name (e.g. "BIGINT", "DECIMAL"), + // matching the Thrift path; "" 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. + // matching the Thrift path. 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 / HasLength report a variable-length column's unbounded length + // (math.MaxInt64), matching Thrift; fixed-width types report (0, false). Length int64 HasLength bool } @@ -58,14 +44,8 @@ var ( // 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. +// here are exactly those ScanCellCached scans, so a column's reported type and its +// scanned value stay in lockstep. func ColumnTypeInfoFor(dt arrow.DataType) ColumnTypeInfo { switch dt.ID() { case arrow.BOOL: @@ -79,11 +59,8 @@ func ColumnTypeInfoFor(dt arrow.DataType) ColumnTypeInfo { 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. + // Databricks SQL has no unsigned types; defensive arm matching ScanCellCached, + // which widens unsigned ints to int64 (driver.Value has no uint64). return ColumnTypeInfo{DatabaseTypeName: "BIGINT", ScanType: scanTypeInt64} case arrow.FLOAT32: return ColumnTypeInfo{DatabaseTypeName: "FLOAT", ScanType: scanTypeFloat32} @@ -100,10 +77,8 @@ func ColumnTypeInfoFor(dt arrow.DataType) ColumnTypeInfo { 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). + // Match Thrift's sql.RawBytes scan type even though the kernel renders the + // value as an exact string — both convert cleanly to a caller's *string/*[]byte. return ColumnTypeInfo{DatabaseTypeName: "DECIMAL", ScanType: scanTypeRawBytes} case arrow.LIST, arrow.LARGE_LIST, arrow.FIXED_SIZE_LIST: return varLen("ARRAY", scanTypeRawBytes) @@ -111,25 +86,20 @@ func ColumnTypeInfoFor(dt arrow.DataType) ColumnTypeInfo { 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. + case arrow.DURATION, arrow.INTERVAL_DAY_TIME: + // INTERVAL DAY TO SECOND: Thrift's prod default pre-formats intervals to text + // and declares STRING_TYPE (verified live), so match STRING; the kernel returns + // it as either arrow.DURATION or arrow.INTERVAL_DAY_TIME, both formatted Go-side + // to the identical string (see ScanCellCached). Without INTERVAL_DAY_TIME here + // its metadata fell through to the default ""/interface{} (comparator diff). 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} + // VOID/NULL columns: the server stringifies them over Thrift (verified live: + // even bare SELECT NULL reports STRING), same as intervals — so match STRING. + return varLen("STRING", scanTypeString) default: // A type ScanCellCached does not handle: report the Thrift default scan type // (*interface{}) and no database name, rather than inventing one. @@ -137,9 +107,8 @@ func ColumnTypeInfoFor(dt arrow.DataType) ColumnTypeInfo { } } -// 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. +// varLen builds a ColumnTypeInfo for a variable-length type, reporting the +// unbounded length (math.MaxInt64) the Thrift path uses for such 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 index ac71ad0c..0dbd2f8a 100644 --- a/internal/arrowscan/coltype_test.go +++ b/internal/arrowscan/coltype_test.go @@ -10,13 +10,10 @@ import ( "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. +// TestColumnTypeInfoFor pins the Arrow → column-metadata mapping so it stays +// byte-identical to the Thrift path — a pure-Go guard the value-parity suites miss +// (they compare scanned values, not ColumnType metadata). Expected values are +// ground truth captured from Thrift on a live warehouse. func TestColumnTypeInfoFor(t *testing.T) { str := reflect.TypeOf("") raw := reflect.TypeOf(sql.RawBytes{}) @@ -52,8 +49,10 @@ func TestColumnTypeInfoFor(t *testing.T) { {"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}, + {"day_time_interval", arrow.FixedWidthTypes.DayTimeInterval, "STRING", str, math.MaxInt64, true}, {"month_interval", arrow.FixedWidthTypes.MonthInterval, "STRING", str, math.MaxInt64, true}, - {"null", arrow.Null, "NULL", nil, 0, false}, + // VOID/NULL: Thrift stringifies to STRING on the wire (verified live), like intervals. + {"null", arrow.Null, "STRING", str, math.MaxInt64, true}, } for _, c := range cases { t.Run(c.name, func(t *testing.T) { @@ -72,12 +71,8 @@ func TestColumnTypeInfoFor(t *testing.T) { } // 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). +// ScanCellCached handles must have a non-fallback ColumnTypeInfoFor entry, so a new +// scanner type without matching metadata is caught here, not at runtime. func TestColumnTypeInfoScanTypeCoversScanner(t *testing.T) { unknown := reflect.TypeOf(new(any)) scannerTypes := []arrow.DataType{ @@ -94,6 +89,7 @@ func TestColumnTypeInfoScanTypeCoversScanner(t *testing.T) { 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.DayTimeInterval, arrow.FixedWidthTypes.MonthInterval, } for _, dt := range scannerTypes { diff --git a/internal/backend/kernel/affectedrows.go b/internal/backend/kernel/affectedrows.go new file mode 100644 index 00000000..ca2bd564 --- /dev/null +++ b/internal/backend/kernel/affectedrows.go @@ -0,0 +1,21 @@ +package kernel + +// This file is intentionally NOT behind the `cgo && databricks_kernel` build tag. +// It holds the pure normalization of the C ABI's affected-row count so the rule is +// unit-testable in the default CGO_ENABLED=0 build (like paramBindArg in bindparams.go), +// separate from the cgo call site in operation.go. + +// normalizeAffectedRows maps the kernel C ABI's modified-row count onto the value +// database/sql's Result.RowsAffected() reports, matching the Thrift path. +// +// The C ABI returns -1 for "not applicable / unknown": DDL, SELECT, or a warehouse +// that doesn't surface the counter (the kernel's num_modified_rows Option is +// None). The Thrift path reports 0 in exactly those cases (TGetOperationStatusResp +// defaults NumModifiedRows to 0), so the -1 sentinel is folded to 0 for parity. +// A real DML count (>= 0) passes through unchanged. +func normalizeAffectedRows(n int64) int64 { + if n < 0 { + return 0 + } + return n +} diff --git a/internal/backend/kernel/affectedrows_test.go b/internal/backend/kernel/affectedrows_test.go new file mode 100644 index 00000000..9df1b5f0 --- /dev/null +++ b/internal/backend/kernel/affectedrows_test.go @@ -0,0 +1,24 @@ +package kernel + +import "testing" + +func TestNormalizeAffectedRows(t *testing.T) { + cases := []struct { + name string + in int64 + want int64 + }{ + {"ddl or select sentinel (-1) folds to Thrift's 0", -1, 0}, + {"any negative folds to 0", -42, 0}, + {"zero real count is preserved", 0, 0}, + {"positive DML count passes through", 7, 7}, + {"large DML count passes through", 1_000_000, 1_000_000}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if got := normalizeAffectedRows(tc.in); got != tc.want { + t.Errorf("normalizeAffectedRows(%d) = %d, want %d", tc.in, got, tc.want) + } + }) + } +} diff --git a/internal/backend/kernel/auth.go b/internal/backend/kernel/auth.go index 29563b28..ddef6510 100644 --- a/internal/backend/kernel/auth.go +++ b/internal/backend/kernel/auth.go @@ -20,21 +20,19 @@ const ( // validateKernelConfig); OpenSession maps it to exactly one // kernel_session_config_set_auth_* call. // Scopes and RedirectPort map to the optional args of set_auth_u2m and are wired -// through to it by setAuth, but no Go path populates them today: the driver exposes -// no user option for U2M scopes or redirect port on either backend (the native -// Thrift path hardcodes both), so resolveKernelAuth leaves them zero and the kernel -// applies its defaults. They are kept — rather than dropped and the setter hardcoded -// to NULL/0 — so kernel.Auth models the full set_auth_u2m surface: adding a future -// WithOAuthRedirectPort / scopes option becomes populating these, not re-plumbing -// the setter. TestSetAuthByMode's "U2M full" case pins that marshalling so the -// dormant path stays correct. +// through to it by setAuth. resolveKernelAuth populates Scopes with the same +// cloud-specific set the Thrift path requests (via oauth.GetScopes) so both +// backends authorize identically; RedirectPort stays zero (no user option, kernel +// default 8020) but is kept so kernel.Auth models the full set_auth_u2m surface — +// a future WithOAuthRedirectPort becomes populating it, not re-plumbing the setter. +// TestSetAuthByMode's "U2M full" case pins the marshalling of both. type Auth struct { Mode AuthMode Token string // PAT ClientID string // M2M + U2M (U2M: the cloud-inferred Go client id) ClientSecret string // M2M - Scopes []string // U2M — dormant (see note above); nil → kernel default scopes - RedirectPort uint16 // U2M — dormant (see note above); 0 → kernel default port (8020) + Scopes []string // U2M — Thrift-parity scopes from oauth.GetScopes; nil → kernel default + RedirectPort uint16 // U2M — no user option today; 0 → kernel default port (8020) } // M2MCredentialsProvider is implemented by the OAuth M2M authenticator to expose diff --git a/internal/backend/kernel/backend.go b/internal/backend/kernel/backend.go index febf37e4..5dedc4b3 100644 --- a/internal/backend/kernel/backend.go +++ b/internal/backend/kernel/backend.go @@ -132,6 +132,20 @@ func (k *KernelBackend) OpenSession(ctx context.Context) error { return err } + // User-Agent so query history attributes the kernel path to this driver. + if k.cfg.UserAgent != "" { + name := newCStr("User-Agent") + val := newCStr(k.cfg.UserAgent) + errSet := call(func() C.KernelStatusCode { + return C.kernel_session_config_set_custom_header(cfg, name.c, val.c) + }) + name.free() + val.free() + if errSet != nil { + return fmt.Errorf("kernel: set_custom_header[User-Agent]: %w", toConnError(errSet)) + } + } + // TLS: crypto/tls's InsecureSkipVerify accepts any server cert, so relax both // chain validation and the hostname check — mapping only one would leave the // kernel path stricter than the Thrift path it mirrors (a self-signed cert diff --git a/internal/backend/kernel/batchiterator_test.go b/internal/backend/kernel/batchiterator_test.go new file mode 100644 index 00000000..833b60b9 --- /dev/null +++ b/internal/backend/kernel/batchiterator_test.go @@ -0,0 +1,91 @@ +//go:build cgo && databricks_kernel + +package kernel + +import ( + "context" + "errors" + "io" + "testing" + + "github.com/apache/arrow/go/v12/arrow" + dbsqlerr "github.com/databricks/databricks-sql-go/errors" + dbsqlrows "github.com/databricks/databricks-sql-go/internal/rows" + dbrows "github.com/databricks/databricks-sql-go/rows" +) + +// TestGetArrowIPCStreamsRejected pins the discoverable sentinel so callers can +// branch on it: the kernel C ABI exports Arrow C Data, not IPC bytes. +func TestGetArrowIPCStreamsRejected(t *testing.T) { + r := &kernelRows{} + it, err := r.GetArrowIPCStreams(context.Background()) + if it != nil { + t.Errorf("GetArrowIPCStreams = %v iterator, want nil", it) + } + if !errors.Is(err, dbsqlerr.ErrNotSupportedByKernel) { + t.Errorf("GetArrowIPCStreams err = %v, want wrapping ErrNotSupportedByKernel", err) + } +} + +// TestBatchIteratorSchema returns the schema captured at construction, and errors +// (rather than panics) when none was captured. +func TestBatchIteratorSchema(t *testing.T) { + sch := arrow.NewSchema([]arrow.Field{{Name: "c", Type: arrow.PrimitiveTypes.Int64}}, nil) + it := &kernelBatchIterator{r: &kernelRows{schema: sch}} + if got, err := it.Schema(); err != nil || got != sch { + t.Errorf("Schema() = (%v, %v), want (%v, nil)", got, err, sch) + } + itNil := &kernelBatchIterator{r: &kernelRows{}} + if _, err := itNil.Schema(); err == nil { + t.Error("Schema() with no schema = nil error, want a failure") + } +} + +// TestBatchIteratorDrainedStream: once the stream is drained (eof), fill reports +// done, so HasNext is false and Next returns io.EOF without touching the C stream. +func TestBatchIteratorDrainedStream(t *testing.T) { + it := &kernelBatchIterator{r: &kernelRows{eof: true}} + if it.HasNext() { + t.Error("HasNext on a drained stream = true, want false") + } + if rec, err := it.Next(); rec != nil || err != io.EOF { + t.Errorf("Next on a drained stream = (%v, %v), want (nil, io.EOF)", rec, err) + } +} + +// TestBatchIteratorFetchErrorSeedsIterationErr is the regression for the telemetry +// gap: the batch path bypasses Next(), so a fetch error must still seed iterationErr +// or OnClose records the failed stream as a successful statement. A cancelled r.ctx +// makes nextBatch return a non-EOF error at the boundary without a live C stream. +func TestBatchIteratorFetchErrorSeedsIterationErr(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + var gotIter error + r := &kernelRows{ + ctx: ctx, + cancel: cancel, + callbacks: &dbsqlrows.TelemetryCallbacks{ + OnClose: func(_ int64, _ int, iterErr, _ error) { gotIter = iterErr }, + }, + } + it := &kernelBatchIterator{r: r} + + _, err := it.Next() + if err == nil { + t.Fatal("Next with a cancelled ctx = nil error, want the fetch error") + } + // The error is sticky: a second Next re-returns it (it never clears). + if _, err2 := it.Next(); err2 != err { + t.Errorf("second Next = %v, want the same sticky error %v", err2, err) + } + if r.iterationErr == nil { + t.Error("iterationErr not seeded on batch-path fetch error — OnClose would report success") + } + it.Close() + if !errors.Is(gotIter, err) { + t.Errorf("OnClose iterErr = %v, want the fetch error %v", gotIter, err) + } +} + +// compile-time guard: kernelBatchIterator must satisfy the public iterator surface. +var _ dbrows.ArrowBatchIterator = (*kernelBatchIterator)(nil) diff --git a/internal/backend/kernel/config.go b/internal/backend/kernel/config.go index 5727407c..e0577891 100644 --- a/internal/backend/kernel/config.go +++ b/internal/backend/kernel/config.go @@ -17,6 +17,10 @@ type Config struct { WarehouseID string // bare warehouse id; preferred over HTTPPath when set Auth Auth // PAT / OAuth M2M / OAuth U2M + // UserAgent is forwarded as the User-Agent header so the kernel path is + // attributed to this driver (not the kernel's built-in UA). Empty leaves it unset. + UserAgent string + // SessionConf carries server-bound session confs verbatim — the same map the // Thrift backend forwards (STATEMENT_TIMEOUT, QUERY_TAGS, TIMEZONE, …). SessionConf map[string]string diff --git a/internal/backend/kernel/operation.go b/internal/backend/kernel/operation.go index 4e0228b9..cc60a5bc 100644 --- a/internal/backend/kernel/operation.go +++ b/internal/backend/kernel/operation.go @@ -194,7 +194,8 @@ func (k *KernelBackend) execute(ctx context.Context, req backend.ExecRequest) (b // Capture the modified-row count and server query id now, while exec is live — // the operation is closed (nulling exec) before these are read on the // ExecContext path, and the query-id pointer is only valid while exec lives. - op.affectedRows = int64(C.kernel_executed_statement_num_modified_rows(exec)) + // normalizeAffectedRows folds the C ABI's -1 sentinel to 0 for Thrift parity. + op.affectedRows = normalizeAffectedRows(int64(C.kernel_executed_statement_num_modified_rows(exec))) // A nil execErr means the statement reached a terminal state server-side (it // committed). We deliberately do NOT re-check ctx.Err() here to convert a // completed statement into a cancellation: for a non-idempotent DML that would diff --git a/internal/backend/kernel/rows.go b/internal/backend/kernel/rows.go index 2020dbda..d04a884e 100644 --- a/internal/backend/kernel/rows.go +++ b/internal/backend/kernel/rows.go @@ -23,8 +23,11 @@ import ( "github.com/apache/arrow/go/v12/arrow" "github.com/apache/arrow/go/v12/arrow/cdata" + dbsqlerr "github.com/databricks/databricks-sql-go/errors" "github.com/databricks/databricks-sql-go/internal/arrowscan" + context2 "github.com/databricks/databricks-sql-go/internal/compat/context" dbsqlrows "github.com/databricks/databricks-sql-go/internal/rows" + dbrows "github.com/databricks/databricks-sql-go/rows" ) // kernelRows implements the same optional column-type interfaces the Thrift path @@ -37,6 +40,9 @@ var ( _ driver.RowsColumnTypeDatabaseTypeName = (*kernelRows)(nil) _ driver.RowsColumnTypeNullable = (*kernelRows)(nil) _ driver.RowsColumnTypeLength = (*kernelRows)(nil) + // Public batch API, matching the Thrift path (rows.rows). GetArrowBatches is + // real; GetArrowIPCStreams rejects (kernel C ABI exports C Data, not IPC bytes). + _ dbrows.Rows = (*kernelRows)(nil) ) // kernelRows implements driver.Rows over the kernel result stream. It pulls one @@ -48,13 +54,18 @@ var ( // per-batch network/decode work once; row reads walk the already-imported // arrow.Record with O(1) indexing. type kernelRows struct { + // ctx is detached from the caller's QueryContext (via context2.WithoutCancel) + // so a submit-gating deadline can't truncate a large CloudFetch stream, matching + // the Thrift path (see ES-1934053); cancel makes it abortable from Close. ctx context.Context + cancel context.CancelFunc op *kernelOp stream *C.kernel_result_stream_t callbacks *dbsqlrows.TelemetryCallbacks cols []string colTypes []arrowscan.ColumnTypeInfo // per-column type metadata (PECOBLR-3692) + schema *arrow.Schema // result-set schema, for GetArrowBatches().Schema() cur arrow.Record // current batch (nil until first Next) rowInCur int // next row index within cur chunkCount int // cumulative batches fetched, for OnChunkFetched @@ -80,7 +91,10 @@ func newKernelRows(ctx context.Context, op *kernelOp, stream *C.kernel_result_st // telemetry. Assign it only on the success path so cleanup Close() on a // schema/import failure does not record a falsely successful CLOSE_STATEMENT; the // construction error itself is surfaced to and recorded by the conn execute path. - r := &kernelRows{ctx: ctx, op: op, stream: stream, keyCache: arrowscan.NewStructKeyCache()} + // Detach from the caller's cancellation but keep its values (auth/logging), so a + // deadline that only gated statement submission can't truncate the result stream. + resultsCtx, resultsCancel := context.WithCancel(context2.WithoutCancel(ctx)) + r := &kernelRows{ctx: resultsCtx, cancel: resultsCancel, op: op, stream: stream, keyCache: arrowscan.NewStructKeyCache()} var csch C.struct_ArrowSchema if err := call(func() C.KernelStatusCode { @@ -95,6 +109,7 @@ func newKernelRows(ctx context.Context, op *kernelOp, stream *C.kernel_result_st r.Close() return nil, fmt.Errorf("kernel: import schema: %w", err) } + r.schema = sch fields := sch.Fields() r.cols = make([]string, len(fields)) // Derive per-column type metadata from the Arrow schema up front (the same @@ -164,6 +179,10 @@ func (r *kernelRows) Close() error { } r.closed = true closeStart := time.Now() + // Release the detached results context so an aborted Close doesn't leak it. + if r.cancel != nil { + r.cancel() + } if r.cur != nil { r.cur.Release() r.cur = nil @@ -226,14 +245,9 @@ func (r *kernelRows) next(dest []driver.Value) error { // nextBatch pulls the next Arrow batch. A released array (release==NULL) is the // kernel's end-of-stream sentinel. func (r *kernelRows) nextBatch() error { - // Honor cancellation at batch boundaries: check ctx before entering the - // blocking C fetch (which cannot itself observe ctx). This does NOT interrupt - // a fetch already in flight — and database/sql's own cancel watcher can't - // either: its Rows.Close takes rs.closemu.Lock(), which blocks until the - // in-progress Next (holding the RLock) returns, so the stream close waits for - // the C call to finish on its own. A single hung CloudFetch batch is therefore - // uninterruptible (the kernel exposes no per-download timeout); mid-fetch - // cancellation would need the execute path's watcher/canceller applied here. + // Stop pulling once Close cancels r.ctx; r.ctx is detached from the caller's + // deadline, so this fires on abort, not on a submit-gating timeout. It does NOT + // interrupt an in-flight C fetch (the kernel exposes no per-download timeout). if r.ctx != nil { if err := r.ctx.Err(); err != nil { return err @@ -284,3 +298,92 @@ func (r *kernelRows) nextBatch() error { klogCtx(r.ctx, "nextBatch: %d rows (chunk %d)", rec.NumRows(), r.chunkCount) return nil } + +// GetArrowBatches exposes the kernel stream as an arrow.Record iterator (the +// public zero-copy batch API), reusing the row scanner's next_batch pull. Callers +// MUST Release() each record. +func (r *kernelRows) GetArrowBatches(context.Context) (dbrows.ArrowBatchIterator, error) { + return &kernelBatchIterator{r: r}, nil +} + +// GetArrowIPCStreams is rejected on the kernel path: the C ABI exports Arrow via +// the C Data Interface (next_batch), not IPC bytes. Use GetArrowBatches, or the +// Thrift backend for IPC streams. +func (r *kernelRows) GetArrowIPCStreams(context.Context) (dbrows.ArrowIPCStreamIterator, error) { + return nil, fmt.Errorf("databricks: GetArrowIPCStreams is %w (kernel exports Arrow C Data, not IPC bytes); use GetArrowBatches", + dbsqlerr.ErrNotSupportedByKernel) +} + +// kernelBatchIterator adapts the kernelRows pull loop to the public +// ArrowBatchIterator: it prefetches one batch (so HasNext is exact) and transfers +// each record's ownership to the caller, who must Release it. +type kernelBatchIterator struct { + r *kernelRows + pending arrow.Record // prefetched, not yet handed out + done bool // stream drained (io.EOF seen) + err error // sticky fetch error, surfaced by Next +} + +var _ dbrows.ArrowBatchIterator = (*kernelBatchIterator)(nil) + +// fill buffers one batch into pending, or sets done/err. It re-uses nextBatch and +// takes ownership of r.cur so the row scanner and Close can't double-release it. +func (it *kernelBatchIterator) fill() { + if it.pending != nil || it.done || it.err != nil { + return + } + if it.r.closed || it.r.eof { // nothing left to pull + it.done = true + return + } + switch err := it.r.nextBatch(); err { + case nil: + it.pending, it.r.cur = it.r.cur, nil + case io.EOF: + it.done = true + default: + it.err = err + // The batch path bypasses Next(), so seed iterationErr here too; else OnClose + // records the failed stream as a successful statement in telemetry. + if it.r.iterationErr == nil { + it.r.iterationErr = err + } + } +} + +// Next returns the next record (caller owns it) or io.EOF when the stream drains. +func (it *kernelBatchIterator) Next() (arrow.Record, error) { + it.fill() + if it.err != nil { + return nil, it.err + } + if it.pending == nil { + return nil, io.EOF + } + rec := it.pending + it.pending = nil + return rec, nil +} + +// HasNext reports whether a following Next would yield a record or an error. +func (it *kernelBatchIterator) HasNext() bool { + it.fill() + return it.pending != nil || it.err != nil +} + +// Close releases any buffered record and tears down the underlying stream. +func (it *kernelBatchIterator) Close() { + if it.pending != nil { + it.pending.Release() + it.pending = nil + } + it.r.Close() //nolint:errcheck +} + +// Schema returns the result-set schema captured at construction. +func (it *kernelBatchIterator) Schema() (*arrow.Schema, error) { + if it.r.schema == nil { + return nil, fmt.Errorf("kernel: no schema available") + } + return it.r.schema, nil +} diff --git a/internal/rows/coltype_thrift_parity_test.go b/internal/rows/coltype_thrift_parity_test.go index be178a64..7635ec00 100644 --- a/internal/rows/coltype_thrift_parity_test.go +++ b/internal/rows/coltype_thrift_parity_test.go @@ -10,24 +10,12 @@ import ( "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). +// TestColumnTypeInfoMatchesThriftMapping cross-checks arrowscan.ColumnTypeInfoFor +// against the Thrift mapping (getScanType / GetDBTypeName / ColumnTypeLength) +// directly, so a change to one that isn't mirrored in the other fails a pure-Go PR +// run — not just the nightly warehouse-gated parity test. It enumerates each shared +// type as an {Arrow type, equivalent Thrift TColumnDesc} pair and asserts they agree. +// White-box (package rows) to reach the unexported getScanType. func TestColumnTypeInfoMatchesThriftMapping(t *testing.T) { // thriftCol builds the minimal TColumnDesc the Thrift mapping functions read. thriftCol := func(id cli_service.TTypeId) *cli_service.TColumnDesc { @@ -63,15 +51,14 @@ func TestColumnTypeInfoMatchesThriftMapping(t *testing.T) { {"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). + // Intervals: the server pre-formats them to text and declares STRING_TYPE over + // Thrift, so pair with STRING_TYPE (the kernel formats the native arrow value to + // the same string). 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}, + // VOID/NULL: the server stringifies it over Thrift as STRING_TYPE (verified + // live — even bare SELECT NULL), same as the interval types above. + {"null", arrow.Null, cli_service.TTypeId_STRING_TYPE}, } for _, c := range cases { diff --git a/kernel_auth_real_test.go b/kernel_auth_real_test.go index d89ab30a..1e677725 100644 --- a/kernel_auth_real_test.go +++ b/kernel_auth_real_test.go @@ -1,8 +1,10 @@ package dbsql import ( + "reflect" "testing" + "github.com/databricks/databricks-sql-go/auth/oauth" "github.com/databricks/databricks-sql-go/auth/oauth/m2m" "github.com/databricks/databricks-sql-go/auth/oauth/u2m" "github.com/databricks/databricks-sql-go/internal/backend/kernel" @@ -64,5 +66,9 @@ func TestResolveKernelAuthRealAuthenticators(t *testing.T) { if got.Mode != kernel.AuthU2M || got.ClientID == "" { t.Errorf("auth = %+v, want mode=U2M with a non-empty (cloud-inferred) clientID", got) } + // Scopes are forwarded from oauth.GetScopes(cfg.Host) for Thrift parity. + if want := oauth.GetScopes(c.Host, nil); !reflect.DeepEqual(got.Scopes, want) { + t.Errorf("U2M scopes = %v, want %v (Thrift parity)", got.Scopes, want) + } }) } diff --git a/kernel_config.go b/kernel_config.go index e4bf996f..0a8053a6 100644 --- a/kernel_config.go +++ b/kernel_config.go @@ -8,9 +8,11 @@ import ( "time" "github.com/databricks/databricks-sql-go/auth/noop" + "github.com/databricks/databricks-sql-go/auth/oauth" "github.com/databricks/databricks-sql-go/auth/pat" dbsqlerr "github.com/databricks/databricks-sql-go/errors" "github.com/databricks/databricks-sql-go/internal/backend/kernel" + "github.com/databricks/databricks-sql-go/internal/client" "github.com/databricks/databricks-sql-go/internal/config" ) @@ -116,6 +118,8 @@ func buildKernelConfig(cfg *config.Config, kauth kernel.Auth) kernel.Config { WarehouseID: cfg.WarehouseID, Auth: kauth, Location: cfg.Location, + // Same UA the Thrift path sends, so query history attributes both alike. + UserAgent: client.BuildUserAgent(cfg), // Initial namespace: no kernel config setter, so the kernel backend applies // these post-connect via USE CATALOG / USE SCHEMA. Catalog: cfg.Catalog, @@ -304,13 +308,15 @@ func resolveKernelAuth(cfg *config.Config) (kernel.Auth, error) { clientID, clientSecret := a.M2MCredentials() return kernel.Auth{Mode: kernel.AuthM2M, ClientID: clientID, ClientSecret: clientSecret}, nil case kernel.U2MCredentialsProvider: - // Go sources only the client id for U2M; kernel.Auth.Scopes / RedirectPort - // are left zero so setAuth passes the kernel's defaults. Go exposes no - // user-facing option for U2M scopes or redirect port on either backend today - // (the native Thrift path hardcodes both), so there is nothing to forward — - // but the kernel.Auth fields + setAuth wiring model the full set_auth_u2m - // surface for if/when such an option is added (see kernel.Auth docs). - return kernel.Auth{Mode: kernel.AuthU2M, ClientID: a.U2MClientID()}, nil + // Forward the SAME cloud-specific scopes the Thrift path requests via + // oauth.GetScopes (offline_access + sql on AWS/GCP, offline_access + + // /user_impersonation on Azure), so both backends authorize against + // the built-in databricks-sql-connector client identically. Without this the + // kernel applied its own default set (all-apis + offline_access), which a + // workspace whose public client isn't granted all-apis rejects with + // access_denied. RedirectPort is still left zero (no user option; kernel + // default 8020). Passing nil to GetScopes yields the pure cloud-default set. + return kernel.Auth{Mode: kernel.AuthU2M, ClientID: a.U2MClientID(), Scopes: oauth.GetScopes(cfg.Host, nil)}, nil case nil, *noop.NoopAuth, *pat.PATAuth: // PAT (or no explicit authenticator). WithAccessToken sets both // cfg.AccessToken and a *pat.PATAuth, but WithAuthenticator(&pat.PATAuth{...}) diff --git a/kernel_config_test.go b/kernel_config_test.go index c152eeab..bf0c97c1 100644 --- a/kernel_config_test.go +++ b/kernel_config_test.go @@ -8,9 +8,11 @@ import ( "testing" "time" + "github.com/databricks/databricks-sql-go/auth/oauth" "github.com/databricks/databricks-sql-go/auth/pat" dbsqlerr "github.com/databricks/databricks-sql-go/errors" "github.com/databricks/databricks-sql-go/internal/backend/kernel" + "github.com/databricks/databricks-sql-go/internal/client" "github.com/databricks/databricks-sql-go/internal/config" ) @@ -181,6 +183,12 @@ func TestValidateKernelConfig(t *testing.T) { if a.Mode != kernel.AuthU2M || a.ClientID != "databricks-sql-connector" { t.Errorf("auth = %+v, want mode=U2M clientID=databricks-sql-connector", a) } + // Scopes must match what the Thrift path requests for this host, so both + // backends authorize against the same client identically (not the kernel's + // all-apis default). baseKernelConfig's host is AWS → [offline_access, sql]. + if want := oauth.GetScopes(c.Host, nil); !reflect.DeepEqual(a.Scopes, want) { + t.Errorf("U2M scopes = %v, want %v (Thrift parity)", a.Scopes, want) + } }) t.Run("last-applied auth wins: M2M then PAT resolves to PAT", func(t *testing.T) { @@ -329,9 +337,11 @@ var kernelConfigFieldDisposition = map[string]string{ "MaxRows": "inert", "UseLz4Compression": "inert", // kernel negotiates compression internally + // Rides in the forwarded User-Agent header (set_custom_header). + "UserAgentEntry": "forwarded", + // Not applicable to the kernel path (Thrift/HTTP-transport or telemetry knobs // that don't reach the kernel binding). - "UserAgentEntry": "inert", // TODO: forward once the kernel exposes a UA setter "EnableTelemetry": "inert", "TelemetryBatchSize": "inert", "TelemetryFlushInterval": "inert", @@ -441,6 +451,11 @@ func TestBuildKernelConfig(t *testing.T) { if kc.Auth.Mode != kauth.Mode || kc.Auth.Token != kauth.Token { t.Errorf("auth not forwarded: got %+v, want %+v", kc.Auth, kauth) } + // UserAgent must be the driver's composed UA, non-empty — else query + // history mis-attributes SEA-path queries to the kernel's built-in UA. + if want := client.BuildUserAgent(c); kc.UserAgent == "" || kc.UserAgent != want { + t.Errorf("UserAgent not forwarded: got %q, want %q", kc.UserAgent, want) + } }) t.Run("MaxChunksInMemory injected into kernel SessionConf", func(t *testing.T) { diff --git a/kernel_e2e_test.go b/kernel_e2e_test.go index 1abaea10..ee2e53a7 100644 --- a/kernel_e2e_test.go +++ b/kernel_e2e_test.go @@ -8,6 +8,7 @@ import ( "database/sql" "database/sql/driver" "errors" + "io" "os" "sync" "sync/atomic" @@ -16,6 +17,7 @@ import ( "github.com/databricks/databricks-sql-go/driverctx" dbsqlerr "github.com/databricks/databricks-sql-go/errors" + dbsqlrows "github.com/databricks/databricks-sql-go/rows" ) // kernelTestDB opens a kernel-backed *sql.DB from the DATABRICKS_PECOTESTING_* @@ -549,3 +551,50 @@ func TestKernelE2EM2M(t *testing.T) { t.Errorf("SELECT 1 = %d, want 1", got) } } + +// TestKernelE2EGetArrowBatches drains the public Arrow batch API over the kernel +// backend end to end — the path driver_e2e_test only exercises through Thrift — +// asserting an exact row count and that each record is released by the caller. +func TestKernelE2EGetArrowBatches(t *testing.T) { + db := kernelTestDB(t) + defer db.Close() + + conn, err := db.Conn(context.Background()) + if err != nil { + t.Fatalf("conn: %v", err) + } + defer conn.Close() + + const want = 100_000 + var driverRows driver.Rows + if err := conn.Raw(func(d any) error { + var qErr error + driverRows, qErr = d.(driver.QueryerContext).QueryContext(context.Background(), "SELECT id FROM range(0, 100000)", nil) + return qErr + }); err != nil { + t.Fatalf("query: %v", err) + } + defer driverRows.Close() + + batches, err := driverRows.(dbsqlrows.Rows).GetArrowBatches(context.Background()) + if err != nil { + t.Fatalf("GetArrowBatches: %v", err) + } + defer batches.Close() + + var rowCount int64 + for batches.HasNext() { + rec, err := batches.Next() + if err == io.EOF { + break + } + if err != nil { + t.Fatalf("batch Next at row %d: %v", rowCount, err) + } + rowCount += rec.NumRows() + rec.Release() // caller owns each record + } + if rowCount != want { + t.Errorf("GetArrowBatches row count = %d, want %d", rowCount, want) + } +} diff --git a/kernel_parity_test.go b/kernel_parity_test.go index 4987e152..ed9e52ed 100644 --- a/kernel_parity_test.go +++ b/kernel_parity_test.go @@ -116,7 +116,9 @@ func TestKernelThriftColumnTypeParity(t *testing.T) { "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" + // VOID: the server declares it STRING_TYPE over Thrift; asserts the kernel + // arrow.NULL→STRING mapping matches live, not just in the pure-Go tests. + "INTERVAL '1' MONTH a_iv_ym, st_point(1,2) a_geom, CAST(NULL AS VOID) a_void" kernelDB := kernelTestDB(t) defer kernelDB.Close()