Skip to content

query tags integration #663

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 2 commits into
base: main
Choose a base branch
from

Conversation

sreekanth-db
Copy link

@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>
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
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy