support passing query attributes to druid context

runtime/drivers/druid/druidsqldriver/druid_api_sql_driver.go

@@ -105,7 +105,7 @@ func (c *sqlConnection) QueryContext(ctx context.Context, query string, args []d
 	return re.RunCtx(ctx, func(ctx context.Context) (driver.Rows, retrier.Action, error) {
 		queryCfg := queryConfigFromContext(ctx)
 
-		dr := newDruidRequest(query, args, queryCfg)
+		dr, queryID := newDruidRequest(query, args, queryCfg)
 		b, err := json.Marshal(dr)
 		if err != nil {
 			return nil, retrier.Fail, err
@@ -116,7 +116,7 @@ func (c *sqlConnection) QueryContext(ctx context.Context, query string, args []d
 		context.AfterFunc(ctx, func() {
 			tctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 			defer cancel()
-			r, err := http.NewRequestWithContext(tctx, http.MethodDelete, urlutil.MustJoinURL(c.dsn, dr.Context.SQLQueryID), http.NoBody)
+			r, err := http.NewRequestWithContext(tctx, http.MethodDelete, urlutil.MustJoinURL(c.dsn, queryID), http.NoBody)
 			if err != nil {
 				return
 			}
@@ -398,7 +398,7 @@ var _ retrier.AdditionalTest = &coordinatorHTTPCheck{}
 // b) if the coordinator has a transient error -> not a hard-failure - the table 'A' can exist
 // c) if the coordinator returns not a transient error (ie access-denied) -> hard-failure - we shouldn't wait until the configuration is changed by someone
 func (chc *coordinatorHTTPCheck) IsHardFailure(ctx context.Context) (bool, error) {
-	dr := newDruidRequest("SELECT * FROM sys.segments LIMIT 1", nil, nil)
+	dr, _ := newDruidRequest("SELECT * FROM sys.segments LIMIT 1", nil, nil)
 	b, err := json.Marshal(dr)
 	if err != nil {
 		return false, err
@@ -452,57 +452,75 @@ func (chc *coordinatorHTTPCheck) IsHardFailure(ctx context.Context) (bool, error
 	}
 }
 
-type DruidQueryContext struct {
-	SQLQueryID                 string `json:"sqlQueryId"`
-	EnableTimeBoundaryPlanning bool   `json:"enableTimeBoundaryPlanning"`
-	UseCache                   *bool  `json:"useCache,omitempty"`
-	PopulateCache              *bool  `json:"populateCache,omitempty"`
-	Priority                   int    `json:"priority,omitempty"`
-}
-
 type DruidParameter struct {
 	Type  string `json:"type"`
 	Value any    `json:"value"`
 }
 
 type DruidRequest struct {
-	Query          string            `json:"query"`
-	Header         bool              `json:"header"`
-	SQLTypesHeader bool              `json:"sqlTypesHeader"`
-	ResultFormat   string            `json:"resultFormat"`
-	Parameters     []DruidParameter  `json:"parameters"`
-	Context        DruidQueryContext `json:"context"`
+	Query          string           `json:"query"`
+	Header         bool             `json:"header"`
+	SQLTypesHeader bool             `json:"sqlTypesHeader"`
+	ResultFormat   string           `json:"resultFormat"`
+	Parameters     []DruidParameter `json:"parameters"`
+	Context        map[string]any   `json:"context"`
 }
 
-func newDruidRequest(query string, args []driver.NamedValue, queryCfg *QueryConfig) *DruidRequest {
+// newDruidRequest builds a Druid SQL API request and returns it along with the generated query ID.
+func newDruidRequest(query string, args []driver.NamedValue, queryCfg *QueryConfig) (*DruidRequest, string) {
 	parameters := make([]DruidParameter, len(args))
 	for i, arg := range args {
 		parameters[i] = DruidParameter{
 			Type:  toType(arg.Value),
 			Value: arg.Value,
 		}
 	}
-	var useCache, populateCache *bool
-	priority := 0
-	if queryCfg != nil {
-		useCache = queryCfg.UseCache
-		populateCache = queryCfg.PopulateCache
-		priority = queryCfg.Priority
-	}
+	queryID := uuid.New().String()
 	return &DruidRequest{
 		Query:          query,
 		Header:         true,
 		SQLTypesHeader: true,
 		ResultFormat:   "arrayLines",
 		Parameters:     parameters,
-		Context: DruidQueryContext{
-			SQLQueryID:                 uuid.New().String(),
-			EnableTimeBoundaryPlanning: true,
-			UseCache:                   useCache,
-			PopulateCache:              populateCache,
-			Priority:                   priority,
-		},
+		Context:        newQueryContext(queryID, queryCfg),
+	}, queryID
+}
+
+// reservedContextKeys are the query context keys set by the driver itself;
+// they cannot be overridden by the pass-through attributes in QueryConfig.Attributes.
+var reservedContextKeys = map[string]bool{
+	"sqlQueryId":                 true,
+	"enableTimeBoundaryPlanning": true,
+	"useCache":                   true,
+	"populateCache":              true,
+	"priority":                   true,
+}
+
+// newQueryContext builds the Druid query context: https://druid.apache.org/docs/latest/querying/query-context/
+func newQueryContext(queryID string, queryCfg *QueryConfig) map[string]any {
+	qctx := map[string]any{
+		"sqlQueryId":                 queryID,
+		"enableTimeBoundaryPlanning": true,
+	}
+	if queryCfg == nil {
+		return qctx
+	}
+	if queryCfg.UseCache != nil {
+		qctx["useCache"] = *queryCfg.UseCache
+	}
+	if queryCfg.PopulateCache != nil {
+		qctx["populateCache"] = *queryCfg.PopulateCache
+	}
+	if queryCfg.Priority != 0 {
+		qctx["priority"] = queryCfg.Priority
+	}
+	for k, v := range queryCfg.Attributes {
+		if k == "" || reservedContextKeys[k] {
+			continue
+		}
+		qctx[k] = v
 	}
+	return qctx
 }
 
 func (s *stmt) Exec(args []driver.Value) (driver.Result, error) {
@@ -517,6 +535,10 @@ type QueryConfig struct {
 	UseCache      *bool
 	PopulateCache *bool
 	Priority      int
+	// Attributes are passed through to the Druid query context as-is,
+	// e.g. for attributing queries to callers in Druid's query logs.
+	// Empty keys and keys that collide with the driver's own context keys are skipped.
+	Attributes map[string]string
 }
 
 type queryCfgCtxKey struct{}

runtime/drivers/druid/druidsqldriver/druid_api_sql_driver_test.go

@@ -0,0 +1,110 @@
+package druidsqldriver
+
+import (
+	"context"
+	"database/sql"
+	"encoding/json"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"sync"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+// testServer is a Druid SQL API stub that captures each request body,
+// and responds with a minimal valid arrayLines result (header, types header, one row).
+type testServer struct {
+	*httptest.Server
+	mu       sync.Mutex
+	requests []DruidRequest
+}
+
+func newTestServer(t *testing.T) *testServer {
+	t.Helper()
+	ts := &testServer{}
+	ts.Server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		body, err := io.ReadAll(r.Body)
+		if err != nil {
+			http.Error(w, err.Error(), http.StatusInternalServerError)
+			return
+		}
+
+		var dr DruidRequest
+		if err := json.Unmarshal(body, &dr); err != nil {
+			http.Error(w, err.Error(), http.StatusBadRequest)
+			return
+		}
+		ts.mu.Lock()
+		ts.requests = append(ts.requests, dr)
+		ts.mu.Unlock()
+
+		w.WriteHeader(http.StatusOK)
+		_, _ = w.Write([]byte("[\"n\"]\n[\"BIGINT\"]\n[1]\n"))
+	}))
+	t.Cleanup(ts.Close)
+	return ts
+}
+
+// query runs a query against the stub server with the given config and returns the captured request's context.
+func (ts *testServer) query(t *testing.T, queryCfg *QueryConfig) map[string]any {
+	t.Helper()
+	db, err := sql.Open("druid", ts.URL)
+	require.NoError(t, err)
+	defer db.Close()
+
+	ctx := context.Background()
+	if queryCfg != nil {
+		ctx = WithQueryConfig(ctx, queryCfg)
+	}
+
+	rows, err := db.QueryContext(ctx, "SELECT 1")
+	require.NoError(t, err)
+	require.NoError(t, rows.Close())
+
+	ts.mu.Lock()
+	defer ts.mu.Unlock()
+	require.NotEmpty(t, ts.requests)
+	return ts.requests[len(ts.requests)-1].Context
+}
+
+func TestQueryContextDefaults(t *testing.T) {
+	qctx := newTestServer(t).query(t, nil)
+
+	require.NotEmpty(t, qctx["sqlQueryId"])
+	require.Equal(t, true, qctx["enableTimeBoundaryPlanning"])
+	// Optional keys must be omitted when unset.
+	require.Len(t, qctx, 2)
+}
+
+func TestQueryContextConfig(t *testing.T) {
+	useCache := false
+	populateCache := true
+	qctx := newTestServer(t).query(t, &QueryConfig{
+		UseCache:      &useCache,
+		PopulateCache: &populateCache,
+		Priority:      3,
+	})
+
+	// A pointer to false must serialize as false, not be omitted.
+	require.Equal(t, false, qctx["useCache"])
+	require.Equal(t, true, qctx["populateCache"])
+	require.Equal(t, float64(3), qctx["priority"])
+}
+
+func TestQueryContextAttributes(t *testing.T) {
+	qctx := newTestServer(t).query(t, &QueryConfig{
+		Priority: 3,
+		Attributes: map[string]string{
+			"userEmail": "user@example.com",
+			// Keys set by the driver itself cannot be overridden.
+			"sqlQueryId": "hijack",
+			"priority":   "999",
+		},
+	})
+
+	require.Equal(t, "user@example.com", qctx["userEmail"])
+	require.NotEqual(t, "hijack", qctx["sqlQueryId"])
+	require.Equal(t, float64(3), qctx["priority"])
+}

runtime/drivers/druid/olap.go

@@ -97,28 +97,15 @@ func (c *connection) Query(ctx context.Context, stmt *drivers.Statement) (res *d
 		ctx, cancelFunc = context.WithTimeout(ctx, stmt.ExecutionTimeout)
 	}
 
-	var queryCfg *druidsqldriver.QueryConfig
-	if stmt.UseCache != nil {
-		queryCfg = &druidsqldriver.QueryConfig{
-			UseCache: stmt.UseCache,
-		}
-	}
-	if stmt.PopulateCache != nil {
-		if queryCfg == nil {
-			queryCfg = &druidsqldriver.QueryConfig{}
-		}
-		queryCfg.PopulateCache = stmt.PopulateCache
+	queryCfg := &druidsqldriver.QueryConfig{
+		UseCache:      stmt.UseCache,
+		PopulateCache: stmt.PopulateCache,
+		Attributes:    stmt.QueryAttributes, // Attributes are passed through to the Druid query context
 	}
-	if !c.config.SkipQueryPriority && stmt.Priority != 0 {
-		if queryCfg == nil {
-			queryCfg = &druidsqldriver.QueryConfig{}
-		}
+	if !c.config.SkipQueryPriority {
 		queryCfg.Priority = stmt.Priority
 	}
-
-	if queryCfg != nil {
-		ctx = druidsqldriver.WithQueryConfig(ctx, queryCfg)
-	}
+	ctx = druidsqldriver.WithQueryConfig(ctx, queryCfg)
 
 	var rows *sqlx.Rows