From b91a148f56cf6a386a492582cbca9196fcbb1100 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Miguel=20P=C3=A9rez=20Pasalodos?= Date: Tue, 28 Jul 2026 10:04:01 +0200 Subject: [PATCH] Stamp caller attribution on the Druid query context Forward the calling user's email and service name from the metrics view's query_attributes to Druid's query context (as rillUserEmail/rillServiceToken), so they appear in Druid's query logs for attribution, similar to Imply Pivot's implyUserEmail. Co-Authored-By: Claude Fable 5 --- .../druidsqldriver/druid_api_sql_driver.go | 12 +++ .../druid_api_sql_driver_test.go | 96 +++++++++++++++++++ runtime/drivers/druid/olap.go | 14 +++ 3 files changed, 122 insertions(+) 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..2dbac65bea66 100644 --- a/runtime/drivers/druid/druidsqldriver/druid_api_sql_driver.go +++ b/runtime/drivers/druid/druidsqldriver/druid_api_sql_driver.go @@ -458,6 +458,8 @@ type DruidQueryContext struct { UseCache *bool `json:"useCache,omitempty"` PopulateCache *bool `json:"populateCache,omitempty"` Priority int `json:"priority,omitempty"` + UserEmail string `json:"rillUserEmail,omitempty"` + ServiceToken string `json:"rillServiceToken,omitempty"` } type DruidParameter struct { @@ -482,13 +484,19 @@ func newDruidRequest(query string, args []driver.NamedValue, queryCfg *QueryConf Value: arg.Value, } } + var useCache, populateCache *bool priority := 0 + userEmail := "" + serviceToken := "" if queryCfg != nil { useCache = queryCfg.UseCache populateCache = queryCfg.PopulateCache priority = queryCfg.Priority + userEmail = queryCfg.UserEmail + serviceToken = queryCfg.ServiceToken } + return &DruidRequest{ Query: query, Header: true, @@ -501,6 +509,8 @@ func newDruidRequest(query string, args []driver.NamedValue, queryCfg *QueryConf UseCache: useCache, PopulateCache: populateCache, Priority: priority, + UserEmail: userEmail, + ServiceToken: serviceToken, }, } } @@ -517,6 +527,8 @@ type QueryConfig struct { UseCache *bool PopulateCache *bool Priority int + UserEmail string + ServiceToken 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..94e54340c9ea --- /dev/null +++ b/runtime/drivers/druid/druidsqldriver/druid_api_sql_driver_test.go @@ -0,0 +1,96 @@ +package druidsqldriver + +import ( + "context" + "database/sql" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/require" +) + +// newTestServer returns a Druid SQL API stub that captures each request body into requests, +// and responds with a minimal valid arrayLines result (header, types header, one row). +func newTestServer(t *testing.T, requests *[]DruidRequest) *httptest.Server { + t.Helper() + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, err := io.ReadAll(r.Body) + require.NoError(t, err) + + var dr DruidRequest + require.NoError(t, json.Unmarshal(body, &dr)) + *requests = append(*requests, dr) + + w.WriteHeader(http.StatusOK) + _, err = w.Write([]byte("[\"n\"]\n[\"BIGINT\"]\n[1]\n")) + require.NoError(t, err) + })) +} + +func TestQueryContext(t *testing.T) { + var requests []DruidRequest + srv := newTestServer(t, &requests) + defer srv.Close() + + db, err := sql.Open("druid", srv.URL) + require.NoError(t, err) + defer db.Close() + + useCache := true + ctx := WithQueryConfig(context.Background(), &QueryConfig{ + UseCache: &useCache, + Priority: 3, + UserEmail: "user@example.com", + ServiceToken: "etl-bot", + }) + + rows, err := db.QueryContext(ctx, "SELECT 1") + require.NoError(t, err) + require.NoError(t, rows.Close()) + + require.Len(t, requests, 1) + qc := requests[0].Context + require.NotEmpty(t, qc.SQLQueryID) + require.True(t, qc.EnableTimeBoundaryPlanning) + require.NotNil(t, qc.UseCache) + require.True(t, *qc.UseCache) + require.Nil(t, qc.PopulateCache) + require.Equal(t, 3, qc.Priority) + require.Equal(t, "user@example.com", qc.UserEmail) + require.Equal(t, "etl-bot", qc.ServiceToken) +} + +func TestQueryContextDefaults(t *testing.T) { + var requests []DruidRequest + srv := newTestServer(t, &requests) + defer srv.Close() + + db, err := sql.Open("druid", srv.URL) + require.NoError(t, err) + defer db.Close() + + rows, err := db.QueryContext(context.Background(), "SELECT 1") + require.NoError(t, err) + require.NoError(t, rows.Close()) + + require.Len(t, requests, 1) + qc := requests[0].Context + require.NotEmpty(t, qc.SQLQueryID) + require.Nil(t, qc.UseCache) + require.Nil(t, qc.PopulateCache) + require.Zero(t, qc.Priority) + require.Empty(t, qc.UserEmail) + + // Fields with omitempty must not be serialized when unset. + b, err := json.Marshal(qc) + require.NoError(t, err) + var m map[string]any + require.NoError(t, json.Unmarshal(b, &m)) + require.NotContains(t, m, "rillUserEmail") + require.NotContains(t, m, "rillServiceToken") + require.NotContains(t, m, "priority") + require.NotContains(t, m, "useCache") +} diff --git a/runtime/drivers/druid/olap.go b/runtime/drivers/druid/olap.go index 301d8aeade97..2a7a9934ca8d 100644 --- a/runtime/drivers/druid/olap.go +++ b/runtime/drivers/druid/olap.go @@ -115,6 +115,20 @@ func (c *connection) Query(ctx context.Context, stmt *drivers.Statement) (res *d } queryCfg.Priority = stmt.Priority } + // The metrics view's query_attributes can identify the caller under the "user_email" and "service_token" keys + // (e.g. user_email: '{{ .user.email }}'); they are stamped on the Druid query context for attribution in Druid's query logs. + if email := stmt.QueryAttributes["user_email"]; email != "" { + if queryCfg == nil { + queryCfg = &druidsqldriver.QueryConfig{} + } + queryCfg.UserEmail = email + } + if svc := stmt.QueryAttributes["service_token"]; svc != "" { + if queryCfg == nil { + queryCfg = &druidsqldriver.QueryConfig{} + } + queryCfg.ServiceToken = svc + } if queryCfg != nil { ctx = druidsqldriver.WithQueryConfig(ctx, queryCfg)