From 532f9b896f53c8f57202b376a2043e1a374039ff Mon Sep 17 00:00:00 2001 From: Parag Jain Date: Thu, 30 Jul 2026 11:06:17 +0530 Subject: [PATCH 1/2] support passing query attributes to druid context --- .../druidsqldriver/druid_api_sql_driver.go | 86 ++++++++----- .../druid_api_sql_driver_test.go | 113 ++++++++++++++++++ runtime/drivers/druid/olap.go | 25 +--- 3 files changed, 173 insertions(+), 51 deletions(-) create mode 100644 runtime/drivers/druid/druidsqldriver/druid_api_sql_driver_test.go diff --git a/runtime/drivers/druid/druidsqldriver/druid_api_sql_driver.go b/runtime/drivers/druid/druidsqldriver/druid_api_sql_driver.go index 53150dc26c6f..3c6fd5ff199e 100644 --- a/runtime/drivers/druid/druidsqldriver/druid_api_sql_driver.go +++ b/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,29 +452,22 @@ 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{ @@ -482,27 +475,52 @@ func newDruidRequest(query string, args []driver.NamedValue, queryCfg *QueryConf 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{} diff --git a/runtime/drivers/druid/druidsqldriver/druid_api_sql_driver_test.go b/runtime/drivers/druid/druidsqldriver/druid_api_sql_driver_test.go new file mode 100644 index 000000000000..5b34a02a7dc4 --- /dev/null +++ b/runtime/drivers/druid/druidsqldriver/druid_api_sql_driver_test.go @@ -0,0 +1,113 @@ +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", + // Empty values (e.g. from an unresolved template) must be omitted. + "empty": "", + }, + }) + + require.Equal(t, "user@example.com", qctx["userEmail"]) + require.NotEqual(t, "hijack", qctx["sqlQueryId"]) + require.Equal(t, float64(3), qctx["priority"]) + require.NotContains(t, qctx, "empty") +} diff --git a/runtime/drivers/druid/olap.go b/runtime/drivers/druid/olap.go index 301d8aeade97..26c5a9511284 100644 --- a/runtime/drivers/druid/olap.go +++ b/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 From 4f11738f59f14675e3c445f8a4b650b8f5a8da6a Mon Sep 17 00:00:00 2001 From: Parag Jain Date: Thu, 30 Jul 2026 16:19:04 +0530 Subject: [PATCH 2/2] fix test --- .../drivers/druid/druidsqldriver/druid_api_sql_driver_test.go | 3 --- 1 file changed, 3 deletions(-) diff --git a/runtime/drivers/druid/druidsqldriver/druid_api_sql_driver_test.go b/runtime/drivers/druid/druidsqldriver/druid_api_sql_driver_test.go index 5b34a02a7dc4..8e6752e07ea0 100644 --- a/runtime/drivers/druid/druidsqldriver/druid_api_sql_driver_test.go +++ b/runtime/drivers/druid/druidsqldriver/druid_api_sql_driver_test.go @@ -101,13 +101,10 @@ func TestQueryContextAttributes(t *testing.T) { // Keys set by the driver itself cannot be overridden. "sqlQueryId": "hijack", "priority": "999", - // Empty values (e.g. from an unresolved template) must be omitted. - "empty": "", }, }) require.Equal(t, "user@example.com", qctx["userEmail"]) require.NotEqual(t, "hijack", qctx["sqlQueryId"]) require.Equal(t, float64(3), qctx["priority"]) - require.NotContains(t, qctx, "empty") }