Skip to content

Conversation

@sreekanth-db
Copy link
Contributor

@sreekanth-db sreekanth-db commented Aug 5, 2025

What type of PR is this?

  • Refactor
  • Feature
  • Bug Fix
  • Other

Description

This PR adds Query Tags support to the Python SQL Connector. Query Tags allow users to attach key-value pairs to SQL executions that appear in the system.query.history table.

Changes Made:

  1. Created Example File (examples/query_tags_example.py):

    • New file demonstrating Query Tags usage with session configuration
  2. Enhanced Integration Test (tests/e2e/test_driver.py):

    • Modified test_ssp_passthrough() to include Query Tags in session configuration
    • Added parametrization to test both Thrift and SEA backends
  3. Enhanced Unit Tests (tests/unit/test_session.py):

    • Updated test_configuration_passthrough() test to verify query tags are properly passed through to the thrift backend
    • Updated test_session_management() to include query tags in session configuration to ensure query tags work correctly in SEA backend

Implementation Details:

Added Query Tags Support in SEA flow (src/databricks/sql/backend/sea/utils/constants.py):

  • Added "QUERY_TAGS": "" to ALLOWED_SESSION_CONF_TO_DEFAULT_VALUES_MAP to enable query tags support in the SEA (Statement Execution API) backend flow

Usage example:

import databricks.sql

connection = databricks.sql.connect(
    server_hostname='host',
    http_path='/sql/1.0/warehouses/xxx',
    access_token='token',
    session_configuration={
        'QUERY_TAGS': 'team:engineering,test:query-tags',
        'ansi_mode': False
    }
)

How is this tested?

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

Testing Details:

Unit Tests:

  • test_configuration_passthrough() - Verifies Query Tags flow in thrift mode
  • test_session_management() - Verifies Query Tags flow in SEA mode
  • All existing tests still pass

E2E Tests:

  • Enhanced test_ssp_passthrough() to include Query Tags in both thirft and SEA modes

Manual Testing:

  • Ran examples/query_tags_example.py with real backend connection
  • Verified Query Tags are accepted and queries execute successfully

Related Tickets & Documents

Signed-off-by: Sreekanth Vadigi <sreekanth.vadigi@databricks.com>
Signed-off-by: Sreekanth Vadigi <sreekanth.vadigi@databricks.com>
Copy link
Contributor

@samikshya-db samikshya-db left a comment

Choose a reason for hiding this comment

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

a minor nit, otherwise LGTM. Thanks for making the change

Signed-off-by: Sreekanth Vadigi <sreekanth.vadigi@databricks.com>
@sreekanth-db sreekanth-db merged commit bcab1df into databricks:main Aug 25, 2025
27 of 29 checks passed
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.

4 participants