diff --git a/connection.go b/connection.go index 5fcb355bb..5f41bbb3b 100644 --- a/connection.go +++ b/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 diff --git a/doc.go b/doc.go index 0c882e09c..5cbda834e 100644 --- a/doc.go +++ b/doc.go @@ -167,6 +167,16 @@ Users can use SetLogger in driver.go to set a customized logger for gosnowflake 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 diff --git a/statement_test.go b/statement_test.go index a565cb046..f76975343 100644 --- a/statement_test.go +++ b/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) + }) +} diff --git a/util.go b/util.go index 32df50e8f..9cd1b1a27 100644 --- a/util.go +++ b/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)