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 0d586b0c2290..8ae7a5a8a1bc 100644 --- a/docs/docs/developers/build/connectors/data-source/databricks.mdx +++ b/docs/docs/developers/build/connectors/data-source/databricks.mdx @@ -56,6 +56,13 @@ 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`.) Ordinary DBSQL warehouses keep using Thrift. +::: + **Step 2: Add credentials to `.env`** ```bash diff --git a/go.mod b/go.mod index 8c8cb04e8704..aaef57712df7 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.14.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.10505.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.57.0 + golang.org/x/net v0.58.0 golang.org/x/oauth2 v0.36.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.40.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.10505.0 // indirect @@ -419,11 +426,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.54.0 // indirect - golang.org/x/mod v0.38.0 // indirect - golang.org/x/telemetry v0.0.0-20260708182218-49f421fb7959 // 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.48.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 78bcb398506d..e8fa9339a88f 100644 --- a/go.sum +++ b/go.sum @@ -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.14.0 h1:pc6/WPEwKujCBlHbz/U2kVWou8BCvo2Hm3GJVsmyogU= -github.com/databricks/databricks-sql-go v1.14.0/go.mod h1:knb3jMMEf7rYm+Hw0i6NFuYC12WX7jkSQNhuz/D5V1Q= +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.54.0 h1:YLIA59K4fiNzHzjnZt2tUJQjQtUWfWbeHBqKtk3eScw= -golang.org/x/crypto v0.54.0/go.mod h1:KWL8ny2AZdGR2cWmzeHrp2azQPGogOv+HeQaVEXC2dk= +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.38.0 h1:MECBjubtXD7yj4HrhIUcywNaGeNVUdfVnxmPajOk4yk= -golang.org/x/mod v0.38.0/go.mod h1:V6Xz0pq8TQ3dGqVQ1FVHuelZpAL0uNhSkk9ogYP3c40= +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.57.0 h1:K5+3DljvIuDG9/Jv9rvyMywYNFCQ9RSUY6OOTTkT+tE= -golang.org/x/net v0.57.0/go.mod h1:KpXc8iv+r3XplLAG/f7Jsf9RPszJzdR0f58q9vGOuEU= +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= @@ -2969,8 +2983,8 @@ 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.47.0 h1:o7XGOvZQCADBQQ4Y7VNq2dRWQR7JmOUW8Kxx4ZsNgWs= golang.org/x/sys v0.47.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= -golang.org/x/telemetry v0.0.0-20260708182218-49f421fb7959 h1:RJhm5l6Fo4rmEIcndxDllNhhf/fAx8qIm4t6A7vpm2A= -golang.org/x/telemetry v0.0.0-20260708182218-49f421fb7959/go.mod h1:LV7u5Oco+Z/g6XI7PqN+EUUUGGkEcmB1uj2ceI0fOVg= +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= @@ -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.40.0 h1:Ub2Z6/xjgF1WrYQz2nuITOEegKFtiIy+rieRJ5lHZKs= -golang.org/x/text v0.40.0/go.mod h1:hpnzDAfGV753zIKo+wk3u1bVKCGPbrnF7+7LBF/UHVY= +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.48.0 h1:3+hClM1aLL5mjMKm5ovokw9epgRXPuu2tILgismM6RE= -golang.org/x/tools v0.48.0/go.mod h1:08xX0orndb/F7jJxGDicx061tyd5pcMto75YMAXr6lk= +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..0abb62c2854b 100644 --- a/runtime/drivers/databricks/databricks.go +++ b/runtime/drivers/databricks/databricks.go @@ -79,6 +79,13 @@ 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.", + }, }, ImplementsOLAP: true, ImplementsWarehouse: true, @@ -94,6 +101,9 @@ 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. + UseKernel bool `mapstructure:"use_kernel"` } func (c *configProperties) validate() error { @@ -124,10 +134,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 +162,32 @@ func (c *configProperties) resolveDSN() string { return u.String() } +// 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 == "" { + return base + "?useKernel=true" + } + parts := strings.Split(query, "&") + kept := parts[:0] + for _, p := range parts { + if strings.HasPrefix(p, "useKernel=") { + continue + } + 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 +230,7 @@ type connection struct { db *sqlx.DB // lazily populated using getDB dbErr error dbMu *semaphore.Weighted + dsn string // set by backendDSN; carries useKernel=true when Lakehouse//RT was detected } // Ping implements drivers.Handle. @@ -303,9 +347,49 @@ 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()) - if c.dbErr != nil { - return nil, c.dbErr + dsn, err := c.backendDSN(ctx) + if err != nil { + return nil, err } + c.db, c.dbErr = sqlx.Open("databricks", dsn) return c.db, c.dbErr } + +// 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 + } + if c.dsn != "" { + return c.dsn, nil + } + + 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 + } + } + + c.dsn = dsn + return dsn, 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..c6e1fc17b2c9 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,42 @@ 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 errors.Is(err, sql.ErrNoRows) { + return nil, drivers.ErrNotFound + } + if err != nil { + 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 +190,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/olap_rt_test.go b/runtime/drivers/databricks/olap_rt_test.go new file mode 100644 index 000000000000..d5ee7520a102 --- /dev/null +++ b/runtime/drivers/databricks/olap_rt_test.go @@ -0,0 +1,102 @@ +//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 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) + + 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 RT supports (e.g. current_version() is UNRESOLVED_ROUTINE on RT). + 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()) + }) + } + }) + + 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()) + }) + } + }) + + 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) + }) + + t.Run("dry_run", func(t *testing.T) { + _, err := olap.Query(t.Context(), &drivers.Statement{Query: "SELECT 1", DryRun: true}) + require.NoError(t, err) + }) +} + +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 +} diff --git a/runtime/drivers/databricks/warehouse.go b/runtime/drivers/databricks/warehouse.go index e9493ba632fb..133738967091 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,16 @@ func (c *connection) QueryAsFiles(ctx context.Context, props map[string]any) (ou return nil, err } - db, err := sql.Open("databricks", c.config.resolveDSN()) + 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 { return nil, err } @@ -82,7 +96,18 @@ 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) + 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} + } + } if err != nil { return nil, err } @@ -125,6 +150,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 {