Add WithQueryTags connector option for session-level map support

Addresses review feedback from jiabin-hu:

1. Add WithQueryTags(map[string]string) as a connector option that
   accepts a structured map and serializes it internally, consistent
   with the statement-level API and the Python connector approach
   (databricks/databricks-sql-python@e916f71).

2. Context values pattern is the idiomatic Go approach for per-request
   metadata in database/sql drivers (same pattern used by ConnId,
   CorrelationId, QueryId, and StagingInfo in this driver).

Co-authored-by: Isaac
Signed-off-by: Jooho Yeo <jooho.yeo@databricks.com>

Author: joohoyeo

connector.go

@@ -235,6 +235,23 @@ func WithSessionParams(params map[string]string) ConnOption {
 	}
 }
 
+// WithQueryTags sets session-level query tags from a map.
+// Tags are serialized and passed as QUERY_TAGS in the session configuration.
+// All queries in the session will carry these tags unless overridden at the statement level.
+// This is the preferred way to set session-level query tags, as it handles serialization
+// and escaping automatically (consistent with the statement-level API).
+func WithQueryTags(tags map[string]string) ConnOption {
+	return func(c *config.Config) {
+		serialized := SerializeQueryTags(tags)
+		if serialized != "" {
+			if c.SessionParams == nil {
+				c.SessionParams = make(map[string]string)
+			}
+			c.SessionParams["QUERY_TAGS"] = serialized
+		}
+	}
+}
+
 // WithSkipTLSHostVerify disables the verification of the hostname in the TLS certificate.
 // WARNING:
 // When this option is used, TLS is susceptible to machine-in-the-middle attacks.

connector_test.go

@@ -268,6 +268,65 @@ func TestNewConnector(t *testing.T) {
 	})
 }
 
+func TestWithQueryTags(t *testing.T) {
+	t.Run("WithQueryTags serializes map into SessionParams QUERY_TAGS", func(t *testing.T) {
+		con, err := NewConnector(
+			WithQueryTags(map[string]string{
+				"team": "data-eng",
+			}),
+		)
+		require.NoError(t, err)
+		coni, ok := con.(*connector)
+		require.True(t, ok)
+		assert.Equal(t, "team:data-eng", coni.cfg.SessionParams["QUERY_TAGS"])
+	})
+
+	t.Run("WithQueryTags with multiple tags", func(t *testing.T) {
+		con, err := NewConnector(
+			WithQueryTags(map[string]string{
+				"team": "eng",
+				"app":  "etl",
+			}),
+		)
+		require.NoError(t, err)
+		coni, ok := con.(*connector)
+		require.True(t, ok)
+		// Map iteration is non-deterministic
+		qt := coni.cfg.SessionParams["QUERY_TAGS"]
+		assert.True(t, qt == "team:eng,app:etl" || qt == "app:etl,team:eng", "got: %s", qt)
+	})
+
+	t.Run("WithQueryTags with empty map does not set QUERY_TAGS", func(t *testing.T) {
+		con, err := NewConnector(
+			WithQueryTags(map[string]string{}),
+		)
+		require.NoError(t, err)
+		coni, ok := con.(*connector)
+		require.True(t, ok)
+		_, exists := coni.cfg.SessionParams["QUERY_TAGS"]
+		assert.False(t, exists)
+	})
+
+	t.Run("WithQueryTags overrides WithSessionParams QUERY_TAGS", func(t *testing.T) {
+		con, err := NewConnector(
+			WithSessionParams(map[string]string{
+				"QUERY_TAGS": "old:value",
+				"ansi_mode":  "false",
+			}),
+			WithQueryTags(map[string]string{
+				"team": "new-team",
+			}),
+		)
+		require.NoError(t, err)
+		coni, ok := con.(*connector)
+		require.True(t, ok)
+		// WithQueryTags should override the QUERY_TAGS from WithSessionParams
+		assert.Equal(t, "team:new-team", coni.cfg.SessionParams["QUERY_TAGS"])
+		// Other session params should be preserved
+		assert.Equal(t, "false", coni.cfg.SessionParams["ansi_mode"])
+	})
+}
+
 type mockRoundTripper struct{}
 
 var _ http.RoundTripper = mockRoundTripper{}

examples/query_tags/main.go

@@ -22,15 +22,17 @@ func main() {
 		log.Fatal(err.Error())
 	}
 
-	// Session-level query tags: applied to all queries in this session.
+	// Connection-level query tags: applied to all queries in this session.
+	// WithQueryTags accepts a map and handles serialization automatically.
 	connector, err := dbsql.NewConnector(
 		dbsql.WithServerHostname(os.Getenv("DATABRICKS_HOST")),
 		dbsql.WithPort(port),
 		dbsql.WithHTTPPath(os.Getenv("DATABRICKS_HTTPPATH")),
 		dbsql.WithAccessToken(os.Getenv("DATABRICKS_ACCESSTOKEN")),
-		dbsql.WithSessionParams(map[string]string{
-			"QUERY_TAGS": "team:engineering,test:query-tags,driver:go",
-			"ansi_mode":  "false",
+		dbsql.WithQueryTags(map[string]string{
+			"team":   "engineering",
+			"test":   "query-tags",
+			"driver": "go",
 		}),
 	)
 	if err != nil {
@@ -40,8 +42,8 @@ func main() {
 	db := sql.OpenDB(connector)
 	defer db.Close()
 
-	// Example 1: Session-level query tags (set during connection)
-	fmt.Println("=== Session-level query tags ===")
+	// Example 1: Connection-level query tags (set during connection)
+	fmt.Println("=== Connection-level query tags ===")
 	ctx := context.Background()
 	var result int
 	err = db.QueryRowContext(ctx, "SELECT 1").Scan(&result)