Support query attributes in the Databricks OLAP driver

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

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=

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.

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

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())
+	})
+}