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

[samikshya-db]

can we club this test into some other e2e test? (as this may not require a separate test tbh)

[sreekanth-db]

We are actually not creating a new test, just updating the existing one to include query tags in both SEA and Thrift modes. If the existing test itself is redundant, i think we can take that up in a separate refactoring PR

[samikshya-db]

I was asking as the following were not present

[
            {
                "use_sea": False,
            },
            {
                "use_sea": True,
            },
        ],
    )

Let's club where this is already present (as this adds to the time taking to run all tests)

[samikshya-db]

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