diff --git a/CHANGELOG.md b/CHANGELOG.md index 5c355a1a..017b3f3d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,7 @@ # Release History ## Unreleased +- Add default-enabled kernel-backed GEOMETRY / GEOGRAPHY results, selected with `WithEnableGeoSpatialSupport` or the `enableGeoSpatialSupport` DSN parameter. Enabled mode returns EWKB with its per-row SRID and logical geo column metadata from `database/sql`; disabled mode returns WKT / EWKT reported as `STRING`. Arrow batch access retains the canonical `struct` value with separate OGC WKB in enabled mode. - **Behavior change:** `m2m.NewAuthenticatorWithScopes` now requests exactly the given scopes instead of always appending `all-apis`, so service principals with scoped OAuth secrets (e.g. `sql` only) can authenticate (databricks/databricks-sql-go#476). Empty scopes and `m2m.NewAuthenticator` still request `all-apis`; callers that relied on it being appended should pass no scopes, or add `all-apis` to their list if they need both. The SEA/kernel backend now forwards custom M2M scopes instead of rejecting them. - Let the kernel telemetry feature flag decide enablement when `enableTelemetry` is unset; explicit `true` and `false` continue to override it. diff --git a/CONNECTION_PARAMETERS.md b/CONNECTION_PARAMETERS.md index 1cf477ee..89aae8f4 100644 --- a/CONNECTION_PARAMETERS.md +++ b/CONNECTION_PARAMETERS.md @@ -155,12 +155,14 @@ ignored with a warning and still stripped from the SEA wire. |---|---|:---:|:---:|---|---| | `useArrowNativeDecimal` | `WithArrowNativeDecimal` | ✅ | ⚠️ | `false` | Thrift: return DECIMAL as native Arrow `decimal128` (lossless string when scanned via `database/sql`). The kernel path already renders DECIMAL as the exact string regardless, so the flag is inert there. | | | `WithKernelDecimalAsFloat(b)` | ❌ | ✅ | `false` | Scan top-level DECIMAL as lossy `float64` instead of the exact string. | +| `enableGeoSpatialSupport` | `WithEnableGeoSpatialSupport(b)` | ❌ | ✅ | `true` | `true`: return top-level GEOMETRY / GEOGRAPHY as EWKB `[]byte` with the per-row SRID embedded and logical geo column metadata; `GetArrowBatches` retains the full `{srid,wkb}` struct with OGC WKB. `false`: return WKT / EWKT strings reported as `STRING`. | -With the default client-result settings, results otherwise render **byte-for-byte +Apart from the geospatial representation selected above, results render **byte-for-byte identically** on both backends (scalars, DECIMAL as exact string, TIMESTAMP / TIMESTAMP_NTZ shifted into the session time zone, -INTERVAL, nested ARRAY / MAP / STRUCT and VARIANT as JSON, GEOMETRY / GEOGRAPHY as WKT, -BINARY as `sql.RawBytes`). +INTERVAL, nested ARRAY / MAP / STRUCT and VARIANT as JSON, and BINARY as +`sql.RawBytes`). Set `enableGeoSpatialSupport=false` on the kernel path to match +the Thrift path's WKT / EWKT geospatial values. ## Cloud Fetch diff --git a/KERNEL_REV b/KERNEL_REV index bda2dea0..0c44de9b 100644 --- a/KERNEL_REV +++ b/KERNEL_REV @@ -1 +1 @@ -e5008635d35af1f83a88057c4cf4d97718e31ab4 +9c2e2e6182fc7285bcfff39b2177941bc804897c diff --git a/README.md b/README.md index f8c473dd..4bae2a09 100644 --- a/README.md +++ b/README.md @@ -270,6 +270,7 @@ session parameter on both backends. |---|---|---|---|---| | `useArrowNativeDecimal` | `WithArrowNativeDecimal` | Thrift only (inert on kernel) | `false` | Thrift: return DECIMAL as native Arrow `decimal128` (lossless string when scanned via `database/sql`). The kernel path already renders DECIMAL as the exact string regardless. | | | `WithKernelDecimalAsFloat(b)` | SEA only | `false` | Scan top-level DECIMAL as lossy `float64` instead of the exact string. | +| `enableGeoSpatialSupport` | `WithEnableGeoSpatialSupport(b)` | SEA only | `true` | `true`: return top-level GEOMETRY / GEOGRAPHY as EWKB `[]byte` with the per-row SRID embedded and logical geo column metadata; `GetArrowBatches` retains the full `{srid,wkb}` struct with OGC WKB. `false`: return WKT / EWKT strings reported as `STRING`. | See [Cloud Fetch](#cloud-fetch), [TLS](#tls), and [Proxy](#proxy) for the remaining groups. Telemetry parameters are covered under [Telemetry](#telemetry). @@ -358,10 +359,12 @@ The Thrift and kernel backends both honor the standard `HTTP_PROXY` / `HTTPS_PRO ## Data types -Results render **byte-for-byte identically** on both backends. Scalars, DECIMAL (exact -string), TIMESTAMP / TIMESTAMP_NTZ (shifted into the session time zone), INTERVAL, nested -ARRAY / MAP / STRUCT and VARIANT (as JSON), and GEOMETRY / GEOGRAPHY (as WKT) are all -supported. BINARY is returned as `sql.RawBytes`. +Apart from the configurable geospatial representation, results render **byte-for-byte +identically** on both backends. Scalars, DECIMAL (exact string), TIMESTAMP / +TIMESTAMP_NTZ (shifted into the session time zone), INTERVAL, and nested ARRAY / MAP / +STRUCT and VARIANT (as JSON) are supported. BINARY is returned as `sql.RawBytes`. +Thrift returns GEOMETRY / GEOGRAPHY as WKT / EWKT; the kernel defaults to EWKB and +matches Thrift when `enableGeoSpatialSupport=false`. Metadata is reached through SQL (`SHOW`, `DESCRIBE`, `information_schema`) on both backends — the driver exposes no `GetCatalogs`/`GetSchemas`/`GetTables`/`GetColumns` API diff --git a/connector.go b/connector.go index e7ac38d6..296d80ec 100644 --- a/connector.go +++ b/connector.go @@ -316,6 +316,10 @@ func withUserConfig(ucfg config.UserConfig) ConnOption { if ucfg.TokenCacheEnabledDSN { kernelExperimental(c).TokenCacheEnabled = true } + if ucfg.EnableGeoSpatialSupportDSN != nil { + enabled := *ucfg.EnableGeoSpatialSupportDSN + kernelExperimental(c).EnableGeoSpatialSupport = &enabled + } } } @@ -671,6 +675,28 @@ func WithKernelDecimalAsFloat(asFloat bool) ConnOption { } } +// WithEnableGeoSpatialSupport selects how the kernel path exposes top-level +// GEOMETRY and GEOGRAPHY result values. True (the default) returns EWKB bytes +// ([]byte), including the per-row SRID, and reports the logical GEOMETRY / +// GEOGRAPHY type through database/sql metadata. False returns WKT / EWKT strings and reports STRING, matching the +// existing string-mode driver contract. The per-row SRID remains present in the +// Arrow struct returned by the driver's GetArrowBatches API in enabled mode; +// that raw Arrow value retains Databricks' separate SRID and OGC WKB children. +// +// The conversion is entirely client-side and is never forwarded to SEA. +// Nested geospatial values follow the driver's existing complex-value contract: +// disabled mode renders WKT / EWKT leaves in JSON, while enabled mode preserves the +// {"srid":...,"wkb":...} shape (with JSON base64 for WKB). +// +// EXPERIMENTAL, kernel-only: the default (Thrift) backend rejects this option at +// connect time. +func WithEnableGeoSpatialSupport(enabled bool) ConnOption { + return func(c *config.Config) { + value := enabled + kernelExperimental(c).EnableGeoSpatialSupport = &value + } +} + // WithKernelTrustedCerts adds a PEM CA-certificate bundle to the kernel's TLS // trust store on top of the system roots — for a corporate re-signing proxy or an // on-prem CA. Required (rather than relying on SSL_CERT_FILE) because the kernel's diff --git a/connector_test.go b/connector_test.go index c2b50079..3b1382dd 100644 --- a/connector_test.go +++ b/connector_test.go @@ -450,6 +450,19 @@ func TestNewConnector(t *testing.T) { assert.Nil(t, coni.cfg.KernelExperimental) }) + t.Run("Connector test enableGeoSpatialSupport=false DSN param preserves explicit false", func(t *testing.T) { + ucfg, err := config.ParseDSN("token:supersecret@databricks-host:443/sql/1.0/endpoints/abc?enableGeoSpatialSupport=false") + require.NoError(t, err) + con, err := NewConnector(withUserConfig(ucfg)) + require.NoError(t, err) + + coni, ok := con.(*connector) + require.True(t, ok) + require.NotNil(t, coni.cfg.KernelExperimental) + require.NotNil(t, coni.cfg.KernelExperimental.EnableGeoSpatialSupport) + assert.False(t, *coni.cfg.KernelExperimental.EnableGeoSpatialSupport) + }) + t.Run("Connector test WithTransport sets HTTPClient in CloudFetchConfig", func(t *testing.T) { host := "databricks-host" accessToken := "token" diff --git a/internal/arrowscan/arrowscan.go b/internal/arrowscan/arrowscan.go index 8703addc..c6705955 100644 --- a/internal/arrowscan/arrowscan.go +++ b/internal/arrowscan/arrowscan.go @@ -36,8 +36,10 @@ import ( // date, timestamp, and top-level decimal (as an exact fixed-point string, // matching the Thrift path — a float64 would lose precision beyond ~17 digits). // Nested types (List/Map/Struct, and VARIANT which -// arrives nested) render to a JSON string byte-identical to the Thrift path; -// GEOMETRY arrives as a WKB/WKT string and is handled by the string arm. INTERVAL +// arrives nested) render to a JSON string byte-identical to the Thrift path. +// Top-level binary GEOMETRY / GEOGRAPHY is selected explicitly through +// ScanOptions and combines its SRID and WKB children into EWKB []byte; string +// mode is handled by the regular string arm. INTERVAL // day-time/year-month arrive as native arrow duration/month-interval and format to // the same string the Thrift path receives pre-formatted from the server. NULLs // map to nil. A genuinely unhandled type returns an error rather than a silently @@ -106,20 +108,39 @@ func (c *StructKeyCache) keyPrefixes(st *arrow.StructType) []string { // StructKeyCache) so struct field-name keys are escaped once per result set // rather than once per row. Pass nil for the un-memoized one-shot behavior. func ScanCellCached(col arrow.Array, row int, loc *time.Location, keys *StructKeyCache) (driver.Value, error) { - return scanCell(col, row, loc, keys, false) + return scanCell(col, row, loc, keys, ScanOptions{}) } // ScanCellCachedDecimalFloat is ScanCellCached that, when decimalAsFloat is true, // scans a TOP-LEVEL Decimal128 to a lossy float64 instead of the exact string // (nested decimals still render exactly). Opt in via WithKernelDecimalAsFloat. func ScanCellCachedDecimalFloat(col arrow.Array, row int, loc *time.Location, keys *StructKeyCache, decimalAsFloat bool) (driver.Value, error) { - return scanCell(col, row, loc, keys, decimalAsFloat) + return scanCell(col, row, loc, keys, ScanOptions{DecimalAsFloat: decimalAsFloat}) } -func scanCell(col arrow.Array, row int, loc *time.Location, keys *StructKeyCache, decimalAsFloat bool) (driver.Value, error) { +// ScanOptions controls top-level result adaptations that cannot be inferred from +// an Arrow array alone. Logical GEOMETRY / GEOGRAPHY identity lives on the +// enclosing schema Field, so kernelRows sets GeospatialBinary only for a field +// carrying that metadata and a physical struct value. +type ScanOptions struct { + DecimalAsFloat bool + GeospatialBinary bool +} + +// ScanCellCachedWithOptions is the full scanner entry point used by the kernel +// rows implementation. Callers that do not need result-specific adaptations +// should continue to use ScanCellCached. +func ScanCellCachedWithOptions(col arrow.Array, row int, loc *time.Location, keys *StructKeyCache, options ScanOptions) (driver.Value, error) { + return scanCell(col, row, loc, keys, options) +} + +func scanCell(col arrow.Array, row int, loc *time.Location, keys *StructKeyCache, options ScanOptions) (driver.Value, error) { if col.IsNull(row) { return nil, nil } + if options.GeospatialBinary { + return scanGeospatialEWKB(col, row) + } switch c := col.(type) { case *array.Null: return nil, nil @@ -195,7 +216,7 @@ func scanCell(col arrow.Array, row int, loc *time.Location, keys *StructKeyCache return inLocation(timestampToTime(int64(c.Value(row)), dt.Unit), loc), nil case *array.Decimal128: dt := col.DataType().(*arrow.Decimal128Type) - if decimalAsFloat { + if options.DecimalAsFloat { // Lossy fast path: float64, no per-cell string. Opt-in only. return c.Value(row).ToFloat64(dt.Scale), nil } @@ -228,6 +249,41 @@ func scanCell(col arrow.Array, row int, loc *time.Location, keys *StructKeyCache } } +// scanGeospatialEWKB validates Reyden's canonical top-level +// struct value and embeds the sibling SRID into EWKB. +// database/sql cannot return the two-child Arrow struct as a driver.Value, and +// returning bare OGC WKB would lose the per-row SRID for GEOMETRY(ANY). +func scanGeospatialEWKB(col arrow.Array, row int) (driver.Value, error) { + values, ok := col.(*array.Struct) + if !ok { + return nil, fmt.Errorf("binary geospatial column has arrow type %s, want struct", col.DataType()) + } + fields := values.DataType().(*arrow.StructType).Fields() + if len(fields) != 2 || fields[0].Name != "srid" || fields[1].Name != "wkb" || + fields[0].Type.ID() != arrow.INT32 || fields[1].Type.ID() != arrow.BINARY { + return nil, fmt.Errorf("binary geospatial column has arrow type %s, want struct", col.DataType()) + } + srids, ok := values.Field(0).(*array.Int32) + if !ok { + return nil, fmt.Errorf("binary geospatial srid child has arrow type %s, want int32", values.Field(0).DataType()) + } + wkbs, ok := values.Field(1).(*array.Binary) + if !ok { + return nil, fmt.Errorf("binary geospatial wkb child has arrow type %s, want binary", values.Field(1).DataType()) + } + if srids.IsNull(row) { + return nil, fmt.Errorf("binary geospatial value at row %d has a null srid child", row) + } + if wkbs.IsNull(row) { + return nil, fmt.Errorf("binary geospatial value at row %d has a null wkb child", row) + } + ewkb, err := wkbToEWKB(wkbs.Value(row), srids.Value(row)) + if err != nil { + return nil, fmt.Errorf("binary geospatial value at row %d: %w", row, err) + } + return ewkb, nil +} + // formatDayTimeInterval renders an arrow duration (in the given time unit) as the // Thrift path's "D HH:MM:SS.nnnnnnnnn" — days, then zero-padded hours:minutes:seconds // with 9 fractional digits, negated with a leading '-'. diff --git a/internal/arrowscan/arrowscan_test.go b/internal/arrowscan/arrowscan_test.go index 31a8774b..e68745eb 100644 --- a/internal/arrowscan/arrowscan_test.go +++ b/internal/arrowscan/arrowscan_test.go @@ -1,7 +1,9 @@ package arrowscan import ( + "bytes" "fmt" + "strings" "testing" "time" @@ -684,3 +686,100 @@ func TestScanCellDecimalAsFloat(t *testing.T) { t.Errorf("decimalAsFloat arm = %v, want ~99.98", fv) } } + +func TestScanCellBinaryGeospatial(t *testing.T) { + pool := memory.NewGoAllocator() + wkb := []byte{ + 0x01, 0x01, 0x00, 0x00, 0x00, + 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xf0, 0x3f, + 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x40, + } + geoType := arrow.StructOf( + arrow.Field{Name: "srid", Type: arrow.PrimitiveTypes.Int32, Nullable: false}, + arrow.Field{Name: "wkb", Type: arrow.BinaryTypes.Binary, Nullable: false}, + ) + + build := func(outerValid, sridValid, wkbValid bool, srid int32) *array.Struct { + b := array.NewStructBuilder(pool, geoType) + defer b.Release() + b.Append(outerValid) + if sridValid { + b.FieldBuilder(0).(*array.Int32Builder).Append(srid) + } else { + b.FieldBuilder(0).(*array.Int32Builder).AppendNull() + } + if wkbValid { + b.FieldBuilder(1).(*array.BinaryBuilder).Append(wkb) + } else { + b.FieldBuilder(1).(*array.BinaryBuilder).AppendNull() + } + return b.NewStructArray() + } + + t.Run("embeds each row SRID in EWKB", func(t *testing.T) { + for _, tc := range []struct { + srid int32 + sridLittle []byte + }{ + {srid: 4326, sridLittle: []byte{0xe6, 0x10, 0x00, 0x00}}, + {srid: 3857, sridLittle: []byte{0x11, 0x0f, 0x00, 0x00}}, + } { + arr := build(true, true, true, tc.srid) + got, err := ScanCellCachedWithOptions(arr, 0, nil, nil, ScanOptions{GeospatialBinary: true}) + arr.Release() + if err != nil { + t.Fatal(err) + } + want := append([]byte{0x01, 0x01, 0x00, 0x00, 0x20}, tc.sridLittle...) + want = append(want, wkb[5:]...) + if !bytes.Equal(got.([]byte), want) { + t.Errorf("SRID %d EWKB = %x, want %x", tc.srid, got, want) + } + } + }) + + t.Run("outer null wins", func(t *testing.T) { + arr := build(false, false, false, 0) + defer arr.Release() + got, err := ScanCellCachedWithOptions(arr, 0, nil, nil, ScanOptions{GeospatialBinary: true}) + if err != nil || got != nil { + t.Errorf("outer-null geo = (%v, %v), want (nil, nil)", got, err) + } + }) + + t.Run("rejects null srid child", func(t *testing.T) { + arr := build(true, false, true, 0) + defer arr.Release() + _, err := ScanCellCachedWithOptions(arr, 0, nil, nil, ScanOptions{GeospatialBinary: true}) + if err == nil || !strings.Contains(err.Error(), "null srid") { + t.Errorf("error = %v, want null-srid validation error", err) + } + }) + + t.Run("rejects null wkb child", func(t *testing.T) { + arr := build(true, true, false, 4326) + defer arr.Release() + _, err := ScanCellCachedWithOptions(arr, 0, nil, nil, ScanOptions{GeospatialBinary: true}) + if err == nil || !strings.Contains(err.Error(), "null wkb") { + t.Errorf("error = %v, want null-wkb validation error", err) + } + }) + + t.Run("rejects a non-canonical struct", func(t *testing.T) { + badType := arrow.StructOf( + arrow.Field{Name: "srid", Type: arrow.PrimitiveTypes.Int64}, + arrow.Field{Name: "bytes", Type: arrow.BinaryTypes.Binary}, + ) + b := array.NewStructBuilder(pool, badType) + defer b.Release() + b.Append(true) + b.FieldBuilder(0).(*array.Int64Builder).Append(4326) + b.FieldBuilder(1).(*array.BinaryBuilder).Append(wkb) + arr := b.NewStructArray() + defer arr.Release() + _, err := ScanCellCachedWithOptions(arr, 0, nil, nil, ScanOptions{GeospatialBinary: true}) + if err == nil || !strings.Contains(err.Error(), "want struct") { + t.Errorf("error = %v, want canonical-shape validation error", err) + } + }) +} diff --git a/internal/arrowscan/coltype.go b/internal/arrowscan/coltype.go index b4be2a1c..d1e8a750 100644 --- a/internal/arrowscan/coltype.go +++ b/internal/arrowscan/coltype.go @@ -4,11 +4,14 @@ import ( "database/sql" "math" "reflect" + "strings" "time" "github.com/apache/arrow/go/v12/arrow" ) +const databricksTypeNameMetadataKey = "databricks.type_name" + // ColumnTypeInfo is the per-column metadata database/sql surfaces through // sql.ColumnType. The kernel derives it from the result's Arrow schema via // ColumnTypeInfoFor — the mapping the value scanner and type reporter share, kept @@ -110,6 +113,53 @@ func ColumnTypeInfoFor(dt arrow.DataType) ColumnTypeInfo { } } +// LogicalGeospatialType returns the top-level Databricks geospatial family +// recorded by the kernel on an Arrow field. The physical value is UTF-8 in +// string mode and a struct in binary mode, so DataType alone cannot preserve +// this database type name. +func LogicalGeospatialType(field arrow.Field) (string, bool) { + name, ok := field.Metadata.GetValue(databricksTypeNameMetadataKey) + if !ok { + return "", false + } + switch strings.ToUpper(strings.TrimSpace(name)) { + case "GEOMETRY": + return "GEOMETRY", true + case "GEOGRAPHY": + return "GEOGRAPHY", true + default: + return "", false + } +} + +// IsBinaryGeospatialField reports whether field is a logical top-level +// GEOMETRY / GEOGRAPHY using the kernel's struct representation. The scanner +// performs the detailed child-shape validation when it reads a non-null value. +func IsBinaryGeospatialField(field arrow.Field) bool { + _, geospatial := LogicalGeospatialType(field) + return geospatial && field.Type.ID() == arrow.STRUCT +} + +// ColumnTypeInfoForField reports logical GEOMETRY / GEOGRAPHY metadata for the +// native struct representation. String mode follows the existing driver +// contract and reports STRING, matching its UTF-8 physical representation. +// Other fields retain the Arrow-only mapping. +func ColumnTypeInfoForField(field arrow.Field) ColumnTypeInfo { + name, geospatial := LogicalGeospatialType(field) + if !geospatial { + return ColumnTypeInfoFor(field.Type) + } + if field.Type.ID() == arrow.STRUCT { + return varLen(name, scanTypeRawBytes) + } + if field.Type.ID() == arrow.STRING || field.Type.ID() == arrow.LARGE_STRING { + return varLen("STRING", scanTypeString) + } + info := ColumnTypeInfoFor(field.Type) + info.DatabaseTypeName = name + return info +} + // 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 { diff --git a/internal/arrowscan/coltype_test.go b/internal/arrowscan/coltype_test.go index 0dbd2f8a..dfb643be 100644 --- a/internal/arrowscan/coltype_test.go +++ b/internal/arrowscan/coltype_test.go @@ -100,3 +100,43 @@ func TestColumnTypeInfoScanTypeCoversScanner(t *testing.T) { } } } + +func TestColumnTypeInfoForFieldReportsGeospatialByRepresentation(t *testing.T) { + raw := reflect.TypeOf(sql.RawBytes{}) + str := reflect.TypeOf("") + metadata := arrow.MetadataFrom(map[string]string{ + "databricks.type_name": "geometry", + "databricks.type_text": "GEOMETRY(ANY)", + }) + geoStruct := arrow.StructOf( + arrow.Field{Name: "srid", Type: arrow.PrimitiveTypes.Int32, Nullable: false}, + arrow.Field{Name: "wkb", Type: arrow.BinaryTypes.Binary, Nullable: false}, + ) + + for _, stringType := range []arrow.DataType{arrow.BinaryTypes.String, arrow.BinaryTypes.LargeString} { + stringField := arrow.Field{Name: "g", Type: stringType, Nullable: true, Metadata: metadata} + stringInfo := ColumnTypeInfoForField(stringField) + if stringInfo.DatabaseTypeName != "STRING" || stringInfo.ScanType != str { + t.Errorf("%s geo metadata = (%q, %v), want (STRING, string)", + stringType, stringInfo.DatabaseTypeName, stringInfo.ScanType) + } + if IsBinaryGeospatialField(stringField) { + t.Errorf("%s geospatial field classified as binary", stringType) + } + } + + binaryField := arrow.Field{Name: "g", Type: geoStruct, Nullable: true, Metadata: metadata} + binaryInfo := ColumnTypeInfoForField(binaryField) + if binaryInfo.DatabaseTypeName != "GEOMETRY" || binaryInfo.ScanType != raw { + t.Errorf("binary geo metadata = (%q, %v), want (GEOMETRY, RawBytes)", binaryInfo.DatabaseTypeName, binaryInfo.ScanType) + } + if !IsBinaryGeospatialField(binaryField) { + t.Error("logical geospatial struct was not classified as binary") + } + + ordinaryStruct := arrow.Field{Name: "s", Type: geoStruct, Nullable: true} + ordinaryInfo := ColumnTypeInfoForField(ordinaryStruct) + if ordinaryInfo.DatabaseTypeName != "STRUCT" || IsBinaryGeospatialField(ordinaryStruct) { + t.Errorf("ordinary struct misclassified: %+v", ordinaryInfo) + } +} diff --git a/internal/arrowscan/ewkb.go b/internal/arrowscan/ewkb.go new file mode 100644 index 00000000..651a8953 --- /dev/null +++ b/internal/arrowscan/ewkb.go @@ -0,0 +1,179 @@ +package arrowscan + +import ( + "encoding/binary" + "fmt" +) + +const ( + ewkbZFlag uint32 = 0x80000000 + ewkbMFlag uint32 = 0x40000000 + ewkbSRIDFlag uint32 = 0x20000000 + ewkbBBoxFlag uint32 = 0x10000000 +) + +// wkbToEWKB converts ISO/OGC WKB dimensional type offsets to EWKB flags at +// every geometry header and embeds the row SRID in the outer header. +func wkbToEWKB(wkb []byte, srid int32) ([]byte, error) { + rewriter := wkbRewriter{ + input: wkb, + output: make([]byte, 0, len(wkb)+4), + } + if _, err := rewriter.rewriteGeometry(true, srid); err != nil { + return nil, err + } + if rewriter.offset != len(wkb) { + return nil, fmt.Errorf("malformed WKB: %d trailing bytes", len(wkb)-rewriter.offset) + } + return rewriter.output, nil +} + +type wkbRewriter struct { + input []byte + offset int + output []byte +} + +func (r *wkbRewriter) rewriteGeometry(includeSRID bool, srid int32) (uint32, error) { + if len(r.input)-r.offset < 5 { + return 0, fmt.Errorf("malformed WKB: geometry header is truncated") + } + + marker := r.input[r.offset] + r.offset++ + var order binary.ByteOrder + switch marker { + case 0: + order = binary.BigEndian + case 1: + order = binary.LittleEndian + default: + return 0, fmt.Errorf("malformed WKB: invalid byte order %d", marker) + } + + rawType, err := r.readUint32(order) + if err != nil { + return 0, err + } + baseType, dimensions, dimensionFlags, err := decodeISOWKBType(rawType) + if err != nil { + return 0, err + } + + r.output = append(r.output, marker) + outputType := baseType | dimensionFlags + if includeSRID { + outputType |= ewkbSRIDFlag + } + r.appendUint32(order, outputType) + if includeSRID { + r.appendUint32(order, uint32(srid)) // #nosec G115 -- EWKB stores the exact int32 bit pattern + } + + switch baseType { + case 1: // POINT + err = r.copyBytes(uint64(dimensions) * 8) + case 2: // LINESTRING + err = r.copyCoordinateSequence(order, dimensions) + case 3: // POLYGON + var rings uint32 + rings, err = r.copyUint32(order) + if err == nil { + for ring := uint32(0); ring < rings; ring++ { + if err = r.copyCoordinateSequence(order, dimensions); err != nil { + break + } + } + } + case 4, 5, 6, 7: // MULTI* / GEOMETRYCOLLECTION + var children uint32 + children, err = r.copyUint32(order) + if err == nil { + for child := uint32(0); child < children; child++ { + var childType uint32 + childType, err = r.rewriteGeometry(false, 0) + if err != nil { + break + } + if baseType != 7 && childType != baseType-3 { + err = fmt.Errorf( + "malformed WKB: geometry type %d contains child type %d", + baseType, + childType, + ) + break + } + } + } + default: + return 0, fmt.Errorf("malformed WKB: unsupported geometry type %d", baseType) + } + if err != nil { + return 0, err + } + return baseType, nil +} + +func decodeISOWKBType(rawType uint32) (baseType uint32, dimensions uint32, flags uint32, err error) { + if rawType&(ewkbZFlag|ewkbMFlag|ewkbSRIDFlag|ewkbBBoxFlag) != 0 { + return 0, 0, 0, fmt.Errorf("malformed WKB: expected ISO/OGC type code, got EWKB flags in %#x", rawType) + } + + switch { + case rawType >= 3000 && rawType < 4000: + baseType, dimensions, flags = rawType-3000, 4, ewkbZFlag|ewkbMFlag + case rawType >= 2000 && rawType < 3000: + baseType, dimensions, flags = rawType-2000, 3, ewkbMFlag + case rawType >= 1000 && rawType < 2000: + baseType, dimensions, flags = rawType-1000, 3, ewkbZFlag + default: + baseType, dimensions = rawType, 2 + } + if baseType < 1 || baseType > 7 { + return 0, 0, 0, fmt.Errorf("malformed WKB: unsupported geometry type %d", rawType) + } + return baseType, dimensions, flags, nil +} + +func (r *wkbRewriter) readUint32(order binary.ByteOrder) (uint32, error) { + if len(r.input)-r.offset < 4 { + return 0, fmt.Errorf("malformed WKB: truncated uint32") + } + value := order.Uint32(r.input[r.offset : r.offset+4]) + r.offset += 4 + return value, nil +} + +func (r *wkbRewriter) appendUint32(order binary.ByteOrder, value uint32) { + var encoded [4]byte + order.PutUint32(encoded[:], value) + r.output = append(r.output, encoded[:]...) +} + +func (r *wkbRewriter) copyUint32(order binary.ByteOrder) (uint32, error) { + value, err := r.readUint32(order) + if err != nil { + return 0, err + } + r.appendUint32(order, value) + return value, nil +} + +func (r *wkbRewriter) copyCoordinateSequence(order binary.ByteOrder, dimensions uint32) error { + points, err := r.copyUint32(order) + if err != nil { + return err + } + return r.copyBytes(uint64(points) * uint64(dimensions) * 8) +} + +func (r *wkbRewriter) copyBytes(count uint64) error { + remaining := r.input[r.offset:] + if count > uint64(len(remaining)) { + return fmt.Errorf("malformed WKB: payload is truncated") + } + payload := remaining[:count] + r.output = append(r.output, payload...) + r.offset += len(payload) + return nil +} diff --git a/internal/arrowscan/ewkb_test.go b/internal/arrowscan/ewkb_test.go new file mode 100644 index 00000000..076f731e --- /dev/null +++ b/internal/arrowscan/ewkb_test.go @@ -0,0 +1,122 @@ +package arrowscan + +import ( + "bytes" + "encoding/binary" + "testing" +) + +func TestWKBToEWKBConvertsISODimensions(t *testing.T) { + for _, tc := range []struct { + name string + isoType uint32 + wantType uint32 + dims int + }{ + {name: "XY", isoType: 1, wantType: ewkbSRIDFlag | 1, dims: 2}, + {name: "Z", isoType: 1001, wantType: ewkbZFlag | ewkbSRIDFlag | 1, dims: 3}, + {name: "M", isoType: 2001, wantType: ewkbMFlag | ewkbSRIDFlag | 1, dims: 3}, + {name: "ZM", isoType: 3001, wantType: ewkbZFlag | ewkbMFlag | ewkbSRIDFlag | 1, dims: 4}, + } { + t.Run(tc.name, func(t *testing.T) { + input := pointWKB(binary.LittleEndian, tc.isoType, tc.dims) + got, err := wkbToEWKB(input, 4326) + if err != nil { + t.Fatal(err) + } + if got[0] != 1 || binary.LittleEndian.Uint32(got[1:5]) != tc.wantType { + t.Fatalf("EWKB header = %x, want type %#x", got[:5], tc.wantType) + } + if binary.LittleEndian.Uint32(got[5:9]) != 4326 { + t.Fatalf("EWKB SRID = %d, want 4326", binary.LittleEndian.Uint32(got[5:9])) + } + if !bytes.Equal(got[9:], input[5:]) { + t.Fatalf("coordinate payload changed: got %x, want %x", got[9:], input[5:]) + } + }) + } +} + +func TestWKBToEWKBHandlesBigEndian(t *testing.T) { + input := pointWKB(binary.BigEndian, 1001, 3) + got, err := wkbToEWKB(input, 3857) + if err != nil { + t.Fatal(err) + } + if got[0] != 0 || binary.BigEndian.Uint32(got[1:5]) != ewkbZFlag|ewkbSRIDFlag|1 { + t.Fatalf("EWKB header = %x", got[:5]) + } + if binary.BigEndian.Uint32(got[5:9]) != 3857 { + t.Fatalf("EWKB SRID = %d, want 3857", binary.BigEndian.Uint32(got[5:9])) + } +} + +func TestWKBToEWKBRewritesNestedCollectionHeaders(t *testing.T) { + input := []byte{1} + input = appendTestUint32(input, binary.LittleEndian, 1007) + input = appendTestUint32(input, binary.LittleEndian, 1) + input = append(input, pointWKB(binary.LittleEndian, 1001, 3)...) + + got, err := wkbToEWKB(input, 4326) + if err != nil { + t.Fatal(err) + } + if outer := binary.LittleEndian.Uint32(got[1:5]); outer != ewkbZFlag|ewkbSRIDFlag|7 { + t.Fatalf("collection type = %#x", outer) + } + if child := binary.LittleEndian.Uint32(got[14:18]); child != ewkbZFlag|1 { + t.Fatalf("nested point type = %#x", child) + } +} + +func TestWKBToEWKBRejectsMalformedPayloads(t *testing.T) { + for _, input := range [][]byte{ + {}, + {2, 1, 0, 0, 0}, + {1, 1, 0, 0, 0}, + {1, 0, 0, 0, 0}, + {1, 8, 0, 0, 0}, + {1, 2, 0, 0, 0, 0xff, 0xff, 0xff, 0xff}, + append(pointWKB(binary.LittleEndian, 1, 2), 0), + } { + if _, err := wkbToEWKB(input, 0); err == nil { + t.Fatalf("wkbToEWKB(%x) succeeded, want error", input) + } + } +} + +func TestWKBRewriterCopyBytesChecksBounds(t *testing.T) { + for _, count := range []uint64{0, 2, 3, 4, ^uint64(0)} { + r := wkbRewriter{input: []byte{0, 1, 2, 3}, offset: 1} + err := r.copyBytes(count) + if count > 3 { + if err == nil || r.offset != 1 || len(r.output) != 0 { + t.Fatalf("copyBytes(%d) = %v, offset=%d, output=%x", count, err, r.offset, r.output) + } + continue + } + if err != nil || !bytes.Equal(r.output, r.input[1:1+count]) || r.offset != 1+len(r.output) { + t.Fatalf("copyBytes(%d) = %v, offset=%d, output=%x", count, err, r.offset, r.output) + } + } +} + +func pointWKB(order binary.ByteOrder, geometryType uint32, dimensions int) []byte { + marker := byte(0) + if order == binary.LittleEndian { + marker = 1 + } + output := appendTestUint32([]byte{marker}, order, geometryType) + for ordinate := 0; ordinate < dimensions; ordinate++ { + var encoded [8]byte + order.PutUint64(encoded[:], uint64(ordinate+1)) + output = append(output, encoded[:]...) + } + return output +} + +func appendTestUint32(output []byte, order binary.ByteOrder, value uint32) []byte { + var encoded [4]byte + order.PutUint32(encoded[:], value) + return append(output, encoded[:]...) +} diff --git a/internal/backend/kernel/backend.go b/internal/backend/kernel/backend.go index bfd17341..610c9539 100644 --- a/internal/backend/kernel/backend.go +++ b/internal/backend/kernel/backend.go @@ -55,6 +55,11 @@ static inline KernelStatusCode go_kernel_set_driver_system_configuration( os_name, os_version, os_arch, client_app_name, locale_name, char_set_encoding, process_name); } + +static inline KernelStatusCode go_kernel_set_enable_geospatial_support( + KernelSessionConfig* config, bool enabled) { + return kernel_session_config_set_enable_geospatial_support(config, enabled); +} */ import "C" @@ -258,6 +263,9 @@ func (k *KernelBackend) OpenSession(ctx context.Context) error { if err := k.applyMaxConnections(cfg); err != nil { return err } + if err := k.applyEnableGeoSpatialSupport(cfg); err != nil { + return err + } // Retry / backoff policy (WithRetries). See applyRetry. if err := k.applyRetry(cfg); err != nil { @@ -430,6 +438,18 @@ func (k *KernelBackend) applyMaxConnections(cfg *C.KernelSessionConfig) error { return nil } +// applyEnableGeoSpatialSupport selects the kernel's client-side representation +// for GEOMETRY / GEOGRAPHY results. The connector supplies the default-enabled +// value explicitly, and an explicit false crosses the ABI unchanged. +func (k *KernelBackend) applyEnableGeoSpatialSupport(cfg *C.KernelSessionConfig) error { + if err := call(func() C.KernelStatusCode { + return C.go_kernel_set_enable_geospatial_support(cfg, C.bool(k.cfg.EnableGeoSpatialSupport)) + }); err != nil { + return fmt.Errorf("kernel: set_enable_geospatial_support: %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 @@ -749,6 +769,18 @@ func trySetMaxConnections(cfg Config) error { return k.applyMaxConnections(c) } +// trySetEnableGeoSpatialSupport exercises the geospatial representation setter +// without opening a network session. It is used only by tagged kernel tests. +func trySetEnableGeoSpatialSupport(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.applyEnableGeoSpatialSupport(c) +} + // trySetTelemetry exercises the telemetry C setter without opening a network // session. It is used only by the tagged kernel tests. func trySetTelemetry(cfg Config) error { diff --git a/internal/backend/kernel/config.go b/internal/backend/kernel/config.go index 04dc7a6f..1e8c254b 100644 --- a/internal/backend/kernel/config.go +++ b/internal/backend/kernel/config.go @@ -13,7 +13,7 @@ import ( // Config is the flat connection config for the kernel backend. The connector // fills it from the driver's config so the user-facing options are unchanged. -// Zero-valued fields are simply not applied. +// Zero-valued fields are simply not applied unless documented otherwise. type Config struct { Host string // workspace hostname, no scheme HTTPPath string // e.g. /sql/1.0/warehouses/abc123 (carries ?o= org routing) @@ -99,6 +99,12 @@ type Config struct { // the exact string (from WithKernelDecimalAsFloat). Kernel still sends Decimal128. DecimalAsFloat bool + // EnableGeoSpatialSupport selects the local GEOMETRY / GEOGRAPHY result + // representation. True requests the canonical struct Arrow + // representation; false requests WKT / EWKT strings. The connector always + // sets this field and defaults it to true. + EnableGeoSpatialSupport bool + // TokenCacheEnabled controls the kernel's on-disk OAuth U2M token-cache persistence // (WithTokenCache / tokenCache DSN param). When false (the default), tokens are held // in memory only; when true, the refresh token is persisted encrypted to diff --git a/internal/backend/kernel/include/databricks_kernel.h b/internal/backend/kernel/include/databricks_kernel.h index 4070486c..6015c143 100644 --- a/internal/backend/kernel/include/databricks_kernel.h +++ b/internal/backend/kernel/include/databricks_kernel.h @@ -201,10 +201,11 @@ typedef struct kernel_session_config_t KernelSessionConfig; typedef struct kernel_session_t kernel_session_t; typedef struct kernel_statement_t kernel_statement_t; typedef struct kernel_executed_statement_t kernel_executed_statement_t; -/* Async-consumption executed handle. Reserved: the async submit path is - * deferred in v0 (see kernel_statement_submit), so no consumer receives a - * live instance yet. */ +/* Async-consumption executed handle returned by kernel_statement_submit. */ typedef struct kernel_executed_async_statement_t kernel_executed_async_statement_t; +/* Execution-scoped canceller detached from an async executed handle. */ +typedef struct kernel_executed_async_statement_canceller_t + kernel_executed_async_statement_canceller_t; typedef struct kernel_result_stream_t kernel_result_stream_t; /* Detached canceller for a sync-executing statement, returned by * kernel_statement_canceller_new. The returned handle is safe to use from @@ -343,6 +344,18 @@ KernelStatusCode kernel_session_config_set_session_conf(KernelSessionConfig* con const char* key, const char* value); +/* Select whether GEOMETRY / GEOGRAPHY results use their native Arrow + * representation. `enabled = true` requests Arrow + * `struct` values; `false` requests WKT / EWKT in + * Arrow UTF-8 values. Omitting this setter uses the enabled kernel default. + * Binary mode requires the native Reyden Arrow path; legacy text-only results + * cannot be reconstructed as WKB. This client-side setting is never sent to + * SEA. A raw `EnableGeoSpatialSupport` entry set through + * kernel_session_config_set_session_conf remains server-bound and is forwarded + * unchanged. */ +KernelStatusCode kernel_session_config_set_enable_geospatial_support( + KernelSessionConfig* config, bool enabled); + /* Append one HTTP header sent on every request. Call once per header * (order preserved); `name` and `value` are both required. */ KernelStatusCode kernel_session_config_set_custom_header(KernelSessionConfig* config, @@ -754,13 +767,7 @@ KernelStatusCode kernel_statement_execute_with_timeout_ms( kernel_statement_t* stmt, uint64_t timeout_ms, kernel_executed_statement_t** out); -/* - * Submit-and-return (async). DEFERRED in v0: this always returns - * `KernelStatusCode_InvalidArgument` (with an explanatory last error) and - * writes nothing to `*out`. Use the synchronous `kernel_statement_execute` - * path. Declared so the symbol/contract is visible; wired when an ODBC / - * Go consumer needs caller-driven polling. - */ +/* Submit-and-return. On success, *out owns an async executed handle. */ KernelStatusCode kernel_statement_submit(kernel_statement_t* stmt, kernel_executed_async_statement_t** out); @@ -887,6 +894,50 @@ const char* kernel_executed_statement_query_id(const kernel_executed_statement_t KernelStatusCode kernel_executed_statement_close(kernel_executed_statement_t* exec); +/* ─── Executed statement (caller-driven async) ─────────────────── */ + +/* Perform at most one status request. Pending, Running, Succeeded, and Closed + * return Success. Failed returns SqlError; Cancelled returns Cancelled. On + * either non-success terminal outcome, *out is still set and the calling + * thread can retrieve the error information with kernel_get_last_error. */ +KernelStatusCode kernel_executed_async_statement_status( + kernel_executed_async_statement_t* exec, KernelStatementStatusKind* out); + +/* Materialise a successful terminal result as the existing owned result-stream + * handle. This is NON-BLOCKING: the caller must first poll + * kernel_executed_async_statement_status until it reports Succeeded (or Closed) + * before calling this. Invoking it before the statement reaches a terminal + * state returns KernelStatusCode_InvalidArgument rather than blocking (v0 + * collapses invalid-state into that code). This is single-consumption: + * a second successful take is rejected. Close the stream with + * kernel_result_stream_close. */ +KernelStatusCode kernel_executed_async_statement_get_result_stream( + kernel_executed_async_statement_t* exec, kernel_result_stream_t** out); + +/* Most recently observed DML count, or -1 when unavailable. */ +int64_t kernel_executed_async_statement_num_modified_rows( + const kernel_executed_async_statement_t* exec); + +/* Borrowed NUL-terminated server statement id, valid until async close. */ +const char* kernel_executed_async_statement_query_id( + const kernel_executed_async_statement_t* exec); + +/* Create an execution-scoped canceller. The returned handle owns no borrow of + * exec and may cancel concurrently with status or result work. Join any + * in-flight cancel before freeing this handle. */ +KernelStatusCode kernel_executed_async_statement_canceller_new( + const kernel_executed_async_statement_t* exec, + kernel_executed_async_statement_canceller_t** out); +KernelStatusCode kernel_executed_async_statement_canceller_cancel( + kernel_executed_async_statement_canceller_t* canceller); +KernelStatusCode kernel_executed_async_statement_canceller_free( + kernel_executed_async_statement_canceller_t* canceller); + +/* Close the server operation and consume the async handle. The handle is + * consumed even when close returns an error. */ +KernelStatusCode kernel_executed_async_statement_close( + kernel_executed_async_statement_t* exec); + /* ─── Result stream ─────────────────────────────────────────────────── * * The schema is exported once; each `next_batch` call exports one Arrow diff --git a/internal/backend/kernel/kernel_test.go b/internal/backend/kernel/kernel_test.go index 06a030c2..cd379412 100644 --- a/internal/backend/kernel/kernel_test.go +++ b/internal/backend/kernel/kernel_test.go @@ -256,6 +256,22 @@ func TestSetTokenCache(t *testing.T) { } } +func TestSetEnableGeoSpatialSupport(t *testing.T) { + for _, tc := range []struct { + name string + cfg Config + }{ + {name: "enabled", cfg: Config{EnableGeoSpatialSupport: true}}, + {name: "disabled", cfg: Config{EnableGeoSpatialSupport: false}}, + } { + t.Run(tc.name, func(t *testing.T) { + if err := trySetEnableGeoSpatialSupport(tc.cfg); err != nil { + t.Errorf("applyEnableGeoSpatialSupport(%s) = %v, want nil", tc.name, err) + } + }) + } +} + func TestSetRequestTimeout(t *testing.T) { for _, tc := range []struct { name string diff --git a/internal/backend/kernel/rows.go b/internal/backend/kernel/rows.go index 635fc10c..d55f77ea 100644 --- a/internal/backend/kernel/rows.go +++ b/internal/backend/kernel/rows.go @@ -65,6 +65,7 @@ type kernelRows struct { cols []string colTypes []arrowscan.ColumnTypeInfo // per-column type metadata (PECOBLR-3692) + geoBinary []bool // top-level logical geo structs become EWKB for database/sql 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 @@ -118,9 +119,11 @@ func newKernelRows(ctx context.Context, op *kernelOp, stream *C.kernel_result_st // (PECOBLR-3692) with no per-call work. Kept in lockstep with the value scanner // (ScanCellCached) via the shared arrowscan.ColumnTypeInfoFor mapper. r.colTypes = make([]arrowscan.ColumnTypeInfo, len(fields)) + r.geoBinary = make([]bool, len(fields)) for i, f := range fields { r.cols[i] = f.Name - r.colTypes[i] = arrowscan.ColumnTypeInfoFor(f.Type) + r.colTypes[i] = arrowscan.ColumnTypeInfoForField(f) + r.geoBinary[i] = arrowscan.IsBinaryGeospatialField(f) } // Construction succeeded — now arm the close telemetry callback so a normal // Close() (after row iteration) records CLOSE_STATEMENT. @@ -232,7 +235,10 @@ func (r *kernelRows) next(dest []driver.Value) error { } rec := r.cur for c := 0; c < len(dest); c++ { - v, err := arrowscan.ScanCellCachedDecimalFloat(rec.Column(c), r.rowInCur, r.op.location, r.keyCache, r.op.decimalAsFloat) + v, err := arrowscan.ScanCellCachedWithOptions(rec.Column(c), r.rowInCur, r.op.location, r.keyCache, arrowscan.ScanOptions{ + DecimalAsFloat: r.op.decimalAsFloat, + GeospatialBinary: r.geoBinary[c], + }) if err != nil { return fmt.Errorf("kernel: scan col %d (%s): %w", c, r.cols[c], err) } diff --git a/internal/config/config.go b/internal/config/config.go index 75635927..cf5ba575 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -137,6 +137,12 @@ type KernelExperimentalConfig struct { // false TokenCacheEnabled calls the setter with enabled=false so the kernel does // NOT persist to disk by default. TokenCacheEnabled bool + + // EnableGeoSpatialSupport selects the kernel's client-side GEOMETRY / + // GEOGRAPHY result representation. True requests the canonical + // struct Arrow value; false requests WKT / EWKT strings. A pointer + // preserves the default-enabled value versus explicitly false. + EnableGeoSpatialSupport *bool } // DeepCopy returns a deep copy of the experimental config, or nil for a nil @@ -158,6 +164,10 @@ func (k *KernelExperimentalConfig) DeepCopy() *KernelExperimentalConfig { DecimalAsFloat: k.DecimalAsFloat, TokenCacheEnabled: k.TokenCacheEnabled, } + if k.EnableGeoSpatialSupport != nil { + enabled := *k.EnableGeoSpatialSupport + cp.EnableGeoSpatialSupport = &enabled + } if k.MaxConnections != nil { max := *k.MaxConnections cp.MaxConnections = &max @@ -302,6 +312,10 @@ type UserConfig struct { // unambiguous. False by default; when true, enables on-disk token-cache persistence // for U2M OAuth. DSN: tokenCache=true. TokenCacheEnabledDSN bool + // EnableGeoSpatialSupportDSN is a DSN-only carrier for + // enableGeoSpatialSupport=true|false. Nil means omitted (the default is + // true); the connector copies a present value into KernelExperimental. + EnableGeoSpatialSupportDSN *bool } // DeepCopy returns a true deep copy of UserConfig @@ -323,7 +337,7 @@ func (ucfg UserConfig) DeepCopy() UserConfig { } - return UserConfig{ + result := UserConfig{ Protocol: ucfg.Protocol, Host: ucfg.Host, Port: ucfg.Port, @@ -352,6 +366,11 @@ func (ucfg UserConfig) DeepCopy() UserConfig { WarehouseID: ucfg.WarehouseID, TokenCacheEnabledDSN: ucfg.TokenCacheEnabledDSN, } + if ucfg.EnableGeoSpatialSupportDSN != nil { + enabled := *ucfg.EnableGeoSpatialSupportDSN + result.EnableGeoSpatialSupportDSN = &enabled + } + return result } var defaultMaxRows = 100000 @@ -516,6 +535,12 @@ func ParseDSN(dsn string) (UserConfig, error) { } ucfg.TokenCacheEnabledDSN = tokenCache } + if enableGeoSpatialSupport, ok, err := params.extractAsBool("enableGeoSpatialSupport"); ok { + if err != nil { + return UserConfig{}, err + } + ucfg.EnableGeoSpatialSupportDSN = &enableGeoSpatialSupport + } // Telemetry parameters if enableTelemetry, ok, err := params.extractAsBool("enableTelemetry"); ok { diff --git a/internal/config/config_test.go b/internal/config/config_test.go index b1e235b0..59c50134 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -693,6 +693,39 @@ func TestParseConfig(t *testing.T) { } } +func TestParseDSNEnableGeoSpatialSupport(t *testing.T) { + base := "token:supersecret@example.cloud.databricks.com:443/sql/1.0/endpoints/abc" + for _, tc := range []struct { + name string + query string + want *bool + }{ + {name: "omitted", query: "", want: nil}, + {name: "enabled", query: "?enableGeoSpatialSupport=true", want: func() *bool { v := true; return &v }()}, + {name: "disabled", query: "?enableGeoSpatialSupport=false", want: func() *bool { v := false; return &v }()}, + } { + t.Run(tc.name, func(t *testing.T) { + got, err := ParseDSN(base + tc.query) + if err != nil { + t.Fatal(err) + } + if tc.want == nil { + if got.EnableGeoSpatialSupportDSN != nil { + t.Errorf("EnableGeoSpatialSupportDSN = %v, want nil", *got.EnableGeoSpatialSupportDSN) + } + return + } + if got.EnableGeoSpatialSupportDSN == nil || *got.EnableGeoSpatialSupportDSN != *tc.want { + t.Errorf("EnableGeoSpatialSupportDSN = %v, want %v", got.EnableGeoSpatialSupportDSN, *tc.want) + } + }) + } + + if _, err := ParseDSN(base + "?enableGeoSpatialSupport=notabool"); err == nil { + t.Error("malformed enableGeoSpatialSupport should fail DSN parsing") + } +} + func TestEffectiveSessionParams(t *testing.T) { t.Run("no metric view leaves params unchanged", func(t *testing.T) { c := &Config{UserConfig: UserConfig{SessionParams: map[string]string{"QUERY_TAGS": "a:1"}}} diff --git a/kernel_config.go b/kernel_config.go index a2a6bf49..b4a1238e 100644 --- a/kernel_config.go +++ b/kernel_config.go @@ -176,6 +176,9 @@ func buildKernelConfig(cfg *config.Config, kauth kernel.Auth) kernel.Config { // this binding's system identity for its own runtime. Telemetry: kernelTelemetryConfig(cfg), DriverSystemConfiguration: kernelDriverSystemConfiguration(cfg), + // Match the public contract and the kernel default. An explicit false + // below selects WKT / EWKT instead. + EnableGeoSpatialSupport: true, } if cfg.ClientQueryTimeout != nil { timeout := *cfg.ClientQueryTimeout @@ -209,6 +212,9 @@ func buildKernelConfig(cfg *config.Config, kauth kernel.Auth) kernel.Config { } // Client-side scan choice: lossy float64 decimals instead of exact strings. kc.DecimalAsFloat = ke.DecimalAsFloat + if ke.EnableGeoSpatialSupport != nil { + kc.EnableGeoSpatialSupport = *ke.EnableGeoSpatialSupport + } // U2M token-cache control (WithTokenCache / tokenCache DSN param): forward to kernel. kc.TokenCacheEnabled = ke.TokenCacheEnabled } diff --git a/kernel_config_test.go b/kernel_config_test.go index f786e310..6b5a59fa 100644 --- a/kernel_config_test.go +++ b/kernel_config_test.go @@ -467,9 +467,10 @@ var kernelConfigFieldDisposition = map[string]string{ "EnableTelemetry": "forwarded", "TelemetryBatchSize": "forwarded", // Forwarded to the kernel's telemetry scheduler on the kernel path. - "TelemetryFlushInterval": "forwarded", - "UseArrowNativeDecimalDSN": "inert", // DSN carrier; kernel renders decimals exactly regardless - "TokenCacheEnabledDSN": "inert", // DSN carrier; forwarded to KernelExperimental.TokenCacheEnabled + "TelemetryFlushInterval": "forwarded", + "UseArrowNativeDecimalDSN": "inert", // DSN carrier; kernel renders decimals exactly regardless + "TokenCacheEnabledDSN": "inert", // DSN carrier; forwarded to KernelExperimental.TokenCacheEnabled + "EnableGeoSpatialSupportDSN": "inert", // DSN carrier; forwarded to KernelExperimental.EnableGeoSpatialSupport // Fields promoted from the embedded CloudFetchConfig. The kernel does // CloudFetch internally (below the C ABI), so none is forwarded — but each is @@ -558,6 +559,26 @@ func TestBuildKernelConfig(t *testing.T) { } }) + t.Run("geospatial support defaults true and preserves explicit values", func(t *testing.T) { + c := baseKernelConfig() + kc := buildKernelConfig(c, kernel.Auth{Mode: kernel.AuthPAT, Token: "dapi-x"}) + if !kc.EnableGeoSpatialSupport { + t.Fatal("omitted EnableGeoSpatialSupport = false, want default true") + } + + WithEnableGeoSpatialSupport(false)(c) + kc = buildKernelConfig(c, kernel.Auth{Mode: kernel.AuthPAT, Token: "dapi-x"}) + if kc.EnableGeoSpatialSupport { + t.Fatal("explicit EnableGeoSpatialSupport=false was not preserved") + } + + WithEnableGeoSpatialSupport(true)(c) + kc = buildKernelConfig(c, kernel.Auth{Mode: kernel.AuthPAT, Token: "dapi-x"}) + if !kc.EnableGeoSpatialSupport { + t.Fatal("explicit EnableGeoSpatialSupport=true was not preserved") + } + }) + t.Run("experimental TLS fields forwarded", func(t *testing.T) { c := baseKernelConfig() c.KernelExperimental = &config.KernelExperimentalConfig{ diff --git a/kernel_e2e_test.go b/kernel_e2e_test.go index 3a6cf952..18eea660 100644 --- a/kernel_e2e_test.go +++ b/kernel_e2e_test.go @@ -192,15 +192,11 @@ func TestKernelE2ETLSSkipVerify(t *testing.T) { // row. Each case selects a single value and compares the scanned result. NULL is // covered as its own case. func TestKernelE2EDataTypes(t *testing.T) { - db := kernelTestDB(t) + // Keep the legacy string expectations explicit; the native default is + // exercised separately by TestKernelE2EGeospatialRepresentations. + db := kernelTestDBWith(t, WithEnableGeoSpatialSupport(false)) 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 @@ -226,18 +222,13 @@ func TestKernelE2EDataTypes(t *testing.T) { {"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/GEOGRAPHY as a WKT/WKB string. + // GEOMETRY/GEOGRAPHY as WKT/EWKT in this explicit string-mode suite. {"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)"}, + {"geography", "st_geogfromtext('POINT(1 2)')", "SRID=4326;POINT(1 2)"}, } for _, c := range cases { @@ -247,17 +238,9 @@ func TestKernelE2EDataTypes(t *testing.T) { if c.name == "variant" { skipOnReyden(t, "VARIANT over inline Arrow is not decodable on the reyden reference engine (reyden-sea gap; dbsql-sea/Thrift read it fine)") } - // Geospatial (st_point/GEOMETRY/GEOGRAPHY) is UNSUPPORTED_FEATURE on the reyden - // real-time engine — a Lakehouse//RT limitation, not a driver gap. - if c.name == "geometry" || c.name == "geography" { - skipOnReyden(t, "geospatial (st_point/GEOMETRY/GEOGRAPHY) is unsupported on the reyden real-time engine — UNSUPPORTED_FEATURE, Lakehouse//RT limitation") - } 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) { diff --git a/kernel_experimental_test.go b/kernel_experimental_test.go index 3f936a57..a17647f1 100644 --- a/kernel_experimental_test.go +++ b/kernel_experimental_test.go @@ -40,6 +40,7 @@ var kernelExperimentalFieldDisposition = map[string]string{ "MaxConnections": "forwarded", // set_max_connections "MaxChunksInMemory": "forwarded", // set_session_conf (cloudfetch_max_chunks_in_memory, client-only) "DecimalAsFloat": "forwarded", // kernel.Config.DecimalAsFloat → kernelOp → arrowscan (client-side scan choice) + "EnableGeoSpatialSupport": "forwarded", // set_enable_geospatial_support; false is preserved via pointer "TokenCacheEnabled": "forwarded", // set_u2m_token_cache_config (enabled, nil passphrase) } @@ -100,6 +101,9 @@ func TestWithKernelTLSOptionsSetExperimental(t *testing.T) { {"decimal as float", WithKernelDecimalAsFloat(true), func(k *config.KernelExperimentalConfig) bool { return k.DecimalAsFloat }}, + {"disable geospatial support", WithEnableGeoSpatialSupport(false), func(k *config.KernelExperimentalConfig) bool { + return k.EnableGeoSpatialSupport != nil && !*k.EnableGeoSpatialSupport + }}, {"token cache", WithTokenCache(true), func(k *config.KernelExperimentalConfig) bool { return k.TokenCacheEnabled }}, @@ -138,6 +142,7 @@ func TestWithKernelOptionsRejectedOnThriftPath(t *testing.T) { {"max connections", WithKernelMaxConnections(37)}, {"max chunks in memory", WithKernelMaxChunksInMemory(4)}, {"decimal as float", WithKernelDecimalAsFloat(true)}, + {"disable geospatial support", WithEnableGeoSpatialSupport(false)}, {"token cache", WithTokenCache(true)}, } for _, tc := range cases { @@ -171,6 +176,24 @@ func TestWithKernelOptionsRejectedOnThriftPath(t *testing.T) { } }) } + + for _, dsn := range []string{ + "token:supersecret@example.cloud.databricks.com:443/sql/1.0/endpoints/12346a5b5b0e123a?enableGeoSpatialSupport=true", + "token:supersecret@example.cloud.databricks.com:443/sql/1.0/endpoints/12346a5b5b0e123a?enableGeoSpatialSupport=false", + } { + ucfg, err := config.ParseDSN(dsn) + if err != nil { + t.Fatalf("ParseDSN: %v", err) + } + c, err := NewConnector(withUserConfig(ucfg)) + if err != nil { + t.Fatalf("NewConnector: %v", err) + } + _, err = c.Connect(context.Background()) + if !errors.Is(err, dbsqlerr.ErrRequiresKernelBackend) { + t.Errorf("DSN %q error = %v, want ErrRequiresKernelBackend", dsn, err) + } + } } // WithKernelTrustedCerts must copy the PEM defensively, so a caller mutating its @@ -218,6 +241,7 @@ func TestWithKernelClientCertificateCopiesPEM(t *testing.T) { // mutation reach another. func TestKernelExperimentalDeepCopy(t *testing.T) { maxConnections := 37 + enableGeoSpatialSupport := false orig := &config.KernelExperimentalConfig{ TLSTrustedCertsPEM: []byte("ca-bundle"), TLSClientCertPEM: []byte("client-cert"), @@ -231,6 +255,7 @@ func TestKernelExperimentalDeepCopy(t *testing.T) { RetryOverallTimeout: 5 * time.Minute, MaxConnections: &maxConnections, MaxChunksInMemory: 4, + EnableGeoSpatialSupport: &enableGeoSpatialSupport, TokenCacheEnabled: true, } cp := orig.DeepCopy() @@ -252,6 +277,11 @@ func TestKernelExperimentalDeepCopy(t *testing.T) { if cp.MaxChunksInMemory != 4 { t.Errorf("DeepCopy lost MaxChunksInMemory: %v", cp.MaxChunksInMemory) } + if cp.EnableGeoSpatialSupport == nil || *cp.EnableGeoSpatialSupport { + t.Errorf("DeepCopy lost EnableGeoSpatialSupport: %v", cp.EnableGeoSpatialSupport) + } else if cp.EnableGeoSpatialSupport == orig.EnableGeoSpatialSupport { + t.Error("DeepCopy aliased EnableGeoSpatialSupport pointer") + } if !cp.TokenCacheEnabled { t.Errorf("DeepCopy lost TokenCacheEnabled: %v", cp.TokenCacheEnabled) } diff --git a/kernel_geospatial_e2e_test.go b/kernel_geospatial_e2e_test.go new file mode 100644 index 00000000..248139d2 --- /dev/null +++ b/kernel_geospatial_e2e_test.go @@ -0,0 +1,159 @@ +//go:build cgo && databricks_kernel + +package dbsql + +import ( + "bytes" + "context" + "database/sql" + "database/sql/driver" + "encoding/base64" + "encoding/hex" + "encoding/json" + "io" + "reflect" + "testing" + + "github.com/apache/arrow/go/v12/arrow/array" + dbsqlrows "github.com/databricks/databricks-sql-go/rows" +) + +// Runs unchanged on the DBSQL and Reyden legs. A server rejection is a failure, +// not a stale geospatial skip. The kernel archive must include the geospatial PR. +func TestKernelE2EGeospatialRepresentations(t *testing.T) { + const query = `SELECT st_geomfromtext('POINT(1 2)',4326) AS geom, + st_geogfromtext('POINT(1 2)') AS geog, st_geomfromtext(NULL,4326) AS missing, + named_struct('g',st_geomfromtext('POINT(1 2)',4326), + 'blob',unhex('0001027F80FF'),'n',st_geomfromtext(NULL,4326)) AS nested` + wkb, _ := hex.DecodeString("0101000000000000000000f03f0000000000000040") + ewkb, _ := hex.DecodeString("0101000020e6100000000000000000f03f0000000000000040") + blob, _ := hex.DecodeString("0001027F80FF") + const text = "SRID=4326;POINT(1 2)" + for _, mode := range []struct { + name string + native bool + extra []ConnOption + }{ + {"default", true, nil}, + {"enabled", true, []ConnOption{WithEnableGeoSpatialSupport(true)}}, + {"disabled", false, []ConnOption{WithEnableGeoSpatialSupport(false)}}, + } { + t.Run(mode.name, func(t *testing.T) { + db := kernelTestDBWith(t, append(mode.extra, WithSessionParams(map[string]string{"timezone": "UTC"}))...) + defer db.Close() + rows, err := db.QueryContext(context.Background(), query) + if err != nil { + t.Fatal(err) + } + defer rows.Close() + columns, err := rows.ColumnTypes() + if err != nil { + t.Fatal(err) + } + for i, family := range []string{"GEOMETRY", "GEOGRAPHY"} { + wantType, wantScan := "STRING", reflect.TypeOf("") + if mode.native { + wantType, wantScan = family, reflect.TypeOf(sql.RawBytes{}) + } + if columns[i].DatabaseTypeName() != wantType || columns[i].ScanType() != wantScan { + t.Fatalf("column %d: type=%s scan=%v, want %s/%v", i, columns[i].DatabaseTypeName(), columns[i].ScanType(), wantType, wantScan) + } + } + if !rows.Next() { + t.Fatalf("missing row: %v", rows.Err()) + } + var geom, geog, missing any + var nested string + if err := rows.Scan(&geom, &geog, &missing, &nested); err != nil { + t.Fatal(err) + } + for _, value := range []any{geom, geog} { + if mode.native { + if data, ok := value.([]byte); !ok || !bytes.Equal(data, ewkb) { + t.Fatalf("got %T %v, want EWKB %x", value, value, ewkb) + } + } else if value != text { + t.Fatalf("got %T %v, want EWKT", value, value) + } + } + if missing != nil { + t.Fatalf("SQL NULL = %#v", missing) + } + var gotNested map[string]any + if err := json.Unmarshal([]byte(nested), &gotNested); err != nil { + t.Fatalf("invalid public JSON %q: %v", nested, err) + } + var geoLeaf any = text + if mode.native { + geoLeaf = map[string]any{"srid": float64(4326), "wkb": base64.StdEncoding.EncodeToString(wkb)} + } + wantNested := map[string]any{"g": geoLeaf, "blob": base64.StdEncoding.EncodeToString(blob), "n": nil} + if !reflect.DeepEqual(gotNested, wantNested) { + t.Fatalf("nested = %s, want %v", nested, wantNested) + } + if rows.Next() || rows.Err() != nil { + t.Fatalf("unexpected trailing row/error: %v", rows.Err()) + } + rows.Close() + + conn, err := db.Conn(context.Background()) + if err != nil { + t.Fatal(err) + } + defer conn.Close() + var raw driver.Rows + if err := conn.Raw(func(c any) error { + var err error + raw, err = c.(driver.QueryerContext).QueryContext(context.Background(), query, nil) + return err + }); err != nil { + t.Fatal(err) + } + defer raw.Close() + batches, err := raw.(dbsqlrows.Rows).GetArrowBatches(context.Background()) + if err != nil { + t.Fatal(err) + } + defer batches.Close() + var count int64 + for batches.HasNext() { + record, err := batches.Next() + if err == io.EOF { + break + } + if err != nil { + t.Fatal(err) + } + func() { + defer record.Release() + for row := 0; row < int(record.NumRows()); row++ { + for col := 0; col < 2; col++ { + if mode.native { + value, ok := record.Column(col).(*array.Struct) + if !ok { + t.Fatalf("raw geo type = %s, want struct", record.Column(col).DataType()) + } + if value.Field(0).(*array.Int32).Value(row) != 4326 || !bytes.Equal(value.Field(1).(*array.Binary).Value(row), wkb) { + t.Fatal("raw Arrow SRID/WKB changed") + } + } else if value, ok := record.Column(col).(*array.String); !ok || value.Value(row) != text { + t.Fatalf("raw text = %v", record.Column(col)) + } + } + if !record.Column(2).IsNull(row) { + t.Fatal("raw Arrow NULL was lost") + } + value := record.Column(3).(*array.Struct) + if !bytes.Equal(value.Field(1).(*array.Binary).Value(row), blob) || !value.Field(2).IsNull(row) { + t.Fatal("raw Arrow binary/null siblings changed") + } + } + count += record.NumRows() + }() + } + if count != 1 { + t.Fatalf("raw Arrow row count = %d, want 1", count) + } + }) + } +} diff --git a/kernel_parity_test.go b/kernel_parity_test.go index ed9e52ed..cdc6fbef 100644 --- a/kernel_parity_test.go +++ b/kernel_parity_test.go @@ -103,7 +103,7 @@ func TestKernelE2EInterval(t *testing.T) { 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 + // collapse to STRING on both backends in explicit string mode, and both // interval types are server/Go-stringified to STRING — included so the parity // covers those collapses too. const query = "SELECT " + @@ -120,7 +120,7 @@ func TestKernelThriftColumnTypeParity(t *testing.T) { // 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) + kernelDB := kernelTestDBWith(t, WithEnableGeoSpatialSupport(false)) defer kernelDB.Close() thriftDB := thriftTestDB(t) defer thriftDB.Close()