Skip to content

Add query_tags parameter support for execute methods#736

Open
jiabin-hu wants to merge 1 commit intodatabricks:mainfrom
jiabin-hu:query_tags
Open

Add query_tags parameter support for execute methods#736
jiabin-hu wants to merge 1 commit intodatabricks:mainfrom
jiabin-hu:query_tags

Conversation

@jiabin-hu
Copy link

@jiabin-hu jiabin-hu commented Feb 7, 2026

Summary

This PR adds per-query query tags functionality that allows users to pass tags as a dictionary to the execute() and execute_async() 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

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.

1 participant