From af9099086caf0668048edc7a54fcf8d0c86aa406 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Miguel=20P=C3=A9rez=20Pasalodos?= Date: Wed, 2 Sep 2026 12:39:28 +0200 Subject: [PATCH] Support query attributes in the Databricks OLAP driver Databricks was the only OLAP driver that ignored a metrics view's query_attributes. #8266 wired them into every driver that existed at the time, but the Databricks OLAP driver landed later in #9255 and was never hooked up, so the attributes were silently dropped. Send them as Databricks query tags, which are recorded in the query_tags column of system.query.history. The driver attaches tags per statement (as a ConfOverlay on the ExecuteStatement request) rather than per session, so they remain correct even though connections are pooled and shared between users. This also avoids depending on the query text: statement_text can be redacted workspace-wide, in which case a SQL comment would carry nothing. Query tags are Public Preview and not available on every workspace, and Databricks fails a statement outright when it doesn't recognise a configuration rather than ignoring it: [CONFIG_NOT_AVAILABLE.WITHOUT_SUGGESTION] Configuration query_tags is not available. Sending tags unconditionally would therefore break every query of any metrics view that sets query_attributes on such a workspace. So the first rejection is detected, the query retried without tags, and the result latched on the connection so later queries skip tagging altogether. Bumps databricks-sql-go to v1.14.0. Per-statement query tags (driverctx.NewContextWithQueryTags) first shipped in v1.11.0, and v1.14.0 is the newest release before v1.15.0 started requiring the databricks-sql-kernel-bindings modules, which add ~585MB of prebuilt platform binaries to the module graph without being reachable from any package we build. Verified against a live SQL warehouse: a Query() with QueryAttributes set is recorded in query history as [{"key":"rillUserEmail","value":"..."},{"key":"rillProbe","value":"..."}] and, with the driver patched to send a configuration the workspace rejects, queries still succeed via the untagged retry. --- go.mod | 4 +- go.sum | 8 +- runtime/drivers/databricks/databricks.go | 5 + runtime/drivers/databricks/olap.go | 75 ++++++++++++++- runtime/drivers/databricks/query_tags_test.go | 91 +++++++++++++++++++ 5 files changed, 174 insertions(+), 9 deletions(-) create mode 100644 runtime/drivers/databricks/query_tags_test.go diff --git a/go.mod b/go.mod index dbd145526c6a..fa9a10d9c721 100644 --- a/go.mod +++ b/go.mod @@ -37,7 +37,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.14.0 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 @@ -214,7 +214,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 diff --git a/go.sum b/go.sum index b7441e182b0b..56a549bfd3d0 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,8 @@ 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.14.0 h1:pc6/WPEwKujCBlHbz/U2kVWou8BCvo2Hm3GJVsmyogU= +github.com/databricks/databricks-sql-go v1.14.0/go.mod h1:knb3jMMEf7rYm+Hw0i6NFuYC12WX7jkSQNhuz/D5V1Q= 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= diff --git a/runtime/drivers/databricks/databricks.go b/runtime/drivers/databricks/databricks.go index 72ed9fc7f444..73f89fa65928 100644 --- a/runtime/drivers/databricks/databricks.go +++ b/runtime/drivers/databricks/databricks.go @@ -6,6 +6,7 @@ import ( "fmt" "net/url" "strings" + "sync/atomic" dbsqllog "github.com/databricks/databricks-sql-go/logger" "github.com/jmoiron/sqlx" @@ -187,6 +188,10 @@ type connection struct { db *sqlx.DB // lazily populated using getDB dbErr error dbMu *semaphore.Weighted + + // Set once the workspace has rejected the query_tags configuration, after which + // query attributes are no longer sent. See connection.queryTagsRejected. + queryTagsUnsupported atomic.Bool } // Ping implements drivers.Handle. diff --git a/runtime/drivers/databricks/olap.go b/runtime/drivers/databricks/olap.go index bb9904eae107..3a1f6942a07b 100644 --- a/runtime/drivers/databricks/olap.go +++ b/runtime/drivers/databricks/olap.go @@ -5,6 +5,7 @@ import ( "fmt" "strings" + "github.com/databricks/databricks-sql-go/driverctx" "github.com/jmoiron/sqlx" runtimev1 "github.com/rilldata/rill/proto/gen/rill/runtime/v1" "github.com/rilldata/rill/runtime/drivers" @@ -14,6 +15,10 @@ import ( var _ drivers.OLAPStore = (*connection)(nil) +// queryTagsConf is the statement configuration the Databricks driver uses to carry +// query tags. It appears in the error returned by workspaces that don't support them. +const queryTagsConf = "query_tags" + // Dialect implements drivers.OLAPStore. func (c *connection) Dialect() drivers.Dialect { return DialectDatabricks @@ -49,19 +54,42 @@ func (c *connection) MayBeScaledToZero(ctx context.Context) bool { // Query implements drivers.OLAPStore. func (c *connection) Query(ctx context.Context, stmt *drivers.Statement) (*drivers.Result, error) { if c.config.LogQueries { - c.logger.Info("databricks query", zap.String("sql", c.Dialect().SanitizeQueryForLogging(stmt.Query)), zap.Any("args", stmt.Args), observability.ZapCtx(ctx)) + fields := []zap.Field{ + zap.String("sql", c.Dialect().SanitizeQueryForLogging(stmt.Query)), + zap.Any("args", stmt.Args), + observability.ZapCtx(ctx), + } + if len(stmt.QueryAttributes) > 0 { + fields = append(fields, zap.Any("query_attributes", stmt.QueryAttributes)) + } + c.logger.Info("databricks query", fields...) } + + // Send the query attributes as Databricks query tags, which are recorded in the + // query_tags column of system.query.history. The driver attaches them per statement + // (as a ConfOverlay on the ExecuteStatement request) rather than per session, so they + // stay correct even though connections are pooled and shared across users. + taggedCtx := c.contextWithQueryTags(ctx, stmt.QueryAttributes) + db, err := c.getDB(ctx) if err != nil { return nil, err } if stmt.DryRun { - _, err = db.ExecContext(ctx, fmt.Sprintf("EXPLAIN %s", stmt.Query), stmt.Args...) + _, err = db.ExecContext(taggedCtx, fmt.Sprintf("EXPLAIN %s", stmt.Query), stmt.Args...) + if c.queryTagsRejected(err) { + _, err = db.ExecContext(ctx, fmt.Sprintf("EXPLAIN %s", stmt.Query), stmt.Args...) + } return nil, err } - rows, err := db.QueryxContext(ctx, stmt.Query, stmt.Args...) + rows, err := db.QueryxContext(taggedCtx, stmt.Query, stmt.Args...) + if c.queryTagsRejected(err) { + // The workspace doesn't support query tags. Retry untagged so a metrics view + // that sets query_attributes still works there. + rows, err = db.QueryxContext(ctx, stmt.Query, stmt.Args...) + } if err != nil { return nil, err } @@ -122,6 +150,47 @@ func (c *connection) Head(ctx context.Context, db, schema, table string, limit i return c.Query(ctx, &drivers.Statement{Query: q}) } +// contextWithQueryTags returns a context carrying attrs as Databricks query tags. +// The context is returned unchanged when there are no attributes, or when this +// connection has already seen the workspace reject query tags. +func (c *connection) contextWithQueryTags(ctx context.Context, attrs map[string]string) context.Context { + if len(attrs) == 0 || c.queryTagsUnsupported.Load() { + return ctx + } + // Copy so a later mutation of the statement's map can't race with the driver + // reading it while the query is in flight. + tags := make(map[string]string, len(attrs)) + for k, v := range attrs { + tags[k] = v + } + return driverctx.NewContextWithQueryTags(ctx, tags) +} + +// queryTagsRejected reports whether err is the workspace refusing the query_tags +// configuration, meaning the query should be retried without tags. +// +// Query tags are sent as a statement configuration, and Databricks fails the whole +// statement when it doesn't recognise a configuration rather than ignoring it: +// +// [CONFIG_NOT_AVAILABLE.WITHOUT_SUGGESTION] Configuration query_tags is not available. +// +// Without this, enabling query_attributes on a metrics view would break every query +// against a workspace that doesn't have the (Public Preview) query tags feature. The +// result is latched on the connection so only the first query pays for the retry. +func (c *connection) queryTagsRejected(err error) bool { + if err == nil { + return false + } + msg := err.Error() + if !strings.Contains(msg, "CONFIG_NOT_AVAILABLE") || !strings.Contains(msg, queryTagsConf) { + return false + } + if c.queryTagsUnsupported.CompareAndSwap(false, true) { + c.logger.Warn("databricks: workspace rejected query tags, dropping query_attributes for this connection", zap.Error(err)) + } + return true +} + func rowsToSchema(r *sqlx.Rows) (*runtimev1.StructType, error) { if r == nil { return nil, nil diff --git a/runtime/drivers/databricks/query_tags_test.go b/runtime/drivers/databricks/query_tags_test.go new file mode 100644 index 000000000000..4cd544338394 --- /dev/null +++ b/runtime/drivers/databricks/query_tags_test.go @@ -0,0 +1,91 @@ +package databricks + +import ( + "context" + "errors" + "testing" + + "github.com/databricks/databricks-sql-go/driverctx" + "github.com/stretchr/testify/require" + "go.uber.org/zap" +) + +func newTestConn() *connection { + return &connection{logger: zap.NewNop()} +} + +func TestContextWithQueryTags(t *testing.T) { + t.Run("NoAttributes", func(t *testing.T) { + // Queries without attributes must be left exactly as they were. + c := newTestConn() + ctx := context.Background() + require.Equal(t, ctx, c.contextWithQueryTags(ctx, nil)) + require.Equal(t, ctx, c.contextWithQueryTags(ctx, map[string]string{})) + require.Nil(t, driverctx.QueryTagsFromContext(c.contextWithQueryTags(ctx, nil))) + }) + + t.Run("Attributes", func(t *testing.T) { + c := newTestConn() + ctx := c.contextWithQueryTags(context.Background(), map[string]string{ + "rillUserEmail": "user@example.com", + "rillProject": "my-project", + }) + + require.Equal(t, map[string]string{ + "rillUserEmail": "user@example.com", + "rillProject": "my-project", + }, driverctx.QueryTagsFromContext(ctx)) + }) + + t.Run("CopiesAttributes", func(t *testing.T) { + // The driver reads the map while the query is in flight, so the tags must not + // alias the caller's map. + c := newTestConn() + attrs := map[string]string{"rillUserEmail": "user@example.com"} + ctx := c.contextWithQueryTags(context.Background(), attrs) + + attrs["rillUserEmail"] = "someone.else@example.com" + delete(attrs, "rillUserEmail") + + require.Equal(t, map[string]string{"rillUserEmail": "user@example.com"}, driverctx.QueryTagsFromContext(ctx)) + }) + + t.Run("SkipsAfterRejection", func(t *testing.T) { + // Once the workspace has rejected query tags, don't pay for a retry on every + // subsequent query. + c := newTestConn() + c.queryTagsUnsupported.Store(true) + + ctx := context.Background() + require.Equal(t, ctx, c.contextWithQueryTags(ctx, map[string]string{"rillUserEmail": "user@example.com"})) + }) +} + +func TestQueryTagsRejected(t *testing.T) { + // The error Databricks returns when the workspace has no query tags support. + unsupported := errors.New("databricks: execution error: failed to execute query: " + + "[CONFIG_NOT_AVAILABLE.WITHOUT_SUGGESTION] Configuration query_tags is not available. SQLSTATE: 42K0I") + + t.Run("Unsupported", func(t *testing.T) { + c := newTestConn() + require.True(t, c.queryTagsRejected(unsupported)) + require.True(t, c.queryTagsUnsupported.Load()) + }) + + t.Run("NilError", func(t *testing.T) { + c := newTestConn() + require.False(t, c.queryTagsRejected(nil)) + require.False(t, c.queryTagsUnsupported.Load()) + }) + + t.Run("UnrelatedErrorsArePassedThrough", func(t *testing.T) { + // A genuine query error must not be retried or latched, otherwise real + // failures would be masked and tags dropped for the rest of the connection. + c := newTestConn() + require.False(t, c.queryTagsRejected(errors.New("[TABLE_OR_VIEW_NOT_FOUND] table not found"))) + // A different unavailable configuration is also not ours to handle. + require.False(t, c.queryTagsRejected(errors.New( + "[CONFIG_NOT_AVAILABLE.WITHOUT_SUGGESTION] Configuration some_other_conf is not available."))) + require.False(t, c.queryTagsUnsupported.Load()) + }) +}