From cde43b8d08db7dfd584e16e95483ca4d12d9aa2e Mon Sep 17 00:00:00 2001 From: Abhishek Dey Date: Fri, 11 Sep 2026 01:42:38 +0530 Subject: [PATCH 1/4] databricks: support Lakehouse//RT, auto-detecting the SEA backend Lakehouse//RT warehouses only speak the Statement Execution API (SEA) and reject the Thrift protocol, so Rill (pinned to databricks-sql-go v1.10.0) could not connect to them at all. Bump databricks-sql-go v1.10.0 -> v1.15.1 (adds the SEA/kernel backend) and auto-detect RT: on first connect the connector probes once and, if the warehouse rejects Thrift, transparently switches to the SEA backend. The decision is cached and shared by both the OLAP path and the warehouse ingest path (effectiveDSN), so an RT warehouse works with no configuration. DBSQL warehouses accept Thrift and are unchanged. A `use_kernel` connector property (default false) is available as an explicit override. The SEA backend exports Arrow C Data rather than IPC streams, so the bulk ingest path in warehouse.go falls back from GetArrowIPCStreams to GetArrowBatches (re-serialized to a self-contained Arrow IPC stream via the driver's Arrow v12 writer, which the existing v18 ipc.Reader consumes) when the driver returns ErrNotSupportedByKernel. The OLAP path already worked over SEA. information_schema.Lookup previously JOINed information_schema.tables and columns; that join forces a shuffle that RT's Photon rejects (PHOTON_INTERNAL_ERROR, retry unsupported), breaking the schema browser. Split it into two filtered point-lookups (no shuffle); equivalent on DBSQL. Adds unit tests for DSN resolution, the RT-detection predicate, and the useKernel DSN helper, and documents the behavior. Co-authored-by: Isaac --- .../connectors/data-source/databricks.mdx | 9 ++ go.mod | 31 +++-- go.sum | 58 +++++---- runtime/drivers/databricks/databricks.go | 112 +++++++++++++++++- runtime/drivers/databricks/databricks_test.go | 69 +++++++++++ .../drivers/databricks/information_schema.go | 37 +++--- runtime/drivers/databricks/warehouse.go | 108 +++++++++++++++-- 7 files changed, 363 insertions(+), 61 deletions(-) create mode 100644 runtime/drivers/databricks/databricks_test.go diff --git a/docs/docs/developers/build/connectors/data-source/databricks.mdx b/docs/docs/developers/build/connectors/data-source/databricks.mdx index 0d586b0c2290..561157c3bffd 100644 --- a/docs/docs/developers/build/connectors/data-source/databricks.mdx +++ b/docs/docs/developers/build/connectors/data-source/databricks.mdx @@ -56,6 +56,15 @@ catalog: "main" # optional schema: "default" # optional ``` +:::info Lakehouse//RT warehouses +Lakehouse//RT (real-time) warehouses only support the Statement Execution API (SEA) +protocol and reject the default Thrift protocol. Rill **auto-detects** this: if a +warehouse rejects Thrift, the connector transparently switches to the SEA backend — no +configuration needed. (You can also force it with `use_kernel: true`.) SEA requires a +Rill build that links the `databricks_kernel` backend; otherwise the connection fails +with a clear error. Ordinary DBSQL warehouses keep using Thrift. +::: + **Step 2: Add credentials to `.env`** ```bash diff --git a/go.mod b/go.mod index dbd145526c6a..74b1dd6ad632 100644 --- a/go.mod +++ b/go.mod @@ -22,6 +22,7 @@ require ( github.com/alicebob/miniredis v2.5.0+incompatible github.com/anthropics/anthropic-sdk-go v1.19.0 github.com/apache/arrow-go/v18 v18.5.1 + github.com/apache/arrow/go/v12 v12.0.1 github.com/aws/aws-sdk-go-v2 v1.41.5 github.com/aws/aws-sdk-go-v2/config v1.32.12 github.com/aws/aws-sdk-go-v2/credentials v1.19.12 @@ -37,7 +38,7 @@ require ( github.com/c2h5oh/datasize v0.0.0-20220606134207-859f65c6625b github.com/confluentinc/confluent-kafka-go/v2 v2.2.0 github.com/coreos/go-oidc/v3 v3.17.0 - github.com/databricks/databricks-sql-go v1.10.0 + github.com/databricks/databricks-sql-go v1.15.1 github.com/dgraph-io/ristretto v0.1.1 github.com/docker/go-connections v0.6.0 github.com/duckdb/duckdb-go/v2 v2.10504.0 @@ -120,12 +121,12 @@ require ( go.uber.org/zap/exp v0.2.0 gocloud.dev v0.36.0 golang.org/x/exp v0.0.0-20260410095643-746e56fc9e2f - golang.org/x/net v0.56.0 + golang.org/x/net v0.58.0 golang.org/x/oauth2 v0.36.0 - golang.org/x/sync v0.21.0 - golang.org/x/sys v0.46.0 - golang.org/x/term v0.44.0 - golang.org/x/text v0.38.0 + golang.org/x/sync v0.22.0 + golang.org/x/sys v0.47.0 + golang.org/x/term v0.45.0 + golang.org/x/text v0.41.0 google.golang.org/api v0.230.0 google.golang.org/genai v1.44.0 google.golang.org/genproto/googleapis/api v0.0.0-20260526163538-3dc84a4a5aaa @@ -140,9 +141,15 @@ require ( ) require ( - github.com/apache/arrow/go/v12 v12.0.1 // indirect github.com/aws/aws-sdk-go-v2/service/signin v1.0.8 // indirect github.com/cenkalti/backoff/v5 v5.0.3 // indirect + github.com/databricks/databricks-sql-kernel-bindings/lib/darwin_amd64 v1.0.0 // indirect + github.com/databricks/databricks-sql-kernel-bindings/lib/darwin_arm64 v1.0.0 // indirect + github.com/databricks/databricks-sql-kernel-bindings/lib/linux_amd64 v1.0.0 // indirect + github.com/databricks/databricks-sql-kernel-bindings/lib/linux_arm v1.0.0 // indirect + github.com/databricks/databricks-sql-kernel-bindings/lib/linux_arm64 v1.0.0 // indirect + github.com/databricks/databricks-sql-kernel-bindings/lib/windows_amd64 v1.0.0 // indirect + github.com/databricks/databricks-sql-kernel-bindings/lib/windows_arm64 v1.0.0 // indirect github.com/dnephin/pflag v1.0.7 // indirect github.com/docker/compose/v5 v5.1.2 // indirect github.com/duckdb/duckdb-go-bindings/lib/darwin-amd64 v0.10504.0 // indirect @@ -214,7 +221,7 @@ require ( github.com/alicebob/gopher-json v0.0.0-20230218143504-906a9b012302 // indirect github.com/andybalholm/brotli v1.2.0 // indirect github.com/apache/arrow/go/v15 v15.0.2 // indirect - github.com/apache/thrift v0.22.0 // indirect + github.com/apache/thrift v0.23.0 // indirect github.com/aws/aws-sdk-go v1.49.0 // indirect github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.8 // indirect github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.20 // indirect @@ -420,11 +427,11 @@ require ( go.opentelemetry.io/proto/otlp v1.10.0 // indirect go.uber.org/goleak v1.3.0 // indirect go.yaml.in/yaml/v3 v3.0.4 // indirect - golang.org/x/crypto v0.53.0 // indirect - golang.org/x/mod v0.36.0 // indirect - golang.org/x/telemetry v0.0.0-20260508192327-42602be52be6 // indirect + golang.org/x/crypto v0.55.0 // indirect + golang.org/x/mod v0.40.0 // indirect + golang.org/x/telemetry v0.0.0-20260811182544-a038080d80e5 // indirect golang.org/x/time v0.14.0 // indirect - golang.org/x/tools v0.45.0 // indirect + golang.org/x/tools v0.49.0 // indirect golang.org/x/xerrors v0.0.0-20240903120638-7835f813f4da // indirect google.golang.org/genproto v0.0.0-20250303144028-a0af3efb3deb // indirect google.golang.org/genproto/googleapis/rpc v0.0.0-20260523011958-0a33c5d7ca68 // indirect diff --git a/go.sum b/go.sum index b7441e182b0b..d73b0ffd70dc 100644 --- a/go.sum +++ b/go.sum @@ -795,8 +795,8 @@ github.com/apache/arrow/go/v12 v12.0.1/go.mod h1:weuTY7JvTG/HDPtMQxEUp7pU73vkLWM github.com/apache/arrow/go/v15 v15.0.2 h1:60IliRbiyTWCWjERBCkO1W4Qun9svcYoZrSLcyOsMLE= github.com/apache/arrow/go/v15 v15.0.2/go.mod h1:DGXsR3ajT524njufqf95822i+KTh+yea1jass9YXgjA= github.com/apache/thrift v0.16.0/go.mod h1:PHK3hniurgQaNMZYaCLEqXKsYK8upmhPbmdP2FXSqgU= -github.com/apache/thrift v0.22.0 h1:r7mTJdj51TMDe6RtcmNdQxgn9XcyfGDOzegMDRg47uc= -github.com/apache/thrift v0.22.0/go.mod h1:1e7J/O1Ae6ZQMTYdy9xa3w9k+XHWPfRvdPyJeynQ+/g= +github.com/apache/thrift v0.23.0 h1:wKR6YnefQSEnxpEfmgTPuJibNG4bF0p2TK34tHLWi3s= +github.com/apache/thrift v0.23.0/go.mod h1:zPt6WxgvTOM6hF92y8C+MkEM5LMxZuk4JcQOiU4Esvs= github.com/armon/circbuf v0.0.0-20150827004946-bbbad097214e/go.mod h1:3U/XgcO3hCbHZ8TKRvWD2dDTCfh9M9ya+I9JpbB7O8o= github.com/armon/consul-api v0.0.0-20180202201655-eb2c6b5be1b6/go.mod h1:grANhF5doyWs3UAsr3K4I6qtAmlQcZDesFNEHPZAzj8= github.com/armon/go-metrics v0.0.0-20180917152333-f0300d1749da/go.mod h1:Q73ZrmVTwzkszR9V5SSuryQ31EELlFMUz1kKyl939pY= @@ -1131,8 +1131,22 @@ github.com/d2g/dhcp4server v0.0.0-20181031114812-7d4a0a7f59a5/go.mod h1:Eo87+Kg/ github.com/d2g/hardwareaddr v0.0.0-20190221164911-e7d9fbe030e4/go.mod h1:bMl4RjIciD2oAxI7DmWRx6gbeqrkoLqv3MV0vzNad+I= github.com/danieljoos/wincred v1.2.3 h1:v7dZC2x32Ut3nEfRH+vhoZGvN72+dQ/snVXo/vMFLdQ= github.com/danieljoos/wincred v1.2.3/go.mod h1:6qqX0WNrS4RzPZ1tnroDzq9kY3fu1KwE7MRLQK4X0bs= -github.com/databricks/databricks-sql-go v1.10.0 h1:U17EKVC+hLP87swFMe2N6UUVektwUgTvT2pMDaDc46g= -github.com/databricks/databricks-sql-go v1.10.0/go.mod h1:qC010ucrtqrNXY2UOcoczbfPD4gJ1jr1y6TL7iqyxPk= +github.com/databricks/databricks-sql-go v1.15.1 h1:KUP/x8+jUTPbki8ChB4dJySB4TLhuQFhqRf2SocqnJc= +github.com/databricks/databricks-sql-go v1.15.1/go.mod h1:+aTNQK37V5esFw0jaMZPJ6GKveMcsnSNPCq3iRA3AOc= +github.com/databricks/databricks-sql-kernel-bindings/lib/darwin_amd64 v1.0.0 h1:BJzNEXT61x2k2yOq/l6XxmyI/oaVDSwA+SdqG/oc2jM= +github.com/databricks/databricks-sql-kernel-bindings/lib/darwin_amd64 v1.0.0/go.mod h1:ceyJVgJAbNmkwc/2YMUNGuw2gd8nHZwZhP2xpN84VBg= +github.com/databricks/databricks-sql-kernel-bindings/lib/darwin_arm64 v1.0.0 h1:MWO7E0fFgSVvZLswQRYwtTchgklXWGDRWtAuxux//RE= +github.com/databricks/databricks-sql-kernel-bindings/lib/darwin_arm64 v1.0.0/go.mod h1:Cr0N6/u4qDUvdbtqm4sVflN4ZXO8qdDcxEBhF5WB9PI= +github.com/databricks/databricks-sql-kernel-bindings/lib/linux_amd64 v1.0.0 h1:bRmxDGVuV9UURuyY6Z+dSBFHqGASWucEKWcZWb5ftBY= +github.com/databricks/databricks-sql-kernel-bindings/lib/linux_amd64 v1.0.0/go.mod h1:KqeZQ/C/GoKAqoGL97/KT3BwLoqsxgSl3qJjje1F7TU= +github.com/databricks/databricks-sql-kernel-bindings/lib/linux_arm v1.0.0 h1:00FOtv6k45rJv/DPFlu94lRb4Bbx6TmzOdTS2nChLOI= +github.com/databricks/databricks-sql-kernel-bindings/lib/linux_arm v1.0.0/go.mod h1:5/HsocwYCRUmO4sZKnLTBb/W+ItCdp3Z8lfssNdfqhM= +github.com/databricks/databricks-sql-kernel-bindings/lib/linux_arm64 v1.0.0 h1:BAvQvVJTsavY2tlsELtDJiViDw16IuKd755nrsLndQQ= +github.com/databricks/databricks-sql-kernel-bindings/lib/linux_arm64 v1.0.0/go.mod h1:/1vEJKUPnrWx566kI/Ifl00GrhSyp798Ztl6rjtymCI= +github.com/databricks/databricks-sql-kernel-bindings/lib/windows_amd64 v1.0.0 h1:zRPyxXJRhzG9w1dYqdS5SMhxq8LF1H/tARhOHPpHHSE= +github.com/databricks/databricks-sql-kernel-bindings/lib/windows_amd64 v1.0.0/go.mod h1:OaUVov84uK1IJz1gNk5TIoae0rj0vx+4N1qqRT88MdU= +github.com/databricks/databricks-sql-kernel-bindings/lib/windows_arm64 v1.0.0 h1:Wb7PH5nuSCpPJNYORcomEGZ1jq1pIQaWMAY2puLFfeg= +github.com/databricks/databricks-sql-kernel-bindings/lib/windows_arm64 v1.0.0/go.mod h1:Z8JD951Nk66P1J8hL6YhkyMp1YAnrRzxE2JMTYRblxc= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM= @@ -2632,8 +2646,8 @@ golang.org/x/crypto v0.0.0-20211108221036-ceb1ce70b4fa/go.mod h1:GvvjBRRGRdwPK5y golang.org/x/crypto v0.0.0-20220622213112-05595931fe9d/go.mod h1:IxCIyHEi3zRg3s0A5j5BB6A9Jmi73HwBIUl50j+osU4= golang.org/x/crypto v0.19.0/go.mod h1:Iy9bg/ha4yyC70EfRS8jz+B6ybOBKMaSxLj6P6oBDfU= golang.org/x/crypto v0.20.0/go.mod h1:Xwo95rrVNIoSMx9wa1JroENMToLWn3RNVrTBpLHgZPQ= -golang.org/x/crypto v0.53.0 h1:QZ4Muo8THX6CizN2vPPd5fBGHyogrdK9fG4wLPFUsto= -golang.org/x/crypto v0.53.0/go.mod h1:DNLU434OwVakk9PzuwV8w62mAJpRJL3vsgcfp4Qnsio= +golang.org/x/crypto v0.55.0 h1:+KWHjbgOaAQ66dh/YlkZKHlz9ZUlq61AFirAR9ntP8M= +golang.org/x/crypto v0.55.0/go.mod h1:uq0V9dE/fzQuJtbnL+2EhWOE63vo164FY8xqEnV9xis= golang.org/x/exp v0.0.0-20180321215751-8460e604b9de/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA= golang.org/x/exp v0.0.0-20180807140117-3d87b88a115f/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA= golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA= @@ -2694,8 +2708,8 @@ golang.org/x/mod v0.5.1/go.mod h1:5OXOZSfqPIIbmVBIIKWRFfZjPR0E5r58TLhUjH0a2Ro= golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4/go.mod h1:jJ57K6gSWd91VN4djpZkiMVwK6gcyfeH4XE8wZrZaV4= golang.org/x/mod v0.7.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs= golang.org/x/mod v0.8.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs= -golang.org/x/mod v0.36.0 h1:JJjpVx6myfUsUdAzZuOSTTmRE0PfZeNWzzvKrP7amb4= -golang.org/x/mod v0.36.0/go.mod h1:moc6ELqsWcOw5Ef3xVprK5ul/MvtVvkIXLziUOICjUQ= +golang.org/x/mod v0.40.0 h1:hUv+3cXcdRHz08UmSiOob7sadHig73uo5bkXxQ/tvUs= +golang.org/x/mod v0.40.0/go.mod h1:0/weTWkPWGBikyTWAX3dkjVztMmBA5hM0DH6BElSupE= golang.org/x/net v0.0.0-20180724234803-3673e40ba225/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= golang.org/x/net v0.0.0-20180826012351-8a410e7b638d/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= golang.org/x/net v0.0.0-20180906233101-161cd47e91fd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= @@ -2776,8 +2790,8 @@ golang.org/x/net v0.7.0/go.mod h1:2Tu9+aMcznHK/AK1HMvgo6xiTLG5rD5rZLDS+rp2Bjs= golang.org/x/net v0.8.0/go.mod h1:QVkue5JL9kW//ek3r6jTKnTFis1tRmNAW2P1shuFdJc= golang.org/x/net v0.10.0/go.mod h1:0qNGK6F8kojg2nk9dLZ2mShWaEBan6FAoqfSigmmuDg= golang.org/x/net v0.21.0/go.mod h1:bIjVDfnllIU7BJ2DNgfnXvpSvtn8VRwhlsaeUTyUS44= -golang.org/x/net v0.56.0 h1:Rw8j/hFzGvJUZwNBXnAtf5sVDVt+65SK2C7IxCxZt5o= -golang.org/x/net v0.56.0/go.mod h1:D3Ku6r+V6JROoZK144D2XfMHFcMq/0zSfLelVTCFKec= +golang.org/x/net v0.58.0 h1:ynWG7rqYi4ccpTEuPZ2QGWHktVEM9DMCj9yzDE0Q7To= +golang.org/x/net v0.58.0/go.mod h1:YwCddHnFlT7eLQqVprV19OnhLGtc5xOKgE0RyqgfWAU= golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U= golang.org/x/oauth2 v0.0.0-20190226205417-e64efc72b421/go.mod h1:gOpvHmFTYa4IltrdGE7lF6nIHvwfUNPOp7c8zoXwtLw= golang.org/x/oauth2 v0.0.0-20190604053449-0f29369cfe45/go.mod h1:gOpvHmFTYa4IltrdGE7lF6nIHvwfUNPOp7c8zoXwtLw= @@ -2824,8 +2838,8 @@ golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4/go.mod h1:RxMgew5VJxzue5/jJ golang.org/x/sync v0.0.0-20220819030929-7fc1605a5dde/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20220929204114-8fcdb60fdcc0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.1.0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= -golang.org/x/sync v0.21.0 h1:HLII4xRRTtCRkxYp4HNFF0Js/Og6q2i++KXbg0gHCwM= -golang.org/x/sync v0.21.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= +golang.org/x/sync v0.22.0 h1:SZjpbeLmrCk4xhRSZFNZW5gFUeCeFgjekvI/+gfScek= +golang.org/x/sync v0.22.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= golang.org/x/sys v0.0.0-20180823144017-11551d06cbcc/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20180905080454-ebe1bf3edb33/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= @@ -2967,10 +2981,10 @@ golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.17.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= -golang.org/x/sys v0.46.0 h1:noSf2Fq6F8DBgS+LysIkx7rIExoNHJsxOAtPp4rthXw= -golang.org/x/sys v0.46.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= -golang.org/x/telemetry v0.0.0-20260508192327-42602be52be6 h1:HjU6IWBiAgRIdAJ9/y1rwCn+UELEmwV+VsTLzj/W4sE= -golang.org/x/telemetry v0.0.0-20260508192327-42602be52be6/go.mod h1:Eqhaxk/wZsWEH8CRxLwj6xzEJbz7k1EFGqx7nyCoabE= +golang.org/x/sys v0.47.0 h1:o7XGOvZQCADBQQ4Y7VNq2dRWQR7JmOUW8Kxx4ZsNgWs= +golang.org/x/sys v0.47.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= +golang.org/x/telemetry v0.0.0-20260811182544-a038080d80e5 h1:ZUSxONxc981v7AW7QUg+I9WwZzSTTJ019ENBYr5pV/Q= +golang.org/x/telemetry v0.0.0-20260811182544-a038080d80e5/go.mod h1:LVehoXe41cL5SCVQilsV7Gg6BNG+Js6P9PhSbYTIUkQ= golang.org/x/term v0.0.0-20201117132131-f5c789dd3221/go.mod h1:Nr5EML6q2oocZ2LXRh80K7BxOlk5/8JxuGnuhpl+muw= golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo= golang.org/x/term v0.0.0-20210220032956-6a3ed077a48d/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo= @@ -2983,8 +2997,8 @@ golang.org/x/term v0.5.0/go.mod h1:jMB1sMXY+tzblOD4FWmEbocvup2/aLOaQEp7JmGp78k= golang.org/x/term v0.6.0/go.mod h1:m6U89DPEgQRMq3DNkDClhWw02AUbt2daBVO4cn4Hv9U= golang.org/x/term v0.8.0/go.mod h1:xPskH00ivmX89bAKVGSKKtLOWNx2+17Eiy94tnKShWo= golang.org/x/term v0.17.0/go.mod h1:lLRBjIVuehSbZlaOtGMbcMncT+aqLLLmKrsjNrUguwk= -golang.org/x/term v0.44.0 h1:0rLvDRCtNj0gZkyIXhCyOb2OAzEhLVqc4B+hrsBhrmc= -golang.org/x/term v0.44.0/go.mod h1:7ze4MdzUzLXpSAoFP1H0bOI9aXDqveSvatT5vKcFh2Y= +golang.org/x/term v0.45.0 h1:NwWyBmoJCbfTHpxrWoZ9C6/VxOf7ic219I8xZZFdrf0= +golang.org/x/term v0.45.0/go.mod h1:9aqxs0blBcrm/n0L9QW0aRVD+ktan8ssZromtqJC43w= golang.org/x/text v0.0.0-20170915032832-14c0d48ead0c/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= golang.org/x/text v0.3.1-0.20180807135948-17ff2d5776d2/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= @@ -3002,8 +3016,8 @@ golang.org/x/text v0.7.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8= golang.org/x/text v0.8.0/go.mod h1:e1OnstbJyHTd6l/uOt8jFFHp6TRDWZR/bV3emEE/zU8= golang.org/x/text v0.9.0/go.mod h1:e1OnstbJyHTd6l/uOt8jFFHp6TRDWZR/bV3emEE/zU8= golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU= -golang.org/x/text v0.38.0 h1:sXmwo9DwP3OK9EZ7PqAdaooSGozfl/3a6/xJcbzPRhE= -golang.org/x/text v0.38.0/go.mod h1:YXZt3QhHUKYT53r2lLKFIVi6Ao1jdzrTR/KQ09qyxF4= +golang.org/x/text v0.41.0 h1:vz/seA0lnX87Othu2f/0L24RcgrXD9/YFTSuGjj3rH8= +golang.org/x/text v0.41.0/go.mod h1:jvf1O8ajNzZqhSrQBPbutR/EB83Cc0CFrezNQIwbb5M= golang.org/x/time v0.0.0-20180412165947-fbb02b2291d2/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ= golang.org/x/time v0.0.0-20181108054448-85acf8d2951c/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ= golang.org/x/time v0.0.0-20190308202827-9d24e82272b4/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ= @@ -3099,8 +3113,8 @@ golang.org/x/tools v0.1.11/go.mod h1:SgwaegtQh8clINPpECJMqnxLv9I09HLqnW3RMqW0CA4 golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc= golang.org/x/tools v0.3.0/go.mod h1:/rWhSS2+zyEVwoJf8YAX6L2f0ntZ7Kn/mGgAWcipA5k= golang.org/x/tools v0.6.0/go.mod h1:Xwgl3UAJ/d3gWutnCtw505GrjyAbvKui8lOU390QaIU= -golang.org/x/tools v0.45.0 h1:18qN3FAooORvApf5XjCXgsuayZOEtXf6JK18I3+ONa8= -golang.org/x/tools v0.45.0/go.mod h1:LuUGqqaXcXMEFEruIVJVm5mgDD8vww/z/SR1gQ4uE/0= +golang.org/x/tools v0.49.0 h1:3NI7VXzL9+1WZD52Dx2ttoPwD5DWrFGpl9mFZDlmisI= +golang.org/x/tools v0.49.0/go.mod h1:SJNXV9DBKT0UbdttsQjbfJlAE/q+y36++zo3uL3N0Oo= golang.org/x/xerrors v0.0.0-20190410155217-1f06c39b4373/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20190513163551-3ee3066db522/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= diff --git a/runtime/drivers/databricks/databricks.go b/runtime/drivers/databricks/databricks.go index 72ed9fc7f444..92f85087c672 100644 --- a/runtime/drivers/databricks/databricks.go +++ b/runtime/drivers/databricks/databricks.go @@ -6,6 +6,7 @@ import ( "fmt" "net/url" "strings" + "sync" dbsqllog "github.com/databricks/databricks-sql-go/logger" "github.com/jmoiron/sqlx" @@ -79,6 +80,14 @@ var spec = drivers.Spec{ Placeholder: "default", Hint: "Schema within the catalog (optional; defaults to the workspace default)", }, + { + Key: "use_kernel", + Type: drivers.BooleanPropertyType, + DisplayName: "Use SEA (Statement Execution API)", + Hint: "Force the SEA backend instead of Thrift. Lakehouse//RT warehouses are " + + "auto-detected and switched to SEA automatically, so this is usually unnecessary. " + + "Requires a Rill build with the databricks_kernel backend (CGO); defaults to false.", + }, }, ImplementsOLAP: true, ImplementsWarehouse: true, @@ -94,6 +103,10 @@ type configProperties struct { Catalog string `mapstructure:"catalog"` Schema string `mapstructure:"schema"` LogQueries bool `mapstructure:"log_queries"` + // UseKernel forces the driver's SEA backend. Lakehouse//RT is auto-detected and + // switched to SEA even when this is false (see getDB), so it's an override. + // Needs a build with the databricks_kernel backend; default false = Thrift. + UseKernel bool `mapstructure:"use_kernel"` } func (c *configProperties) validate() error { @@ -124,10 +137,17 @@ func (c *configProperties) validate() error { func (c *configProperties) resolveDSN() string { if c.DSN != "" { + if c.UseKernel { + return withUseKernel(c.DSN) + } return c.DSN } params := url.Values{} params.Set("timezone", "UTC") + // Opt in to the SEA backend (required for Lakehouse//RT); default is Thrift. + if c.UseKernel { + params.Set("useKernel", "true") + } if c.Catalog != "" { params.Set("catalog", c.Catalog) } @@ -145,6 +165,35 @@ func (c *configProperties) resolveDSN() string { return u.String() } +// withUseKernel forces useKernel=true on a resolved DSN. It is only called when SEA is +// required (explicit use_kernel, or RT auto-detect), so any existing useKernel in the +// DSN — including useKernel=false — is dropped in favor of true. It splits at the first +// '?' (the query separator) and rewrites only the query part; it avoids url.Parse +// because the driver also accepts scheme-less DSNs (token:...@host), which url.Parse +// would misread (treating "token" as the scheme). +func withUseKernel(dsn string) string { + base, query, hasQuery := strings.Cut(dsn, "?") + if !hasQuery || query == "" { + return base + "?useKernel=true" + } + parts := strings.Split(query, "&") + kept := parts[:0] + for _, p := range parts { + if strings.HasPrefix(p, "useKernel=") { + continue // drop any existing useKernel (true or false) + } + kept = append(kept, p) + } + kept = append(kept, "useKernel=true") + return base + "?" + strings.Join(kept, "&") +} + +// rtRequiresSEA reports whether err is a Lakehouse//RT warehouse rejecting the +// Thrift protocol (the driver surfaces the server's message verbatim). +func rtRequiresSEA(err error) bool { + return err != nil && strings.Contains(err.Error(), "not supported for Thrift protocol") +} + func (d driver) Open(_, instanceID string, config map[string]any, st *storage.Client, ac *activity.Client, logger *zap.Logger) (drivers.Handle, error) { if instanceID == "" { return nil, errors.New("databricks driver can't be shared") @@ -187,6 +236,12 @@ type connection struct { db *sqlx.DB // lazily populated using getDB dbErr error dbMu *semaphore.Weighted + + // resolvedDSN is the DSN to connect with, computed by effectiveDSN (auto-detects + // Lakehouse//RT and upgrades to the SEA backend). Empty until a *definitive* + // probe outcome is reached, so a transient probe failure is not cached. + dsnMu sync.Mutex + resolvedDSN string } // Ping implements drivers.Handle. @@ -303,9 +358,64 @@ func (c *connection) getDB(ctx context.Context) (*sqlx.DB, error) { return c.db, c.dbErr } - c.db, c.dbErr = sqlx.Open("databricks", c.config.resolveDSN()) + // Only build (and permanently cache) the shared pool once the backend is + // *definitively* known. On a non-definitive result we surface the probe error + // (which may be a real misconfiguration or a transient blip) without caching it, + // so the caller can retry and we re-probe. c.db is thus never swapped or closed + // here, so a pool already handed to concurrent callers is never closed underneath. + dsn, definitive, probeErr := c.effectiveDSN(ctx) + if !definitive { + if probeErr != nil { + return nil, probeErr + } + return nil, errors.New("databricks: could not determine warehouse protocol") + } + c.db, c.dbErr = sqlx.Open("databricks", dsn) if c.dbErr != nil { return nil, c.dbErr } return c.db, c.dbErr } + +// effectiveDSN resolves the DSN to connect with, auto-detecting Lakehouse//RT: if SEA +// wasn't explicitly requested (use_kernel) and the warehouse rejects Thrift, it upgrades +// the DSN to the SEA backend. It returns whether the result is definitive (safe to +// cache) and, when not, the probe error that prevented a decision — callers surface that +// error and re-resolve next time rather than lock onto the wrong backend or mask a real +// failure. It's the single shared RT-detection point for the OLAP (getDB) and warehouse +// ingest paths, so RT works without any configuration. +func (c *connection) effectiveDSN(ctx context.Context) (dsn string, definitive bool, probeErr error) { + c.dsnMu.Lock() + defer c.dsnMu.Unlock() + if c.resolvedDSN != "" { + return c.resolvedDSN, true, nil + } + + base := c.config.resolveDSN() + if c.config.UseKernel { + c.resolvedDSN = base // already carries useKernel=true + return base, true, nil + } + + probe, err := sqlx.Open("databricks", base) + if err != nil { + return base, false, err // can't even open; don't cache + } + defer probe.Close() //nolint:errcheck // best-effort close of the probe connection + + // Only cache a *definitive* outcome. A clean ping means Thrift works; the RT + // rejection means switch to SEA. Any other ping error is returned uncached (with + // the error) — it may be transient (a blip, re-probed next call) or persistent (a + // bad token / wrong host), and surfacing it verbatim avoids both locking onto the + // wrong backend and masking a real misconfiguration behind auto-detection. + switch perr := probe.PingContext(ctx); { + case perr == nil: + c.resolvedDSN = base + case rtRequiresSEA(perr): + c.logger.Info("databricks: warehouse rejected Thrift (Lakehouse//RT); switching to the SEA backend") + c.resolvedDSN = withUseKernel(base) + default: + return base, false, perr // transient or persistent — surface it, don't cache + } + return c.resolvedDSN, true, nil +} diff --git a/runtime/drivers/databricks/databricks_test.go b/runtime/drivers/databricks/databricks_test.go new file mode 100644 index 000000000000..80c0e6354a6a --- /dev/null +++ b/runtime/drivers/databricks/databricks_test.go @@ -0,0 +1,69 @@ +package databricks + +import ( + "errors" + "fmt" + "strings" + "testing" + + "github.com/stretchr/testify/require" +) + +// TestResolveDSN pins the backend-selection behavior of resolveDSN: +// - by default (use_kernel unset) the DSN carries no useKernel param, so ordinary +// DBSQL warehouses keep using the Thrift backend (backward compatible); +// - with use_kernel=true the DSN carries useKernel=true, selecting the SEA/kernel +// backend required by Lakehouse//RT; +// - a raw DSN is passed through untouched unless use_kernel opts in, and is never +// given a duplicate useKernel param. +func TestResolveDSN(t *testing.T) { + t.Run("params: default has no useKernel", func(t *testing.T) { + c := &configProperties{Host: "h.cloud.databricks.com", HTTPPath: "/sql/1.0/warehouses/w", Token: "t"} + require.NotContains(t, c.resolveDSN(), "useKernel") + }) + + t.Run("params: use_kernel adds useKernel=true", func(t *testing.T) { + c := &configProperties{Host: "h.cloud.databricks.com", HTTPPath: "/sql/1.0/warehouses/w", Token: "t", UseKernel: true} + require.Contains(t, c.resolveDSN(), "useKernel=true") + }) + + t.Run("dsn: passed through unchanged by default", func(t *testing.T) { + dsn := "token:tok@h.cloud.databricks.com:443/sql/1.0/warehouses/w" + c := &configProperties{DSN: dsn} + require.Equal(t, dsn, c.resolveDSN()) + }) + + t.Run("dsn: use_kernel appends with ? when no query", func(t *testing.T) { + c := &configProperties{DSN: "token:tok@h.cloud.databricks.com:443/sql/1.0/warehouses/w", UseKernel: true} + require.Equal(t, "token:tok@h.cloud.databricks.com:443/sql/1.0/warehouses/w?useKernel=true", c.resolveDSN()) + }) + + t.Run("dsn: use_kernel appends with & when query exists", func(t *testing.T) { + c := &configProperties{DSN: "token:tok@h.cloud.databricks.com:443/sql/1.0/warehouses/w?catalog=main", UseKernel: true} + require.Equal(t, "token:tok@h.cloud.databricks.com:443/sql/1.0/warehouses/w?catalog=main&useKernel=true", c.resolveDSN()) + }) + + t.Run("dsn: no duplicate useKernel", func(t *testing.T) { + c := &configProperties{DSN: "token:tok@h.cloud.databricks.com:443/sql/1.0/warehouses/w?useKernel=true", UseKernel: true} + require.Equal(t, 1, strings.Count(c.resolveDSN(), "useKernel=")) + }) +} + +// TestRTRequiresSEA pins the Lakehouse//RT autodetect trigger: only the server's +// "not supported for Thrift protocol" message (however wrapped) flips to SEA; other +// connection errors must not. +func TestRTRequiresSEA(t *testing.T) { + require.True(t, rtRequiresSEA(errors.New("Lakehouse/RT is not supported for Thrift protocol. Please update..."))) + require.True(t, rtRequiresSEA(fmt.Errorf("wrap: %w", errors.New("...is not supported for Thrift protocol")))) + require.False(t, rtRequiresSEA(errors.New("HTTP 403: Invalid access token"))) + require.False(t, rtRequiresSEA(errors.New("connection refused"))) + require.False(t, rtRequiresSEA(nil)) +} + +func TestWithUseKernel(t *testing.T) { + require.Equal(t, "dsn?useKernel=true", withUseKernel("dsn")) + require.Equal(t, "dsn?x=1&useKernel=true", withUseKernel("dsn?x=1")) + require.Equal(t, "dsn?useKernel=true", withUseKernel("dsn?useKernel=true")) // no dup + require.Equal(t, "dsn?useKernel=true", withUseKernel("dsn?useKernel=false")) // overrides false + require.Equal(t, "dsn?a=1&b=2&useKernel=true", withUseKernel("dsn?a=1&useKernel=false&b=2")) // drops mid, re-adds +} diff --git a/runtime/drivers/databricks/information_schema.go b/runtime/drivers/databricks/information_schema.go index 08cead19288b..6c7b9fc85209 100644 --- a/runtime/drivers/databricks/information_schema.go +++ b/runtime/drivers/databricks/information_schema.go @@ -2,6 +2,8 @@ package databricks import ( "context" + "database/sql" + "errors" "fmt" runtimev1 "github.com/rilldata/rill/proto/gen/rill/runtime/v1" @@ -137,34 +139,39 @@ func (c *connection) ListTables(ctx context.Context, database, databaseSchema st func (c *connection) Lookup(ctx context.Context, database, databaseSchema, name string) (*drivers.OlapTable, error) { prefix := catalogPrefix(database) - q := fmt.Sprintf(` - SELECT - CASE WHEN t.table_type = 'VIEW' THEN true ELSE false END AS is_view, - c.column_name, - c.data_type - FROM %sinformation_schema.tables t - JOIN %sinformation_schema.columns c - ON t.table_schema = c.table_schema AND t.table_name = c.table_name - WHERE t.table_schema = ? AND t.table_name = ? - ORDER BY c.ordinal_position - `, prefix, prefix) conn, err := c.getDB(ctx) if err != nil { return nil, err } - rows, err := conn.QueryContext(ctx, q, databaseSchema, name) + // Query the table type and columns separately rather than JOINing + // information_schema.tables and information_schema.columns. The join forces a + // shuffle that Lakehouse//RT's Photon rejects (PHOTON_INTERNAL_ERROR), after + // which RT refuses the retry; two filtered point-lookups avoid the shuffle and + // are equivalent on DBSQL. + var tableType string + err = conn.QueryRowContext(ctx, + fmt.Sprintf("SELECT table_type FROM %sinformation_schema.tables WHERE table_schema = ? AND table_name = ?", prefix), + databaseSchema, name, + ).Scan(&tableType) + if err != nil && !errors.Is(err, sql.ErrNoRows) { + return nil, err + } + + rows, err := conn.QueryContext(ctx, + fmt.Sprintf("SELECT column_name, data_type FROM %sinformation_schema.columns WHERE table_schema = ? AND table_name = ? ORDER BY ordinal_position", prefix), + databaseSchema, name, + ) if err != nil { return nil, err } defer rows.Close() - var isView bool var fields []*runtimev1.StructType_Field var colName, colType string for rows.Next() { - if err := rows.Scan(&isView, &colName, &colType); err != nil { + if err := rows.Scan(&colName, &colType); err != nil { return nil, err } fields = append(fields, &runtimev1.StructType_Field{ @@ -180,7 +187,7 @@ func (c *connection) Lookup(ctx context.Context, database, databaseSchema, name Database: database, DatabaseSchema: databaseSchema, Name: name, - View: isView, + View: tableType == "VIEW", Schema: &runtimev1.StructType{Fields: fields}, }, nil } diff --git a/runtime/drivers/databricks/warehouse.go b/runtime/drivers/databricks/warehouse.go index e9493ba632fb..040dcf5c6660 100644 --- a/runtime/drivers/databricks/warehouse.go +++ b/runtime/drivers/databricks/warehouse.go @@ -1,6 +1,7 @@ package databricks import ( + "bytes" "context" "database/sql" sqld "database/sql/driver" @@ -15,7 +16,11 @@ import ( "github.com/apache/arrow-go/v18/parquet" "github.com/apache/arrow-go/v18/parquet/compress" "github.com/apache/arrow-go/v18/parquet/pqarrow" + // v12 Arrow IPC writer (kernel exports v12 records); aliased to avoid clashing + // with the v18 arrow/ipc import above. + dbipc "github.com/apache/arrow/go/v12/arrow/ipc" "github.com/c2h5oh/datasize" + dbsqlerr "github.com/databricks/databricks-sql-go/errors" dbsqlrows "github.com/databricks/databricks-sql-go/rows" "github.com/mitchellh/mapstructure" "github.com/rilldata/rill/runtime/drivers" @@ -48,7 +53,18 @@ func (c *connection) QueryAsFiles(ctx context.Context, props map[string]any) (ou return nil, err } - db, err := sql.Open("databricks", c.config.resolveDSN()) + // effectiveDSN auto-detects Lakehouse//RT and selects the SEA backend if needed, + // shared with the OLAP path so ingest works without configuration. On a + // non-definitive result, defer like getDB (surface the probe error / retry) rather + // than ingest over a possibly-wrong backend. + dsn, definitive, probeErr := c.effectiveDSN(ctx) + if !definitive { + if probeErr != nil { + return nil, probeErr + } + return nil, errors.New("databricks: could not determine warehouse protocol") + } + db, err := sql.Open("databricks", dsn) if err != nil { return nil, err } @@ -82,7 +98,20 @@ func (c *connection) QueryAsFiles(ctx context.Context, props map[string]any) (ou } }() - ipcStreams, err := rows.(dbsqlrows.Rows).GetArrowIPCStreams(ctx) + // Thrift/DBSQL uses native IPC streams. The SEA backend (Lakehouse//RT) doesn't + // implement them and returns ErrNotSupportedByKernel, so fall back to + // GetArrowBatches via kernelIPCStreams; the parquet path is identical for both. + dr := rows.(dbsqlrows.Rows) + selfDescribing := false // whether each stream carries its own schema (kernel adapter) + ipcStreams, err := dr.GetArrowIPCStreams(ctx) + if errors.Is(err, dbsqlerr.ErrNotSupportedByKernel) { + var batches dbsqlrows.ArrowBatchIterator + batches, err = dr.GetArrowBatches(ctx) + if err == nil { + ipcStreams = &kernelIPCStreams{batches: batches} + selfDescribing = true + } + } if err != nil { return nil, err } @@ -102,12 +131,13 @@ func (c *connection) QueryAsFiles(ctx context.Context, props map[string]any) (ou } return &fileIterator{ - db: db, - conn: conn, - rows: rows, - ipcStreams: ipcStreams, - logger: c.logger, - tempDir: tempDir, + db: db, + conn: conn, + rows: rows, + ipcStreams: ipcStreams, + selfDescribing: selfDescribing, + logger: c.logger, + tempDir: tempDir, }, nil } @@ -116,8 +146,12 @@ type fileIterator struct { conn *sql.Conn rows sqld.Rows ipcStreams dbsqlrows.ArrowIPCStreamIterator - logger *zap.Logger - tempDir string + // selfDescribing is true when each stream carries its own schema message (the + // kernel adapter). For native Thrift streams it's false, and subsequent streams + // are read with ipc.WithSchema to validate cross-stream schema consistency. + selfDescribing bool + logger *zap.Logger + tempDir string totalRecords int64 downloaded bool @@ -125,6 +159,51 @@ type fileIterator struct { var _ drivers.FileIterator = &fileIterator{} +// kernelIPCStreams adapts the SEA backend's ArrowBatchIterator to the +// ArrowIPCStreamIterator this file consumes, re-serializing each batch to a +// self-contained IPC stream with the driver's v12 writer. +type kernelIPCStreams struct { + batches dbsqlrows.ArrowBatchIterator +} + +var _ dbsqlrows.ArrowIPCStreamIterator = &kernelIPCStreams{} + +func (k *kernelIPCStreams) HasNext() bool { return k.batches.HasNext() } + +func (k *kernelIPCStreams) Next() (io.Reader, error) { + rec, err := k.batches.Next() + if err != nil { + return nil, err // propagates io.EOF + } + defer rec.Release() + + var buf bytes.Buffer + w := dbipc.NewWriter(&buf, dbipc.WithSchema(rec.Schema())) + if err := w.Write(rec); err != nil { + _ = w.Close() + return nil, err + } + if err := w.Close(); err != nil { + return nil, err + } + return &buf, nil +} + +func (k *kernelIPCStreams) Close() { k.batches.Close() } + +func (k *kernelIPCStreams) SchemaBytes() ([]byte, error) { + sc, err := k.batches.Schema() + if err != nil { + return nil, err + } + var buf bytes.Buffer + w := dbipc.NewWriter(&buf, dbipc.WithSchema(sc)) + if err := w.Close(); err != nil { + return nil, err + } + return buf.Bytes(), nil +} + // Close implements drivers.FileIterator. func (f *fileIterator) Close() error { if f.ipcStreams != nil { @@ -238,7 +317,14 @@ func (f *fileIterator) Next(ctx context.Context) ([]string, error) { return nil, err } - rdr, err := ipc.NewReader(stream, ipc.WithSchema(schema)) + // Native (Thrift) subsequent streams are validated against the first stream's + // schema; the kernel adapter's streams are self-describing, so skip WithSchema + // there (each carries its own schema message). + var opts []ipc.Option + if !f.selfDescribing { + opts = append(opts, ipc.WithSchema(schema)) + } + rdr, err := ipc.NewReader(stream, opts...) if err != nil { return nil, err } From 4a76a00b9a50da3d4a1345b735714dc5b12b5a7b Mon Sep 17 00:00:00 2001 From: Abhishek Dey Date: Sun, 13 Sep 2026 10:51:31 +0530 Subject: [PATCH 2/4] databricks: add live Lakehouse//RT OLAP test (disabled by default) Adds TestOLAP_LakehouseRT, a live integration test that drives the full Rill Databricks OLAP path against a Lakehouse//RT (Reyden) SQL warehouse over the auto-detected SEA backend. The connector config passes only the DSN (no use_kernel), so a passing query proves the Thrift->SEA autodetection end to end. Because Lakehouse//RT speaks only SEA, the test requires the SEA-via-kernel backend and carries //go:build databricks_kernel (CGO), excluding it from the default Thrift-only build and the standard `go test -short ./...` CI job. It is otherwise gated like the existing Databricks/Snowflake live tests: a leading t.Skip disables it by default, testmode.Expensive keeps it out of normal runs, and it needs RILL_RUNTIME_DATABRICKS_RT_TEST_DSN to point at an RT warehouse. Co-authored-by: Isaac --- runtime/drivers/databricks/olap_rt_test.go | 130 +++++++++++++++++++++ 1 file changed, 130 insertions(+) create mode 100644 runtime/drivers/databricks/olap_rt_test.go diff --git a/runtime/drivers/databricks/olap_rt_test.go b/runtime/drivers/databricks/olap_rt_test.go new file mode 100644 index 000000000000..475b33cf2521 --- /dev/null +++ b/runtime/drivers/databricks/olap_rt_test.go @@ -0,0 +1,130 @@ +//go:build databricks_kernel + +package databricks_test + +import ( + "os" + "testing" + + "github.com/rilldata/rill/runtime/drivers" + "github.com/rilldata/rill/runtime/pkg/activity" + "github.com/rilldata/rill/runtime/storage" + "github.com/rilldata/rill/runtime/testruntime/testmode" + "github.com/stretchr/testify/require" + "go.uber.org/zap" +) + +// TestOLAP_LakehouseRT exercises the full Rill Databricks OLAP path against a live +// Lakehouse//RT (Reyden) SQL warehouse. +// +// Lakehouse//RT rejects the Thrift/HiveServer2 protocol and speaks only SEA (the +// Statement Execution API), which the driver reaches through its SEA-via-kernel +// backend. That backend is compiled in only with `-tags databricks_kernel` +// (CGO_ENABLED=1), so this file carries that build constraint and is excluded from +// the default (Thrift-only) build and from the standard `go test -short ./...` CI job. +// +// The connector config deliberately carries NO use_kernel flag: the driver probes on +// connect and auto-switches Thrift -> SEA when the warehouse rejects Thrift. A query +// that succeeds here therefore proves the RT autodetection path end to end — Thrift +// would have been rejected at connect. +// +// Gated like the other Databricks live tests: it only runs under +// RILL_RUNTIME_TEST_MODE=expensive (skipped under -short) and when +// RILL_RUNTIME_DATABRICKS_RT_TEST_DSN points at a Lakehouse//RT warehouse, e.g. +// +// RILL_RUNTIME_DATABRICKS_RT_TEST_DSN='token:@:443/sql/1.0/warehouses/' +// +// Run with: +// +// CGO_ENABLED=1 RILL_RUNTIME_TEST_MODE=expensive \ +// go test -tags databricks_kernel -run TestOLAP_LakehouseRT ./runtime/drivers/databricks/... +func TestOLAP_LakehouseRT(t *testing.T) { + t.Skip("skipping Lakehouse//RT live test; needs a databricks_kernel build and an RT warehouse (set RILL_RUNTIME_DATABRICKS_RT_TEST_DSN)") + testmode.Expensive(t) + + dsn := os.Getenv("RILL_RUNTIME_DATABRICKS_RT_TEST_DSN") + if dsn == "" { + t.Skip("RILL_RUNTIME_DATABRICKS_RT_TEST_DSN not configured") + } + + _, olap := acquireTestDatabricksRT(t, dsn) + + // Only assert queries Lakehouse//RT actually supports. RT is a SEA reference + // engine and genuinely omits some routines the full DBSQL warehouse ships (e.g. + // current_version() returns UNRESOLVED_ROUTINE / SQLSTATE 42883) — a documented + // RT-vs-DBSQL divergence, not a driver or connectivity failure. + t.Run("scalar_values", func(t *testing.T) { + tests := []struct { + query string + result map[string]any + }{ + {"SELECT TRUE AS bool_val", map[string]any{"bool_val": true}}, + {"SELECT FALSE AS bool_val", map[string]any{"bool_val": false}}, + {"SELECT 'hello' AS string_val", map[string]any{"string_val": "hello"}}, + } + for _, test := range tests { + t.Run(test.query, func(t *testing.T) { + rows, err := olap.Query(t.Context(), &drivers.Statement{Query: test.query}) + require.NoError(t, err) + defer rows.Close() + for rows.Next() { + res := make(map[string]any) + require.NoError(t, rows.MapScan(res)) + require.Equal(t, test.result, res) + } + require.NoError(t, rows.Err()) + }) + } + }) + + // Session/context routines RT does implement: prove real query execution against + // the warehouse (not just a literal the client could fold locally). + t.Run("session_routines", func(t *testing.T) { + for _, q := range []string{ + "SELECT current_catalog() AS v", + "SELECT current_schema() AS v", + "SELECT current_user() AS v", + } { + t.Run(q, func(t *testing.T) { + rows, err := olap.Query(t.Context(), &drivers.Statement{Query: q}) + require.NoError(t, err) + defer rows.Close() + require.True(t, rows.Next(), "expected one row") + res := make(map[string]any) + require.NoError(t, rows.MapScan(res)) + require.NotEmpty(t, res["v"], "expected a non-empty value from %s", q) + require.NoError(t, rows.Err()) + }) + } + }) + + // Schema inspection over the SEA backend. + t.Run("query_schema", func(t *testing.T) { + schema, err := olap.QuerySchema(t.Context(), "SELECT 1 AS int_col, 'x' AS str_col", nil) + require.NoError(t, err) + require.Len(t, schema.Fields, 2) + require.Equal(t, "int_col", schema.Fields[0].Name) + require.Equal(t, "str_col", schema.Fields[1].Name) + }) + + // DryRun issues an EXPLAIN over the SEA backend and must not error. + t.Run("dry_run", func(t *testing.T) { + _, err := olap.Query(t.Context(), &drivers.Statement{Query: "SELECT 1", DryRun: true}) + require.NoError(t, err) + }) +} + +// acquireTestDatabricksRT opens the Databricks driver against a Lakehouse//RT DSN. +// It passes only the DSN (no use_kernel): the driver must auto-detect RT and switch +// to SEA on its own. +func acquireTestDatabricksRT(t *testing.T, dsn string) (drivers.Handle, drivers.OLAPStore) { + cfg := map[string]any{"dsn": dsn} + conn, err := drivers.Open("databricks", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) + require.NoError(t, err) + t.Cleanup(func() { conn.Close() }) + + olap, ok := conn.AsOLAP("default") + require.True(t, ok) + + return conn, olap +} From 3df307d9b941d8a665bccb15ea8078d6d5bda4b3 Mon Sep 17 00:00:00 2001 From: Abhishek Dey Date: Sun, 13 Sep 2026 11:06:56 +0530 Subject: [PATCH 3/4] databricks: trim comments in Lakehouse//RT test Co-authored-by: Isaac --- runtime/drivers/databricks/olap_rt_test.go | 42 ++++------------------ 1 file changed, 7 insertions(+), 35 deletions(-) diff --git a/runtime/drivers/databricks/olap_rt_test.go b/runtime/drivers/databricks/olap_rt_test.go index 475b33cf2521..d5ee7520a102 100644 --- a/runtime/drivers/databricks/olap_rt_test.go +++ b/runtime/drivers/databricks/olap_rt_test.go @@ -14,30 +14,12 @@ import ( "go.uber.org/zap" ) -// TestOLAP_LakehouseRT exercises the full Rill Databricks OLAP path against a live -// Lakehouse//RT (Reyden) SQL warehouse. -// -// Lakehouse//RT rejects the Thrift/HiveServer2 protocol and speaks only SEA (the -// Statement Execution API), which the driver reaches through its SEA-via-kernel -// backend. That backend is compiled in only with `-tags databricks_kernel` -// (CGO_ENABLED=1), so this file carries that build constraint and is excluded from -// the default (Thrift-only) build and from the standard `go test -short ./...` CI job. -// -// The connector config deliberately carries NO use_kernel flag: the driver probes on -// connect and auto-switches Thrift -> SEA when the warehouse rejects Thrift. A query -// that succeeds here therefore proves the RT autodetection path end to end — Thrift -// would have been rejected at connect. -// -// Gated like the other Databricks live tests: it only runs under -// RILL_RUNTIME_TEST_MODE=expensive (skipped under -short) and when -// RILL_RUNTIME_DATABRICKS_RT_TEST_DSN points at a Lakehouse//RT warehouse, e.g. -// -// RILL_RUNTIME_DATABRICKS_RT_TEST_DSN='token:@:443/sql/1.0/warehouses/' -// -// Run with: -// -// CGO_ENABLED=1 RILL_RUNTIME_TEST_MODE=expensive \ -// go test -tags databricks_kernel -run TestOLAP_LakehouseRT ./runtime/drivers/databricks/... +// TestOLAP_LakehouseRT runs the Rill Databricks OLAP path against a live Lakehouse//RT +// warehouse. RT speaks only SEA, reached via the SEA-via-kernel backend, so this file is +// tagged databricks_kernel (out of the default build and `go test -short` CI). Config +// passes only the DSN (no use_kernel): a passing query proves Thrift->SEA autodetect. +// Run: CGO_ENABLED=1 RILL_RUNTIME_TEST_MODE=expensive RILL_RUNTIME_DATABRICKS_RT_TEST_DSN=... +// go test -tags databricks_kernel -run TestOLAP_LakehouseRT ./runtime/drivers/databricks/... func TestOLAP_LakehouseRT(t *testing.T) { t.Skip("skipping Lakehouse//RT live test; needs a databricks_kernel build and an RT warehouse (set RILL_RUNTIME_DATABRICKS_RT_TEST_DSN)") testmode.Expensive(t) @@ -49,10 +31,7 @@ func TestOLAP_LakehouseRT(t *testing.T) { _, olap := acquireTestDatabricksRT(t, dsn) - // Only assert queries Lakehouse//RT actually supports. RT is a SEA reference - // engine and genuinely omits some routines the full DBSQL warehouse ships (e.g. - // current_version() returns UNRESOLVED_ROUTINE / SQLSTATE 42883) — a documented - // RT-vs-DBSQL divergence, not a driver or connectivity failure. + // Only assert queries RT supports (e.g. current_version() is UNRESOLVED_ROUTINE on RT). t.Run("scalar_values", func(t *testing.T) { tests := []struct { query string @@ -77,8 +56,6 @@ func TestOLAP_LakehouseRT(t *testing.T) { } }) - // Session/context routines RT does implement: prove real query execution against - // the warehouse (not just a literal the client could fold locally). t.Run("session_routines", func(t *testing.T) { for _, q := range []string{ "SELECT current_catalog() AS v", @@ -98,7 +75,6 @@ func TestOLAP_LakehouseRT(t *testing.T) { } }) - // Schema inspection over the SEA backend. t.Run("query_schema", func(t *testing.T) { schema, err := olap.QuerySchema(t.Context(), "SELECT 1 AS int_col, 'x' AS str_col", nil) require.NoError(t, err) @@ -107,16 +83,12 @@ func TestOLAP_LakehouseRT(t *testing.T) { require.Equal(t, "str_col", schema.Fields[1].Name) }) - // DryRun issues an EXPLAIN over the SEA backend and must not error. t.Run("dry_run", func(t *testing.T) { _, err := olap.Query(t.Context(), &drivers.Statement{Query: "SELECT 1", DryRun: true}) require.NoError(t, err) }) } -// acquireTestDatabricksRT opens the Databricks driver against a Lakehouse//RT DSN. -// It passes only the DSN (no use_kernel): the driver must auto-detect RT and switch -// to SEA on its own. func acquireTestDatabricksRT(t *testing.T, dsn string) (drivers.Handle, drivers.OLAPStore) { cfg := map[string]any{"dsn": dsn} conn, err := drivers.Open("databricks", "", "default", cfg, storage.MustNew(t.TempDir(), nil), activity.NewNoopClient(), zap.NewNop()) From 6eb57b1e5564d4b9c449f31c1acbe2f2c2d5c695 Mon Sep 17 00:00:00 2001 From: Anshul Khandelwal <12948312+k-anshul@users.noreply.github.com> Date: Wed, 16 Sep 2026 17:29:33 +0530 Subject: [PATCH 4/4] simplifications + ship databricks static bindings --- .github/workflows/cli-release.yml | 2 +- .github/workflows/rill-cloud.yml | 4 +- Makefile | 4 +- .../connectors/data-source/databricks.mdx | 4 +- runtime/drivers/databricks/databricks.go | 110 +++++++----------- .../drivers/databricks/information_schema.go | 5 +- runtime/drivers/databricks/warehouse.go | 50 +++----- 7 files changed, 69 insertions(+), 110 deletions(-) diff --git a/.github/workflows/cli-release.yml b/.github/workflows/cli-release.yml index 9772b73a2b3c..6ff37fcb5b04 100644 --- a/.github/workflows/cli-release.yml +++ b/.github/workflows/cli-release.yml @@ -70,7 +70,7 @@ jobs: VERSION='${{ github.ref_name }}' fi - go build -o rill \ + go build -tags databricks_kernel -o rill \ -mod=readonly \ -ldflags="-s -w -X main.Version=${VERSION} -X main.Commit=${{ github.sha }} -X main.BuildDate=$(date +%FT%TZ)" \ cli/main.go diff --git a/.github/workflows/rill-cloud.yml b/.github/workflows/rill-cloud.yml index 2ec633a6949f..910f156b6a1d 100644 --- a/.github/workflows/rill-cloud.yml +++ b/.github/workflows/rill-cloud.yml @@ -96,9 +96,9 @@ jobs: make runtime.examples.embed if [ ${RELEASE} == "true" ]; then - go build -o rill -mod=readonly -ldflags="-s -w -X main.Version=$(scripts/versiontag.sh) -X main.Commit=${GITHUB_SHA} -X main.BuildDate=$(date +%FT%T)" cli/main.go + go build -tags databricks_kernel -o rill -mod=readonly -ldflags="-s -w -X main.Version=$(scripts/versiontag.sh) -X main.Commit=${GITHUB_SHA} -X main.BuildDate=$(date +%FT%T)" cli/main.go else - go build -o rill -mod=readonly -ldflags="-s -w -X main.Commit=${GITHUB_SHA} -X main.BuildDate=$(date +%FT%T)" cli/main.go + go build -tags databricks_kernel -o rill -mod=readonly -ldflags="-s -w -X main.Commit=${GITHUB_SHA} -X main.BuildDate=$(date +%FT%T)" cli/main.go fi docker buildx build --platform linux/${{ matrix.arch }} -t ${IMAGE}:${GITHUB_SHA}-${{ matrix.arch }} . --push --provenance=false --sbom=false diff --git a/Makefile b/Makefile index 9247013862ee..11fb16ae1dca 100644 --- a/Makefile +++ b/Makefile @@ -4,11 +4,11 @@ all: cli .PHONE: cli-only cli-only: go run scripts/embed_duckdb_ext/main.go - go build -o rill cli/main.go + go build -tags databricks_kernel -o rill cli/main.go .PHONY: cli cli: cli.prepare - go build -o rill cli/main.go + go build -tags databricks_kernel -o rill cli/main.go .PHONY: cli.prepare cli.prepare: runtime.examples.embed diff --git a/docs/docs/developers/build/connectors/data-source/databricks.mdx b/docs/docs/developers/build/connectors/data-source/databricks.mdx index 561157c3bffd..8ae7a5a8a1bc 100644 --- a/docs/docs/developers/build/connectors/data-source/databricks.mdx +++ b/docs/docs/developers/build/connectors/data-source/databricks.mdx @@ -60,9 +60,7 @@ schema: "default" # optional Lakehouse//RT (real-time) warehouses only support the Statement Execution API (SEA) protocol and reject the default Thrift protocol. Rill **auto-detects** this: if a warehouse rejects Thrift, the connector transparently switches to the SEA backend — no -configuration needed. (You can also force it with `use_kernel: true`.) SEA requires a -Rill build that links the `databricks_kernel` backend; otherwise the connection fails -with a clear error. Ordinary DBSQL warehouses keep using Thrift. +configuration needed. (You can also force it with `use_kernel: true`.) Ordinary DBSQL warehouses keep using Thrift. ::: **Step 2: Add credentials to `.env`** diff --git a/runtime/drivers/databricks/databricks.go b/runtime/drivers/databricks/databricks.go index 92f85087c672..0abb62c2854b 100644 --- a/runtime/drivers/databricks/databricks.go +++ b/runtime/drivers/databricks/databricks.go @@ -6,7 +6,6 @@ import ( "fmt" "net/url" "strings" - "sync" dbsqllog "github.com/databricks/databricks-sql-go/logger" "github.com/jmoiron/sqlx" @@ -85,8 +84,7 @@ var spec = drivers.Spec{ Type: drivers.BooleanPropertyType, DisplayName: "Use SEA (Statement Execution API)", Hint: "Force the SEA backend instead of Thrift. Lakehouse//RT warehouses are " + - "auto-detected and switched to SEA automatically, so this is usually unnecessary. " + - "Requires a Rill build with the databricks_kernel backend (CGO); defaults to false.", + "auto-detected and switched to SEA automatically, so this is usually unnecessary.", }, }, ImplementsOLAP: true, @@ -105,7 +103,6 @@ type configProperties struct { LogQueries bool `mapstructure:"log_queries"` // UseKernel forces the driver's SEA backend. Lakehouse//RT is auto-detected and // switched to SEA even when this is false (see getDB), so it's an override. - // Needs a build with the databricks_kernel backend; default false = Thrift. UseKernel bool `mapstructure:"use_kernel"` } @@ -165,12 +162,9 @@ func (c *configProperties) resolveDSN() string { return u.String() } -// withUseKernel forces useKernel=true on a resolved DSN. It is only called when SEA is -// required (explicit use_kernel, or RT auto-detect), so any existing useKernel in the -// DSN — including useKernel=false — is dropped in favor of true. It splits at the first -// '?' (the query separator) and rewrites only the query part; it avoids url.Parse -// because the driver also accepts scheme-less DSNs (token:...@host), which url.Parse -// would misread (treating "token" as the scheme). +// withUseKernel forces useKernel=true on a resolved DSN, dropping any existing useKernel param. +// It rewrites only the query part after the first '?' rather than using url.Parse, +// because the driver also accepts scheme-less DSNs (token:...@host) which url.Parse misreads. func withUseKernel(dsn string) string { base, query, hasQuery := strings.Cut(dsn, "?") if !hasQuery || query == "" { @@ -180,7 +174,7 @@ func withUseKernel(dsn string) string { kept := parts[:0] for _, p := range parts { if strings.HasPrefix(p, "useKernel=") { - continue // drop any existing useKernel (true or false) + continue } kept = append(kept, p) } @@ -236,12 +230,7 @@ type connection struct { db *sqlx.DB // lazily populated using getDB dbErr error dbMu *semaphore.Weighted - - // resolvedDSN is the DSN to connect with, computed by effectiveDSN (auto-detects - // Lakehouse//RT and upgrades to the SEA backend). Empty until a *definitive* - // probe outcome is reached, so a transient probe failure is not cached. - dsnMu sync.Mutex - resolvedDSN string + dsn string // set by backendDSN; carries useKernel=true when Lakehouse//RT was detected } // Ping implements drivers.Handle. @@ -358,64 +347,49 @@ func (c *connection) getDB(ctx context.Context) (*sqlx.DB, error) { return c.db, c.dbErr } - // Only build (and permanently cache) the shared pool once the backend is - // *definitively* known. On a non-definitive result we surface the probe error - // (which may be a real misconfiguration or a transient blip) without caching it, - // so the caller can retry and we re-probe. c.db is thus never swapped or closed - // here, so a pool already handed to concurrent callers is never closed underneath. - dsn, definitive, probeErr := c.effectiveDSN(ctx) - if !definitive { - if probeErr != nil { - return nil, probeErr - } - return nil, errors.New("databricks: could not determine warehouse protocol") + dsn, err := c.backendDSN(ctx) + if err != nil { + return nil, err } c.db, c.dbErr = sqlx.Open("databricks", dsn) - if c.dbErr != nil { - return nil, c.dbErr - } return c.db, c.dbErr } -// effectiveDSN resolves the DSN to connect with, auto-detecting Lakehouse//RT: if SEA -// wasn't explicitly requested (use_kernel) and the warehouse rejects Thrift, it upgrades -// the DSN to the SEA backend. It returns whether the result is definitive (safe to -// cache) and, when not, the probe error that prevented a decision — callers surface that -// error and re-resolve next time rather than lock onto the wrong backend or mask a real -// failure. It's the single shared RT-detection point for the OLAP (getDB) and warehouse -// ingest paths, so RT works without any configuration. -func (c *connection) effectiveDSN(ctx context.Context) (dsn string, definitive bool, probeErr error) { - c.dsnMu.Lock() - defer c.dsnMu.Unlock() - if c.resolvedDSN != "" { - return c.resolvedDSN, true, nil +// backendDSN returns the DSN to connect with, auto-detecting Lakehouse//RT on first call: +// unless SEA was requested explicitly (use_kernel), it pings over Thrift and switches to SEA if the warehouse rejects it. +// The probe pool is closed right away so callers that only need the DSN (QueryAsFiles) leave no connection open. +// Other ping errors are cached in dbErr so callers queued behind dbMu do not each re-probe; +// a cancelled or timed-out probe is not cached since the caller's ctx, not the warehouse, failed. +// Caller must hold dbMu. +func (c *connection) backendDSN(ctx context.Context) (string, error) { + if c.dbErr != nil { + return "", c.dbErr } - - base := c.config.resolveDSN() - if c.config.UseKernel { - c.resolvedDSN = base // already carries useKernel=true - return base, true, nil + if c.dsn != "" { + return c.dsn, nil } - probe, err := sqlx.Open("databricks", base) - if err != nil { - return base, false, err // can't even open; don't cache - } - defer probe.Close() //nolint:errcheck // best-effort close of the probe connection - - // Only cache a *definitive* outcome. A clean ping means Thrift works; the RT - // rejection means switch to SEA. Any other ping error is returned uncached (with - // the error) — it may be transient (a blip, re-probed next call) or persistent (a - // bad token / wrong host), and surfacing it verbatim avoids both locking onto the - // wrong backend and masking a real misconfiguration behind auto-detection. - switch perr := probe.PingContext(ctx); { - case perr == nil: - c.resolvedDSN = base - case rtRequiresSEA(perr): - c.logger.Info("databricks: warehouse rejected Thrift (Lakehouse//RT); switching to the SEA backend") - c.resolvedDSN = withUseKernel(base) - default: - return base, false, perr // transient or persistent — surface it, don't cache + dsn := c.config.resolveDSN() + if !c.config.UseKernel { + probe, err := sqlx.Open("databricks", dsn) + if err != nil { + c.dbErr = err + return "", err + } + defer probe.Close() + err = probe.PingContext(ctx) + switch { + case rtRequiresSEA(err): + c.logger.Info("databricks: warehouse rejected Thrift (Lakehouse//RT); switching to the SEA backend") + dsn = withUseKernel(dsn) + case err != nil: + if ctx.Err() == nil { + c.dbErr = err + } + return "", err + } } - return c.resolvedDSN, true, nil + + c.dsn = dsn + return dsn, nil } diff --git a/runtime/drivers/databricks/information_schema.go b/runtime/drivers/databricks/information_schema.go index 6c7b9fc85209..c6e1fc17b2c9 100644 --- a/runtime/drivers/databricks/information_schema.go +++ b/runtime/drivers/databricks/information_schema.go @@ -155,7 +155,10 @@ func (c *connection) Lookup(ctx context.Context, database, databaseSchema, name fmt.Sprintf("SELECT table_type FROM %sinformation_schema.tables WHERE table_schema = ? AND table_name = ?", prefix), databaseSchema, name, ).Scan(&tableType) - if err != nil && !errors.Is(err, sql.ErrNoRows) { + if errors.Is(err, sql.ErrNoRows) { + return nil, drivers.ErrNotFound + } + if err != nil { return nil, err } diff --git a/runtime/drivers/databricks/warehouse.go b/runtime/drivers/databricks/warehouse.go index 040dcf5c6660..133738967091 100644 --- a/runtime/drivers/databricks/warehouse.go +++ b/runtime/drivers/databricks/warehouse.go @@ -53,16 +53,14 @@ func (c *connection) QueryAsFiles(ctx context.Context, props map[string]any) (ou return nil, err } - // effectiveDSN auto-detects Lakehouse//RT and selects the SEA backend if needed, - // shared with the OLAP path so ingest works without configuration. On a - // non-definitive result, defer like getDB (surface the probe error / retry) rather - // than ingest over a possibly-wrong backend. - dsn, definitive, probeErr := c.effectiveDSN(ctx) - if !definitive { - if probeErr != nil { - return nil, probeErr - } - return nil, errors.New("databricks: could not determine warehouse protocol") + err = c.dbMu.Acquire(ctx, 1) + if err != nil { + return nil, err + } + dsn, err := c.backendDSN(ctx) + c.dbMu.Release(1) + if err != nil { + return nil, err } db, err := sql.Open("databricks", dsn) if err != nil { @@ -102,14 +100,12 @@ func (c *connection) QueryAsFiles(ctx context.Context, props map[string]any) (ou // implement them and returns ErrNotSupportedByKernel, so fall back to // GetArrowBatches via kernelIPCStreams; the parquet path is identical for both. dr := rows.(dbsqlrows.Rows) - selfDescribing := false // whether each stream carries its own schema (kernel adapter) ipcStreams, err := dr.GetArrowIPCStreams(ctx) if errors.Is(err, dbsqlerr.ErrNotSupportedByKernel) { var batches dbsqlrows.ArrowBatchIterator batches, err = dr.GetArrowBatches(ctx) if err == nil { ipcStreams = &kernelIPCStreams{batches: batches} - selfDescribing = true } } if err != nil { @@ -131,13 +127,12 @@ func (c *connection) QueryAsFiles(ctx context.Context, props map[string]any) (ou } return &fileIterator{ - db: db, - conn: conn, - rows: rows, - ipcStreams: ipcStreams, - selfDescribing: selfDescribing, - logger: c.logger, - tempDir: tempDir, + db: db, + conn: conn, + rows: rows, + ipcStreams: ipcStreams, + logger: c.logger, + tempDir: tempDir, }, nil } @@ -146,12 +141,8 @@ type fileIterator struct { conn *sql.Conn rows sqld.Rows ipcStreams dbsqlrows.ArrowIPCStreamIterator - // selfDescribing is true when each stream carries its own schema message (the - // kernel adapter). For native Thrift streams it's false, and subsequent streams - // are read with ipc.WithSchema to validate cross-stream schema consistency. - selfDescribing bool - logger *zap.Logger - tempDir string + logger *zap.Logger + tempDir string totalRecords int64 downloaded bool @@ -317,14 +308,7 @@ func (f *fileIterator) Next(ctx context.Context) ([]string, error) { return nil, err } - // Native (Thrift) subsequent streams are validated against the first stream's - // schema; the kernel adapter's streams are self-describing, so skip WithSchema - // there (each carries its own schema message). - var opts []ipc.Option - if !f.selfDescribing { - opts = append(opts, ipc.WithSchema(schema)) - } - rdr, err := ipc.NewReader(stream, opts...) + rdr, err := ipc.NewReader(stream, ipc.WithSchema(schema)) if err != nil { return nil, err }