allow clients to set QUERY_TAG parameter via context

connection.go

@@ -112,6 +112,9 @@ func (sc *snowflakeConn) exec(
 	if key := ctx.Value(multiStatementCount); key != nil {
 		req.Parameters[string(multiStatementCount)] = key
 	}
+	if tag := ctx.Value(queryTag); tag != nil {
+		req.Parameters[string(queryTag)] = tag
+	}
 	logger.WithContext(ctx).Infof("parameters: %v", req.Parameters)
 
 	// handle bindings, if required

statement_test.go

@@ -830,3 +830,29 @@ func TestStatementQueryExecs(t *testing.T) {
 		}
 	})
 }
+
+func TestWithQueryTag(t *testing.T) {
+	runDBTest(t, func(dbt *DBTest) {
+		testQueryTag := "TEST QUERY TAG"
+		ctx := WithQueryTag(context.Background(), testQueryTag)
+
+		// This query itself will be part of the history and will have the query tag
+		rows := dbt.mustQueryContext(
+			ctx,
+			"SELECT QUERY_TAG FROM table(information_schema.query_history_by_session())")
+		if !rows.Next() {
+			t.Fatal("no rows")
+		}
+		var tag sql.NullString
+		err := rows.Scan(&tag)
+		if err != nil {
+			t.Error(err)
+		}
+		if !tag.Valid {
+			t.Fatal("no tag set")
+		}
+		if tag.String != testQueryTag {
+			t.Fatalf("expected tag '%s' but got '%s'", testQueryTag, tag.String)
+		}
+	})
+}

util.go

@@ -30,6 +30,7 @@ const (
 	arrowBatches            contextKey = "ARROW_BATCHES"
 	arrowAlloc              contextKey = "ARROW_ALLOC"
 	enableOriginalTimestamp contextKey = "ENABLE_ORIGINAL_TIMESTAMP"
+	queryTag                contextKey = "QUERY_TAG"
 )
 
 const (
@@ -114,6 +115,12 @@ func WithOriginalTimestamp(ctx context.Context) context.Context {
 	return context.WithValue(ctx, enableOriginalTimestamp, true)
 }
 
+// WithQueryTag returns a context that will set the given tag as the QUERY_TAG
+// parameter on any queries that are run
+func WithQueryTag(ctx context.Context, tag string) context.Context {
+	return context.WithValue(ctx, queryTag, tag)
+}
+
 // Get the request ID from the context if specified, otherwise generate one
 func getOrGenerateRequestIDFromContext(ctx context.Context) UUID {
 	requestID, ok := ctx.Value(snowflakeRequestIDKey).(UUID)