Skip to content

Add statement-level query_tags support for SEA backend#754

Merged
sreekanth-db merged 2 commits intomainfrom
sea-query-tags
Mar 16, 2026
Merged

Add statement-level query_tags support for SEA backend#754
sreekanth-db merged 2 commits intomainfrom
sea-query-tags

Conversation

@sreekanth-db
Copy link
Contributor

Summary

The query_tags feature was previously implemented only for the Thrift backend. Connection-level query tags (via session_configuration["QUERY_TAGS"]) already worked for SEA since it goes through _filter_session_configuration. This PR adds the missing statement-level (per-execute) query tags support for the SEA backend.

Root cause

SeaDatabricksClient.execute_command() accepted the query_tags parameter but had a # TODO comment and silently ignored it — the tags were never included in the ExecuteStatementRequest.

Wire format

The SEA REST API expects query_tags as an array of {key, value} objects, unlike Thrift which uses a serialized string in confOverlay:

"query_tags": [
  {"key": "team", "value": "data-eng"},
  {"key": "experimental"}
]
  • None values are passed through as-is — server handles them as key-only tags
  • Empty or None dict omits the field entirely from the request
  • No client-side escaping needed — the {key, value} JSON format handles special characters natively (unlike the serialized string format used in Thrift/session path)

Changes

  • src/databricks/sql/backend/sea/models/requests.py: Add query_tags field to ExecuteStatementRequest dataclass; serialize to [{key, value}] array in to_dict()
  • src/databricks/sql/backend/sea/backend.py: Pass query_tags to ExecuteStatementRequest constructor; remove TODO comment

Tests

Unit tests added (tests/unit/test_sea_backend.py)

  • test_execute_command_query_tags_string_values — string values serialized correctly to [{key, value}]
  • test_execute_command_query_tags_none_valueNone values passed through as {"key": k, "value": None}
  • test_execute_command_no_query_tags_omitted — field absent when query_tags=None
  • test_execute_command_empty_query_tags_omitted — field absent when query_tags={}
  • test_execute_command_async_query_tags — tags included in async (execute_async) path

CI checks verified locally

poetry run black --check src           ✓ all files unchanged
poetry run mypy --non-interactive src  ✓ no issues in 59 source files
poetry run python -m pytest tests/unit ✓ all passing

E2E test

Ran examples/query_tags_example.py with use_sea=True against a live Databricks SQL warehouse. All 4 examples passed — connection-level tags, per-query tags (including None values and special characters), async execution, and executemany. Verified tags stored correctly in query history UI including special characters (colons, commas, backslashes) — no escaping needed for the JSON {key, value} format.

Signed-off-by: Sreekanth Vadigi <sreekanth.vadigi@databricks.com>
Signed-off-by: Sreekanth Vadigi <sreekanth.vadigi@databricks.com>
@sreekanth-db sreekanth-db merged commit ca4d7bc into main Mar 16, 2026
35 of 37 checks passed
@sreekanth-db sreekanth-db deleted the sea-query-tags branch March 16, 2026 09:20
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.

2 participants