Skip to content

Add statement-level query tag support#341

Merged
vikrantpuppala merged 4 commits into
databricks:mainfrom
joohoyeo-dev:statement-level-query-tags
Apr 6, 2026
Merged

Add statement-level query tag support#341
vikrantpuppala merged 4 commits into
databricks:mainfrom
joohoyeo-dev:statement-level-query-tags

Conversation

@joohoyeo-dev

Copy link
Copy Markdown
Contributor

Summary

  • Adds per-statement query tag support via driverctx.NewContextWithQueryTags, allowing users to attach query tags to individual SQL statements through context
  • Tags are serialized into TExecuteStatementReq.ConfOverlay["query_tags"], consistent with the Python (#736) and NodeJS (#339) connector implementations
  • Previously only session-level query tags were supported (set once via WithSessionParams at connection time)

Usage

ctx := driverctx.NewContextWithQueryTags(context.Background(), map[string]string{
    "team": "data-eng",
    "app":  "etl-pipeline",
})
rows, err := db.QueryContext(ctx, "SELECT * FROM table")

Changes

File Description
driverctx/ctx.go NewContextWithQueryTags, QueryTagsFromContext, propagation in NewContextFromBackground
query_tags.go (new) SerializeQueryTags — map to wire format with escaping
connection.go Read tags from context → serialize → set ConfOverlay["query_tags"]
driverctx/ctx_test.go 5 tests for context helpers
query_tags_test.go (new) 13 tests for serialization (escaping, edge cases)
connection_test.go 6 integration tests verifying ConfOverlay behavior
examples/query_tags/main.go Updated with session + statement-level examples

Test plan

  • Unit tests for SerializeQueryTags covering nil, empty, single/multi tags, escaping of \, :, , in values and keys
  • Unit tests for NewContextWithQueryTags / QueryTagsFromContext including nil context, missing key, timeout preservation, background propagation
  • Integration tests verifying ConfOverlay["query_tags"] is correctly set (or absent) in captured TExecuteStatementReq
  • Verify existing tests still pass (CI)

This pull request was AI-assisted by Isaac.

@joohoyeo-dev

joohoyeo-dev commented Apr 2, 2026

Copy link
Copy Markdown
Contributor Author

⏺ Test 1 — Session-level only:
// Expected in Query History: {"team": "platform", "test": "go-connector-pr"}
db.QueryRowContext(context.Background(), "SELECT 1")

Test 2 — Statement-level tags:
// Expected in Query History: {"team": "query-tags-test", "application": "go-connector-e2e", "env": "staging"}
ctx := driverctx.NewContextWithQueryTags(context.Background(), map[string]string{
"team": "query-tags-test",
"application": "go-connector-e2e",
"env": "staging",
})
db.QueryRowContext(ctx, "SELECT 2")

Test 3 — Different statement-level tags:
// Expected in Query History: {"team": "analytics", "job": "daily-report"}
ctx2 := driverctx.NewContextWithQueryTags(context.Background(), map[string]string{
"team": "analytics",
"job": "daily-report",
})
db.QueryRowContext(ctx2, "SELECT 3")

Test 4 — No statement-level tags (session-level still applies):
// Expected in Query History: {"team": "platform", "test": "go-connector-pr"}
db.QueryRowContext(context.Background(), "SELECT 4")

@joohoyeo-dev

Copy link
Copy Markdown
Contributor Author
Screenshot 2026-04-01 at 9 05 50 PM Screenshot 2026-04-01 at 9 06 06 PM Screenshot 2026-04-01 at 9 06 11 PM Screenshot 2026-04-01 at 9 06 14 PM

@jiabin-hu jiabin-hu left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Were you able to run the sample code successfully?

EDIT: nvm saw your run results

Comment thread driverctx/ctx.go
// NewContextWithQueryTags creates a new context with per-statement query tags.
// These tags are serialized and passed via confOverlay as "query_tags" in TExecuteStatementReq.
// They apply only to the statement executed with this context and do not persist across queries.
func NewContextWithQueryTags(ctx context.Context, queryTags map[string]string) context.Context {

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not sure if this is the best way to create a context object with an optional param in Go. I'm not familiar w/ Go. Perhaps ask Claude to double check.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good question! This is actually the idiomatic Go pattern for per-request metadata in database/sql drivers. Since the QueryContext(ctx, query, args) interface is fixed by the standard library, context values are the standard way to pass per-request options.

This driver already uses the same pattern for other per-request data:

  • driverctx.NewContextWithConnId
  • driverctx.NewContextWithCorrelationId
  • driverctx.NewContextWithQueryId
  • driverctx.NewContextWithStagingInfo

So NewContextWithQueryTags follows the established convention.

Comment thread examples/query_tags/main.go Outdated

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Either in this PR or as follow up, we should allow user to pass in a map as query tags so it's consistent with the statement level. I did it in Python and plan to do in NodeJS, too.

databricks/databricks-sql-python@e916f71

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Great point — addressed in 606f44e. Added WithQueryTags(map[string]string) as a connector option that accepts a structured map and handles serialization internally, consistent with the statement-level API and your Python approach (databricks/databricks-sql-python@e916f71).

// Before: pre-serialized string
dbsql.WithSessionParams(map[string]string{
    "QUERY_TAGS": "team:engineering,app:etl",
})

// After: structured map (preferred)
dbsql.WithQueryTags(map[string]string{
    "team": "engineering",
    "app":  "etl",
})

WithSessionParams still works for backward compatibility. If both are used, WithQueryTags takes precedence (applied after in the options chain). Tests added in connector_test.go.

@joohoyeo-dev

Copy link
Copy Markdown
Contributor Author

checked the tags still work as expected in Logfood

@vikrantpuppala vikrantpuppala left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

please fix failing checks, lgtm otherwise, thanks!

joohoyeo added 3 commits April 2, 2026 08:41
Previously, query tags could only be set at the session level via
WithSessionParams during connection creation. This adds per-statement
query tag support, allowing different tags for each query execution.

Users pass query tags through context using the new
driverctx.NewContextWithQueryTags function. The tags are serialized
into the TExecuteStatementReq.ConfOverlay["query_tags"] field,
consistent with the Python and NodeJS connector implementations.

Co-authored-by: Isaac
Signed-off-by: Jooho Yeo <jooho.yeo@databricks.com>
Verifies that session-level tags (TOpenSessionReq.Configuration) and
statement-level tags (TExecuteStatementReq.ConfOverlay) are independent:
session params don't leak into ConfOverlay, and statement-level tags
are correctly set even when session-level tags exist.

Co-authored-by: Isaac
Signed-off-by: Jooho Yeo <jooho.yeo@databricks.com>
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>
@joohoyeo-dev joohoyeo-dev force-pushed the statement-level-query-tags branch from 606f44e to 04704e1 Compare April 2, 2026 15:50
@vikrantpuppala vikrantpuppala merged commit f850b47 into databricks:main Apr 6, 2026
3 checks passed
@vikrantpuppala vikrantpuppala mentioned this pull request Apr 21, 2026
2 tasks
vikrantpuppala added a commit that referenced this pull request Apr 21, 2026
## Summary
Bump `DriverVersion` to `1.11.0` and add the v1.11.0 section to
`CHANGELOG.md`.

### Changes since v1.10.0
- Enable telemetry by default with DSN-controlled priority (#320, #321,
#322, #349)
- Add SPOG (Custom URL) routing support via `x-databricks-org-id` header
(#347)
- Add statement-level query tag support (#341)
- Add AI coding agent detection to User-Agent header (#326)
- Fix CloudFetch returning stale column names from cached results (#351)
- Fix resource leak: close staging Rows in execStagingOperation (#325)

Internal/infra-only changes are omitted from the user-facing notes (CI
hardening, dependabot bumps, CODEOWNERS).

## Test plan
- [x] `go build ./...` clean
- [x] `go test ./... -count=1 -short` passes locally

## Next steps after merge
1. Tag the merge commit as `v1.11.0` and push the tag
2. Trigger `peco-databricks-sql-go` in
secure-public-registry-releases-eng with `ref=v1.11.0`, `dry-run=true`
to verify
3. Re-run with `dry-run=false` for the actual release

NO_CHANGELOG=true

This pull request was AI-assisted by Isaac.

Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>
atzoum pushed a commit to rudderlabs/sqlconnect-go that referenced this pull request Jun 9, 2026
….10.0 to 1.12.0 (#502)

Bumps
[github.com/databricks/databricks-sql-go](https://github.com/databricks/databricks-sql-go)
from 1.10.0 to 1.12.0.
<details>
<summary>Release notes</summary>
<p><em>Sourced from <a
href="https://github.com/databricks/databricks-sql-go/releases">github.com/databricks/databricks-sql-go's
releases</a>.</em></p>
<blockquote>
<h2>v1.12.0</h2>
<ul>
<li><code>databricks/databricks-sql-go#355</code><a
href="https://redirect.github.com/databricks/databricks-sql-go/issues/361">#361</a>)</li>
<li><code>databricks/databricks-sql-go#354</code><a
href="https://redirect.github.com/databricks/databricks-sql-go/issues/364">#364</a>)</li>
<li><code>databricks/databricks-sql-go#360</code><a
href="https://redirect.github.com/databricks/databricks-sql-go/issues/363">#363</a>)</li>
</ul>
<h2>v1.11.1</h2>
<ul>
<li><code>databricks/databricks-sql-go#357</code></li>
</ul>
<h2>v1.11.0</h2>
<ul>
<li><code>databricks/databricks-sql-go#320</code><a
href="https://redirect.github.com/databricks/databricks-sql-go/issues/321">#321</a>,
<a
href="https://redirect.github.com/databricks/databricks-sql-go/issues/322">#322</a>,
<a
href="https://redirect.github.com/databricks/databricks-sql-go/issues/349">#349</a>)</li>
<li>Add SPOG (Custom URL) routing support via
<code>x-databricks-org-id</code><code>databricks/databricks-sql-go#347</code></li>
<li><code>databricks/databricks-sql-go#341</code></li>
<li><code>databricks/databricks-sql-go#326</code></li>
<li><code>databricks/databricks-sql-go#351</code></li>
<li><code>databricks/databricks-sql-go#325</code></li>
</ul>
</blockquote>
</details>
<details>
<summary>Changelog</summary>
<p><em>Sourced from <a
href="https://github.com/databricks/databricks-sql-go/blob/main/CHANGELOG.md">github.com/databricks/databricks-sql-go's
changelog</a>.</em></p>
<blockquote>
<h2>v1.12.0 (2026-05-25)</h2>
<ul>
<li><code>databricks/databricks-sql-go#355</code><a
href="https://redirect.github.com/databricks/databricks-sql-go/issues/361">#361</a>)</li>
<li><code>databricks/databricks-sql-go#354</code><a
href="https://redirect.github.com/databricks/databricks-sql-go/issues/364">#364</a>)</li>
<li><code>databricks/databricks-sql-go#360</code><a
href="https://redirect.github.com/databricks/databricks-sql-go/issues/363">#363</a>)</li>
</ul>
<h2>v1.11.1 (2026-05-20)</h2>
<ul>
<li><code>databricks/databricks-sql-go#357</code></li>
</ul>
<h2>v1.11.0 (2026-04-16)</h2>
<ul>
<li><code>databricks/databricks-sql-go#320</code><a
href="https://redirect.github.com/databricks/databricks-sql-go/issues/321">#321</a>,
<a
href="https://redirect.github.com/databricks/databricks-sql-go/issues/322">#322</a>,
<a
href="https://redirect.github.com/databricks/databricks-sql-go/issues/349">#349</a>)</li>
<li>Add SPOG (Custom URL) routing support via
<code>x-databricks-org-id</code><code>databricks/databricks-sql-go#347</code></li>
<li><code>databricks/databricks-sql-go#341</code></li>
<li><code>databricks/databricks-sql-go#326</code></li>
<li><code>databricks/databricks-sql-go#351</code></li>
<li><code>databricks/databricks-sql-go#325</code></li>
</ul>
</blockquote>
</details>
<details>
<summary>Commits</summary>
<ul>
<li><a
href="https://github.com/databricks/databricks-sql-go/commit/c3be94a518f150367c56e063c9320069962667f3"><code>c3be94a</code></a>
Prepare for v1.12.0 release (<a
href="https://redirect.github.com/databricks/databricks-sql-go/issues/365">#365</a>)</li>
<li><a
href="https://github.com/databricks/databricks-sql-go/commit/7b961f22057c20926a95b9d50d1195f24f3ee7bc"><code>7b961f2</code></a>
Telemetry: normalize host key for per-host client + breaker registries
(<a
href="https://redirect.github.com/databricks/databricks-sql-go/issues/364">#364</a>)</li>
<li><a
href="https://github.com/databricks/databricks-sql-go/commit/de57a476f80bb02f93fcbfe0db07e4b1fa478831"><code>de57a47</code></a>
Telemetry: stop retrying into 429s, honour Retry-After, fix userAgent
(<a
href="https://redirect.github.com/databricks/databricks-sql-go/issues/354">#354</a>)</li>
<li><a
href="https://github.com/databricks/databricks-sql-go/commit/dc3c4f9f891eec92549012f702dd1892b71ad072"><code>dc3c4f9</code></a>
[ES-1911239] Retry transient S3 errors on staging PUT/GET/REMOVE (<a
href="https://redirect.github.com/databricks/databricks-sql-go/issues/361">#361</a>)</li>
<li><a
href="https://github.com/databricks/databricks-sql-go/commit/e248c0495747b8213deb3e1eb19d1233c7b822f7"><code>e248c04</code></a>
Bump golang-jwt, x/net, protobuf to clear Go-1.20-compatible CVEs (<a
href="https://redirect.github.com/databricks/databricks-sql-go/issues/363">#363</a>)</li>
<li><a
href="https://github.com/databricks/databricks-sql-go/commit/a97b104260857d9af2f0565cf0d6aa5f3d0b89c5"><code>a97b104</code></a>
[ES-1892645] Retry transient S3 errors in CloudFetch downloads (<a
href="https://redirect.github.com/databricks/databricks-sql-go/issues/355">#355</a>)</li>
<li><a
href="https://github.com/databricks/databricks-sql-go/commit/b31339e0aca915851f64b0936aee7f4252d69ad7"><code>b31339e</code></a>
[SIRT-1753] Bump go-jose/go-jose/v3 to v3.0.5 (CVE-2026-34986) (<a
href="https://redirect.github.com/databricks/databricks-sql-go/issues/360">#360</a>)</li>
<li><a
href="https://github.com/databricks/databricks-sql-go/commit/dd8a79f4c2717bb35bcee6f64b061a89659dc328"><code>dd8a79f</code></a>
Prepare for v1.11.1 release (<a
href="https://redirect.github.com/databricks/databricks-sql-go/issues/358">#358</a>)</li>
<li><a
href="https://github.com/databricks/databricks-sql-go/commit/2557e6fd17cb2b43c0e5420e003d3c5375eeaf1e"><code>2557e6f</code></a>
Fix CloudFetch goroutine leak that retains Arrow buffers after Close (<a
href="https://redirect.github.com/databricks/databricks-sql-go/issues/357">#357</a>)</li>
<li><a
href="https://github.com/databricks/databricks-sql-go/commit/f4d99924f71c94faffaa896ce8c3e427b303fb7b"><code>f4d9992</code></a>
Prepare for v1.11.0 release (<a
href="https://redirect.github.com/databricks/databricks-sql-go/issues/352">#352</a>)</li>
<li>Additional commits viewable in <a
href="https://github.com/databricks/databricks-sql-go/compare/v1.10.0...v1.12.0">compare
view</a></li>
</ul>
</details>
<br />


[![Dependabot compatibility
score](https://dependabot-badges.githubapp.com/badges/compatibility_score?dependency-name=github.com/databricks/databricks-sql-go&package-manager=go_modules&previous-version=1.10.0&new-version=1.12.0)](https://docs.github.com/en/github/managing-security-vulnerabilities/about-dependabot-security-updates#about-compatibility-scores)

Dependabot will resolve any conflicts with this PR as long as you don't
alter it yourself. You can also trigger a rebase manually by commenting
`@dependabot rebase`.

[//]: # (dependabot-automerge-start)
[//]: # (dependabot-automerge-end)

---

<details>
<summary>Dependabot commands and options</summary>
<br />

You can trigger Dependabot actions by commenting on this PR:
- `@dependabot rebase` will rebase this PR
- `@dependabot recreate` will recreate this PR, overwriting any edits
that have been made to it
- `@dependabot show <dependency name> ignore conditions` will show all
of the ignore conditions of the specified dependency
- `@dependabot ignore this major version` will close this PR and stop
Dependabot creating any more for this major version (unless you reopen
the PR or upgrade to it yourself)
- `@dependabot ignore this minor version` will close this PR and stop
Dependabot creating any more for this minor version (unless you reopen
the PR or upgrade to it yourself)
- `@dependabot ignore this dependency` will close this PR and stop
Dependabot creating any more for this dependency (unless you reopen the
PR or upgrade to it yourself)


</details>

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants