Repository: snowflakedb/gosnowflake
Pull request #427 (2021-05-11 22:18:37)

allow clients to set QUERY_TAG parameter via context
### Description
1. Add new `WithQueryTag` method to allow clients to pass the `QUERY_TAG` parameter to a query via the context
2. Add test to verify that we're setting the tag correctly

### Checklist
- [X] Code compiles correctly
- [ ] Run ``make fmt`` to fix inconsistent formats
- [X] Run ``make lint`` to get lint errors and fix all of them
- [X] Created tests which fail without the change (if possible)
- [x] All tests passing
- [ ] Extended the README / documentation, if necessary


GregOwen: @sfc-gh-jbahk here's the second of the two PRs I mentioned. This one should hopefully be pretty straightforward.
GregOwen: @sfc-gh-jbahk Sorry for the delay in getting back to you. Some of our customers have asked us to include Sigma-level metadata (username, Sigma document) for each Snowflake query so that they can do automated attribution of their Sigma-initiated Snowflake usage. Right now, we add this metadata as a JSON blob in a comment at the end of the query, but our customers have asked us to move this information into the `Query Tag` column so that it's easier for them to extract and analyze within Snowflake.
sfc-gh-jbahk: @GregOwen I see potential conflicts with this approach wherein the statement parameters that the customers might pass into the statement could either throw errors, not accepted by our server or conflict with our internal parameters. Another issue would be security threats including, but not limited to, injections. Could you give me a more specific example? Thanks.
GregOwen: That makes sense. We probably don't need this to be as wide an interface as it currently is.
Our requirements are:
1. we can pass a JSON blob containing metadata per-query
2. the per-query metadata will be associated with the query in a way that Sigma customers can parse and analyze within Snowflake (in particular, we don't want the JSON blob to be part of a comment within the text of the query itself)

The way that we're currently using it is something like this:
```
params := make(map[string]string)
params["QUERY_TAG"] = fmt.Sprintf(`{"request-id":"%s","email":"%s"}`, requestId, email) // possibly Sigma document ID, etc.
ctx = WithContextStatementParameters(ctx, params)

// execute the query with the given context
```
sfc-gh-jbahk: @GregOwen I see. Where would these parameters be used? Will it be inserted as a value of a column?
GregOwen: @sfc-gh-jbahk The idea is that the metadata will show up as a `Query Tag` column in the customer's History view. Here's an example:

<img width="1790" alt="Screen Shot 2021-05-13 at 1 50 22 PM" src="https://user-images.githubusercontent.com/1291832/118191398-5937f480-b3f9-11eb-9138-f65609736eef.png">

This way a customer's Snowflake admin can use Snowflake to ask questions like "how many queries did Sigma run?" or "which Sigma documents are responsible for my increased Snowflake usage?".

sfc-gh-jbahk: @GregOwen I've reviewed this with my team and think it should be okay to move forward but if you can change the code to be less generic such that the statement context doesn't accept anything the user inputs but rather tailor it towards a query tag. Maybe this can be of some use: https://docs.snowflake.com/en/sql-reference/parameters.html#query-tag.
GregOwen: @sfc-gh-jbahk sounds great - I reworked it to be more precisely targeted at `QUERY_TAG` and added a test
sfc-gh-jbahk: @GregOwen can you rebase and push this branch? 
GregOwen: @sfc-gh-jbahk done
GregOwen: Test failures look like they're related to test infrastructure/secrets:

```
Run ./ci/test.sh
  ./ci/test.sh
  shell: /usr/bin/bash --noprofile --norc -e -o pipefail {0}
  env:
    GOROOT: /opt/hostedtoolcache/go/1.15.13/x64
    PARAMETERS_SECRET: 
    CLOUD_PROVIDER: AWS
gpg: gcry_kdf_derive failed: Invalid data
gpg: decryption failed: No secret key
Error: Process completed with exit code 2.
```
tampajohn: @sfc-gh-jbahk any idea if/when this is planned to be merged?  Or what needs to be done to get it mergeable?
sfc-gh-jbahk: @sfc-gh-cshi can you take a look at this PR and see if it's something the go driver needs/wants? cc: @sfc-gh-hkapre 
DmitryMaletin: @sfc-gh-jbahk  is there any plan to complete the merge? 
sfc-gh-pfus: Hi @tampajohn ! Do you still need this? Can you rebase and solve conflicts and align to new test helper function?
sfc-gh-pfus: Sorry @tampajohn , I meant @GregOwen of course!
GregOwen: @sfc-gh-pfus yes, we would still like to merge this PR. Can you confirm that somebody at Snowflake will be responsive to changes to this PR? Last time we got it to the 1-yard line but couldn't get it merged because each new commit requires manual re-approval to run the CI tests.
sfc-gh-pfus: Yes, now go driver migrated to another team. I merged/closed/handled 13 PRs this week, so I believe that we can handle one more :)
github-actions[bot]: **CLA Assistant Lite bot:** <br/>Thank you for your submission, we really appreciate it. Like many open-source projects, we ask that you sign our [Contributor License Agreement](https://github.com/snowflakedb/CLA/blob/main/README.md) before we can accept your contribution. You can sign the CLA by just posting a Pull Request Comment same as the below format.<br/>
   - - -
   I have read the CLA Document and I hereby sign the CLA
   - - -
   <sub>You can retrigger this bot by commenting **recheck** in this Pull Request</sub>
GregOwen: One sec; making sure that Sigma knows about the CLA and is comfortable with me signing it
debugmiller: Would love to see this merged in!
GregOwen: I have read the CLA Document and I hereby sign the CLA
GregOwen: @sfc-gh-pfus I've rebased and we've signed the CLA (it took a while to get legal approval for that). Could you please kick off the tests?
sfc-gh-pfus: Hi @GregOwen ! Thanks again for your contribution. I left some two suggestions and in the meantime started jobs.
sfc-gh-pfus: Your test hangs. I believe this is because you forgot to defer rows closing.
GregOwen: @sfc-gh-pfus I've pushed changes that should address your comments; let me know if they look good to you
sfc-gh-pfus: Hi @GregOwen ! All of your checks are green (I fixed a flaky test in the meantime). Can you rebase your PR? I can't do it myself.
GregOwen: @sfc-gh-pfus rebased. It looks like we may get stuck in a "need to rebase" --> "need approval to run tests" loop since we're in different timezones. I'm based in EST; should we try to coordinate tomorrow so that I make my changes first thing in the morning my time (which I'm guessing will line up with the end of your workday)?

Patch:
diff --git a/connection.go b/connection.go
index 5fcb355..5f41bbb 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 0c882e0..5cbda83 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/util.go b/util.go
index 32df50e..9cd1b1a 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)


Test patch:
diff --git a/statement_test.go b/statement_test.go
index a565cb0..f769753 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)
+	})
+}