Stamp caller attribution (rillUserEmail, rillServiceToken) on the Druid query context

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{}

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

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)