Add statement-level query tag support via context

Previously, query tags could only be set at the session level via
WithSessionParams during connection creation. This adds per-statement
query tag support, allowing different tags for each query execution.

Users pass query tags through context using the new
driverctx.NewContextWithQueryTags function. The tags are serialized
into the TExecuteStatementReq.ConfOverlay["query_tags"] field,
consistent with the Python and NodeJS connector implementations.

Co-authored-by: Isaac

Author: joohoyeo

connection.go

@@ -323,6 +323,17 @@ func (c *conn) executeStatement(ctx context.Context, query string, args []driver
 		req.Parameters = parameters
 	}
 
+	// Add per-statement query tags if provided via context
+	if queryTags := driverctx.QueryTagsFromContext(ctx); len(queryTags) > 0 {
+		serialized := SerializeQueryTags(queryTags)
+		if serialized != "" {
+			if req.ConfOverlay == nil {
+				req.ConfOverlay = make(map[string]string)
+			}
+			req.ConfOverlay["query_tags"] = serialized
+		}
+	}
+
 	resp, err := c.client.ExecuteStatement(ctx, &req)
 	var log *logger.DBSQLLogger
 	log, ctx = client.LoggerAndContext(ctx, resp)

connection_test.go

@@ -10,6 +10,7 @@ import (
 	"github.com/apache/thrift/lib/go/thrift"
 	"github.com/pkg/errors"
 
+	"github.com/databricks/databricks-sql-go/driverctx"
 	dbsqlerr "github.com/databricks/databricks-sql-go/errors"
 	"github.com/databricks/databricks-sql-go/internal/cli_service"
 	"github.com/databricks/databricks-sql-go/internal/client"
@@ -493,6 +494,121 @@ func TestConn_executeStatement_ProtocolFeatures(t *testing.T) {
 	}
 }
 
+func TestConn_executeStatement_QueryTags(t *testing.T) {
+	t.Parallel()
+
+	makeTestConn := func(captureReq *(*cli_service.TExecuteStatementReq)) *conn {
+		executeStatement := func(ctx context.Context, req *cli_service.TExecuteStatementReq) (r *cli_service.TExecuteStatementResp, err error) {
+			*captureReq = req
+			return &cli_service.TExecuteStatementResp{
+				Status: &cli_service.TStatus{
+					StatusCode: cli_service.TStatusCode_SUCCESS_STATUS,
+				},
+				OperationHandle: &cli_service.TOperationHandle{
+					OperationId: &cli_service.THandleIdentifier{
+						GUID:   []byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16},
+						Secret: []byte("secret"),
+					},
+				},
+				DirectResults: &cli_service.TSparkDirectResults{
+					OperationStatus: &cli_service.TGetOperationStatusResp{
+						Status: &cli_service.TStatus{
+							StatusCode: cli_service.TStatusCode_SUCCESS_STATUS,
+						},
+						OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_FINISHED_STATE),
+					},
+				},
+			}, nil
+		}
+
+		return &conn{
+			session: getTestSession(),
+			client: &client.TestClient{
+				FnExecuteStatement: executeStatement,
+			},
+			cfg: config.WithDefaults(),
+		}
+	}
+
+	t.Run("query tags from context are set in ConfOverlay", func(t *testing.T) {
+		var capturedReq *cli_service.TExecuteStatementReq
+		testConn := makeTestConn(&capturedReq)
+
+		ctx := driverctx.NewContextWithQueryTags(context.Background(), map[string]string{
+			"team": "engineering",
+			"app":  "etl",
+		})
+
+		_, err := testConn.executeStatement(ctx, "SELECT 1", nil)
+		assert.NoError(t, err)
+		assert.NotNil(t, capturedReq.ConfOverlay)
+		// Map iteration is non-deterministic, so check both possible orderings
+		queryTags := capturedReq.ConfOverlay["query_tags"]
+		assert.True(t,
+			queryTags == "team:engineering,app:etl" || queryTags == "app:etl,team:engineering",
+			"unexpected query_tags value: %s", queryTags)
+	})
+
+	t.Run("no query tags in context means no ConfOverlay", func(t *testing.T) {
+		var capturedReq *cli_service.TExecuteStatementReq
+		testConn := makeTestConn(&capturedReq)
+
+		_, err := testConn.executeStatement(context.Background(), "SELECT 1", nil)
+		assert.NoError(t, err)
+		assert.Nil(t, capturedReq.ConfOverlay)
+	})
+
+	t.Run("empty query tags map means no ConfOverlay", func(t *testing.T) {
+		var capturedReq *cli_service.TExecuteStatementReq
+		testConn := makeTestConn(&capturedReq)
+
+		ctx := driverctx.NewContextWithQueryTags(context.Background(), map[string]string{})
+
+		_, err := testConn.executeStatement(ctx, "SELECT 1", nil)
+		assert.NoError(t, err)
+		assert.Nil(t, capturedReq.ConfOverlay)
+	})
+
+	t.Run("single query tag", func(t *testing.T) {
+		var capturedReq *cli_service.TExecuteStatementReq
+		testConn := makeTestConn(&capturedReq)
+
+		ctx := driverctx.NewContextWithQueryTags(context.Background(), map[string]string{
+			"team": "data-eng",
+		})
+
+		_, err := testConn.executeStatement(ctx, "SELECT 1", nil)
+		assert.NoError(t, err)
+		assert.Equal(t, "team:data-eng", capturedReq.ConfOverlay["query_tags"])
+	})
+
+	t.Run("query tags with special characters in values", func(t *testing.T) {
+		var capturedReq *cli_service.TExecuteStatementReq
+		testConn := makeTestConn(&capturedReq)
+
+		ctx := driverctx.NewContextWithQueryTags(context.Background(), map[string]string{
+			"url": "http://host:8080",
+		})
+
+		_, err := testConn.executeStatement(ctx, "SELECT 1", nil)
+		assert.NoError(t, err)
+		assert.Equal(t, `url:http\://host\:8080`, capturedReq.ConfOverlay["query_tags"])
+	})
+
+	t.Run("query tags with empty value", func(t *testing.T) {
+		var capturedReq *cli_service.TExecuteStatementReq
+		testConn := makeTestConn(&capturedReq)
+
+		ctx := driverctx.NewContextWithQueryTags(context.Background(), map[string]string{
+			"flag": "",
+		})
+
+		_, err := testConn.executeStatement(ctx, "SELECT 1", nil)
+		assert.NoError(t, err)
+		assert.Equal(t, "flag", capturedReq.ConfOverlay["query_tags"])
+	})
+}
+
 func TestConn_pollOperation(t *testing.T) {
 	t.Parallel()
 	t.Run("pollOperation returns finished state response when query finishes", func(t *testing.T) {

driverctx/ctx.go

@@ -15,6 +15,7 @@ const (
 	QueryIdCallbackKey
 	ConnIdCallbackKey
 	StagingAllowedLocalPathKey
+	QueryTagsContextKey
 )
 
 type IdCallbackFunc func(string)
@@ -107,16 +108,40 @@ func NewContextWithStagingInfo(ctx context.Context, stagingAllowedLocalPath []st
 	return context.WithValue(ctx, StagingAllowedLocalPathKey, stagingAllowedLocalPath)
 }
 
+// NewContextWithQueryTags creates a new context with per-statement query tags.
+// These tags are serialized and passed via confOverlay as "query_tags" in TExecuteStatementReq.
+// They apply only to the statement executed with this context and do not persist across queries.
+func NewContextWithQueryTags(ctx context.Context, queryTags map[string]string) context.Context {
+	return context.WithValue(ctx, QueryTagsContextKey, queryTags)
+}
+
+// QueryTagsFromContext retrieves the per-statement query tags stored in context.
+func QueryTagsFromContext(ctx context.Context) map[string]string {
+	if ctx == nil {
+		return nil
+	}
+
+	queryTags, ok := ctx.Value(QueryTagsContextKey).(map[string]string)
+	if !ok {
+		return nil
+	}
+	return queryTags
+}
+
 func NewContextFromBackground(ctx context.Context) context.Context {
 	connId := ConnIdFromContext(ctx)
 	corrId := CorrelationIdFromContext(ctx)
 	queryId := QueryIdFromContext(ctx)
 	stagingPaths := StagingPathsFromContext(ctx)
+	queryTags := QueryTagsFromContext(ctx)
 
 	newCtx := NewContextWithConnId(context.Background(), connId)
 	newCtx = NewContextWithCorrelationId(newCtx, corrId)
 	newCtx = NewContextWithQueryId(newCtx, queryId)
 	newCtx = NewContextWithStagingInfo(newCtx, stagingPaths)
+	if queryTags != nil {
+		newCtx = NewContextWithQueryTags(newCtx, queryTags)
+	}
 
 	return newCtx
 }

driverctx/ctx_test.go

@@ -8,6 +8,57 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
+func TestNewContextWithQueryTags(t *testing.T) {
+	t.Run("stores and retrieves query tags", func(t *testing.T) {
+		tags := map[string]string{"team": "engineering", "app": "etl"}
+		ctx := NewContextWithQueryTags(context.Background(), tags)
+		result := QueryTagsFromContext(ctx)
+		assert.Equal(t, tags, result)
+	})
+
+	t.Run("returns nil for context without query tags", func(t *testing.T) {
+		result := QueryTagsFromContext(context.Background())
+		assert.Nil(t, result)
+	})
+
+	t.Run("returns nil for nil context", func(t *testing.T) {
+		result := QueryTagsFromContext(nil)
+		assert.Nil(t, result)
+	})
+
+	t.Run("it maintains timeout", func(t *testing.T) {
+		ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond)
+		defer cancel()
+		tags := map[string]string{"team": "eng"}
+		ctx1 := NewContextWithQueryTags(ctx, tags)
+		result := QueryTagsFromContext(ctx1)
+		assert.Equal(t, tags, result)
+		dl, ok := ctx.Deadline()
+		dl1, ok1 := ctx1.Deadline()
+		assert.Equal(t, dl, dl1)
+		assert.True(t, ok)
+		assert.True(t, ok1)
+	})
+
+	t.Run("NewContextFromBackground preserves query tags", func(t *testing.T) {
+		tags := map[string]string{"team": "eng"}
+		ctx := NewContextWithConnId(context.Background(), "conn-1")
+		ctx = NewContextWithCorrelationId(ctx, "corr-1")
+		ctx = NewContextWithQueryTags(ctx, tags)
+
+		newCtx := NewContextFromBackground(ctx)
+		assert.Equal(t, tags, QueryTagsFromContext(newCtx))
+		assert.Equal(t, "conn-1", ConnIdFromContext(newCtx))
+		assert.Equal(t, "corr-1", CorrelationIdFromContext(newCtx))
+	})
+
+	t.Run("NewContextFromBackground without query tags", func(t *testing.T) {
+		ctx := NewContextWithConnId(context.Background(), "conn-1")
+		newCtx := NewContextFromBackground(ctx)
+		assert.Nil(t, QueryTagsFromContext(newCtx))
+	})
+}
+
 func TestNewContextWithCorrelationId(t *testing.T) {
 	t.Run("base case", func(t *testing.T) {
 

examples/query_tags/main.go

@@ -9,6 +9,7 @@ import (
 	"strconv"
 
 	dbsql "github.com/databricks/databricks-sql-go"
+	"github.com/databricks/databricks-sql-go/driverctx"
 	"github.com/joho/godotenv"
 )
 
@@ -21,6 +22,7 @@ func main() {
 		log.Fatal(err.Error())
 	}
 
+	// Session-level query tags: applied to all queries in this session.
 	connector, err := dbsql.NewConnector(
 		dbsql.WithServerHostname(os.Getenv("DATABRICKS_HOST")),
 		dbsql.WithPort(port),
@@ -38,16 +40,41 @@ func main() {
 	db := sql.OpenDB(connector)
 	defer db.Close()
 
+	// Example 1: Session-level query tags (set during connection)
+	fmt.Println("=== Session-level query tags ===")
 	ctx := context.Background()
 	var result int
 	err = db.QueryRowContext(ctx, "SELECT 1").Scan(&result)
 	if err != nil {
-		if err == sql.ErrNoRows {
-			fmt.Println("not found")
-			return
-		} else {
-			fmt.Printf("err: %v\n", err)
-		}
+		log.Printf("err: %v\n", err)
+	} else {
+		fmt.Println(result)
+	}
+
+	// Example 2: Statement-level query tags (per-query override via context)
+	fmt.Println("=== Statement-level query tags ===")
+	ctx = driverctx.NewContextWithQueryTags(context.Background(), map[string]string{
+		"team":        "data-eng",
+		"application": "etl-pipeline",
+		"env":         "production",
+	})
+	err = db.QueryRowContext(ctx, "SELECT 2").Scan(&result)
+	if err != nil {
+		log.Printf("err: %v\n", err)
+	} else {
+		fmt.Println(result)
+	}
+
+	// Example 3: Different query tags for a different statement
+	fmt.Println("=== Different statement-level query tags ===")
+	ctx = driverctx.NewContextWithQueryTags(context.Background(), map[string]string{
+		"team": "analytics",
+		"job":  "daily-report",
+	})
+	err = db.QueryRowContext(ctx, "SELECT 3").Scan(&result)
+	if err != nil {
+		log.Printf("err: %v\n", err)
+	} else {
+		fmt.Println(result)
 	}
-	fmt.Println(result)
 }

query_tags.go

@@ -0,0 +1,32 @@
+package dbsql
+
+import "strings"
+
+// SerializeQueryTags converts a map of query tags to the wire format string.
+// The format is comma-separated key:value pairs (e.g., "team:engineering,app:etl").
+//
+// Escaping rules (consistent with Python and NodeJS connectors):
+//   - Keys: only backslashes are escaped
+//   - Values: backslashes, colons, and commas are escaped with a leading backslash
+//   - Empty string values result in just the key being emitted (no colon)
+//
+// Returns empty string if the map is nil or empty.
+func SerializeQueryTags(tags map[string]string) string {
+	if len(tags) == 0 {
+		return ""
+	}
+
+	parts := make([]string, 0, len(tags))
+	for k, v := range tags {
+		escapedKey := strings.ReplaceAll(k, `\`, `\\`)
+		if v == "" {
+			parts = append(parts, escapedKey)
+		} else {
+			escapedValue := strings.ReplaceAll(v, `\`, `\\`)
+			escapedValue = strings.ReplaceAll(escapedValue, `:`, `\:`)
+			escapedValue = strings.ReplaceAll(escapedValue, `,`, `\,`)
+			parts = append(parts, escapedKey+":"+escapedValue)
+		}
+	}
+	return strings.Join(parts, ",")
+}