Test pr 427

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

doc.go

@@ -167,6 +167,16 @@
 In order to enable debug logging for the driver, user could use SetLogLevel("debug") in SFLogger interface
 as shown in demo code at cmd/logger.go. To redirect the logs SFlogger.SetOutput method could do the work.
 
+# Query tag
+
+A custom query tag can be set in the context. Each query run with this context
+will include the custom query tag as metadata that will appear in the Query Tag
+column in the Query History log. For example:
+
+	queryTag := "my custom query tag"
+	ctxWithQueryTag := WithQueryTag(ctx, queryTag)
+	rows, err := db.QueryContext(ctxWithQueryTag, query)
+
 # Query request ID
 
 A specific query request ID can be set in the context and will be passed through
@@ -949,7 +959,7 @@
 db.Query() function:
 
 	db.Query("GET file://<local_file> <stage_identifier> <optional_parameters>")
 
 "<local_file>" should include the file path as well as the name. Snowflake recommends using
 an absolute path rather than a relative path. For example:
 

statement_test.go

@@ -830,3 +830,23 @@ 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())")
+		defer rows.Close()
+
+		assertTrueF(t, rows.Next())
+		var tag sql.NullString
+		err := rows.Scan(&tag)
+		assertNilF(t, err)
+		assertTrueF(t, tag.Valid, "no QUERY_TAG set")
+		assertEqualF(t, tag.String, testQueryTag)
+	})
+}

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)