Skip to content

Add query_tags parameter support for execute methods#736

Merged
sreekanth-db merged 6 commits intodatabricks:mainfrom
jiabin-hu:query_tags
Mar 2, 2026
Merged

Add query_tags parameter support for execute methods#736
sreekanth-db merged 6 commits intodatabricks:mainfrom
jiabin-hu:query_tags

Conversation

@jiabin-hu
Copy link
Copy Markdown
Contributor

@jiabin-hu jiabin-hu commented Feb 7, 2026 •
edited
Loading

Summary

This PR adds per-query query tags functionality that allows users to pass tags as a dictionary to the execute.*() methods.

Changes

  • Added serialize_query_tags() utility function in utils.py to convert dict to the required format "key1:value1,key2:value2"
  • Added query_tags parameter to Cursor.execute() and Cursor.execute_async() methods
  • Updated DatabricksClient interface and ThriftDatabricksClient implementation
  • Query tags are serialized and passed via internal confOverlay with key "query_tags"
  • Added comprehensive unit tests for query_tags serialization (8 new tests)
  • Updated query_tags example to demonstrate new per-query tags usage

Key Features

  • None value handling: If value is None, omits colon and value (e.g., "key1:value1,key2,key3:value3")
  • Special character escaping: Escapes :, ,, \ in values with leading backslash
  • Keys not escaped: Assumes keys are controlled identifiers

Usage Example

cursor.execute(
    "SELECT * FROM table",
    query_tags={"team": "data-eng", "application": "etl"}
)

Testing

  • All existing tests pass (122/122 client tests, 66/66 thrift backend tests)
  • Added 8 new unit tests for query_tags serialization
  • All tests pass with no regressions

The branch is ready and all tests are passing! 🎉

What type of PR is this?

  • Refactor
  • Feature
  • Bug Fix
  • Other

Description

How is this tested?

  • Unit tests
  • E2E Tests
  • Manually
  • N/A

Related Tickets & Documents

sreekanth-db
sreekanth-db approved these changes Feb 10, 2026
View reviewed changes
Copy link
Copy Markdown
Contributor

@sreekanth-db sreekanth-db left a comment

Choose a reason for hiding this comment

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

LGTM, requesting @jprakash-db and @gopalldb to take a look

sreekanth-db
sreekanth-db reviewed Feb 10, 2026
View reviewed changes
Copy link
Copy Markdown
Contributor

@sreekanth-db sreekanth-db left a comment

Choose a reason for hiding this comment

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

Few PR checks are failing, please take a look at them @jiabin-hu

gopalldb
gopalldb reviewed Feb 11, 2026
View reviewed changes
@jiabin-hu
Add statement level query tag support by introducing it as a paramete…
9988722 
…r on execute* methods

Signed-off-by: Jiabin Hu <jiabin.hu@databricks.com>
@jiabin-hu
Add query_tags support to executemany method
99cb5f5 
- Added query_tags parameter to executemany() method
- Query tags are applied to all queries in the batch
- Updated example to demonstrate executemany usage with query_tags
- All tests pass (122/122 client tests)

Signed-off-by: Jiabin Hu <jiabin.hu@databricks.com>
@jiabin-hu
add example that doesn't have tag
0947c63 
Signed-off-by: Jiabin Hu <jiabin.hu@databricks.com>
@jiabin-hu
fix presubmit errors
3819a28 
Signed-off-by: Jiabin Hu <jiabin.hu@databricks.com>
@jiabin-hu
another lint
9492990 
Signed-off-by: Jiabin Hu <jiabin.hu@databricks.com>
@jiabin-hu
address review comments
23b6591 
Signed-off-by: Jiabin Hu <jiabin.hu@databricks.com>
@jiabin-hu
Copy link
Copy Markdown
Contributor Author

jiabin-hu commented Mar 1, 2026

jenkins merge

@sreekanth-db sreekanth-db merged commit 38097f2 into databricks:main Mar 2, 2026
68 of 72 checks passed
@joohoyeo-dev joohoyeo-dev mentioned this pull request Apr 2, 2026
Merged
4 tasks
vikrantpuppala added a commit to databricks/databricks-sql-go that referenced this pull request Apr 6, 2026
@joohoyeo-dev @joohoyeo @vikrantpuppala
Add statement-level query tag support (#341)
f850b47 
## 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](databricks/databricks-sql-python#736))
and NodeJS
([#339](databricks/databricks-sql-nodejs#339))
connector implementations
- Previously only session-level query tags were supported (set once via
`WithSessionParams` at connection time)

## Usage
```go
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
- [x] Unit tests for `SerializeQueryTags` covering nil, empty,
single/multi tags, escaping of `\`, `:`, `,` in values and keys
- [x] Unit tests for `NewContextWithQueryTags` / `QueryTagsFromContext`
including nil context, missing key, timeout preservation, background
propagation
- [x] 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.

---------

Signed-off-by: Jooho Yeo <jooho.yeo@databricks.com>
Co-authored-by: Jooho Yeo <jooho.yeo@databricks.com>
Co-authored-by: Vikrant Puppala <vikrant.puppala@databricks.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@sreekanth-db sreekanth-db sreekanth-db approved these changes

@gopalldb gopalldb gopalldb approved these changes

@jprakash-db jprakash-db Awaiting requested review from jprakash-db

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@jiabin-hu @gopalldb @sreekanth-db