query tags integration

[sreekanth-db]

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

[nikhilsuri-db]

Should we add a test with invalid query param format & see if we are getting the required error?

[sreekanth-db]

In the backend design docs, it is mentioned that invalid query tags will be ignored without raising error so that actual query execution won't result in failure. Tested manually in local by passing invalid/non-parsable query tags - backend doesn't throw error

[nikhilsuri-db]

Similarly two scenarios we can add here

1. Not passing the query tags config param
2. Passing non parsable query tags -> "team:marketing,test"

[sreekanth-db]

1. _filter_session_configuration method loops through provided session config fields and filters out invalid params, so separate unit test for not passing query tags would be an overkill. For e2e tests, we already have few tests in which we don't pass query tags as part of session config.

2. We don't do any parsing of query tags in the driver and we send the user provided value as is to the backend. As mentioned in the other comment, backend would ignore the non parsable query tags without raising the error