query tags integration

Signed-off-by: Sreekanth Vadigi <[email protected]>

Author: sreekanth-db

examples/query_tags_example.py

@@ -0,0 +1,30 @@
+import os
+import databricks.sql as sql
+
+"""
+This example demonstrates how to use Query Tags.
+
+Query Tags are key-value pairs that can be attached to SQL executions and will appear
+in the system.query.history table for analytical purposes.
+
+Format: "key1:value1,key2:value2,key3:value3"
+"""
+
+print("=== Query Tags Example ===\n")
+
+with sql.connect(
+    server_hostname=os.getenv("DATABRICKS_SERVER_HOSTNAME"),
+    http_path=os.getenv("DATABRICKS_HTTP_PATH"),
+    access_token=os.getenv("DATABRICKS_TOKEN"),
+    session_configuration={
+        'QUERY_TAGS': 'team:engineering,test:query-tags',
+        'ansi_mode': False
+    }
+) as connection:
+    
+    with connection.cursor() as cursor:
+        cursor.execute("SELECT 1")
+        result = cursor.fetchone()
+        print(f"  Result: {result[0]}")
+
+print("\n=== Query Tags Example Complete ===")

tests/e2e/test_driver.py

@@ -851,7 +851,7 @@ def test_socket_timeout_user_defined(self):
     def test_ssp_passthrough(self):
         for enable_ansi in (True, False):
             with self.cursor(
-                {"session_configuration": {"ansi_mode": enable_ansi}}
+                {"session_configuration": {"ansi_mode": enable_ansi, "QUERY_TAGS": "team:marketing,dashboard:abc123,driver:python"}}
             ) as cursor:
                 cursor.execute("SET ansi_mode")
                 assert list(cursor.fetchone()) == ["ansi_mode", str(enable_ansi)]

tests/unit/test_session.py

@@ -162,6 +162,18 @@ def test_configuration_passthrough(self, mock_client_class):
         call_kwargs = mock_client_class.return_value.open_session.call_args[1]
         assert call_kwargs["session_configuration"] == mock_session_config
 
+    @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
+    def test_query_tags_configuration_passthrough(self, mock_client_class):
+        """Test that Query Tags are properly passed through to the backend."""
+        query_tags_config = {"QUERY_TAGS": "team:marketing,dashboard:abc123"}
+        databricks.sql.connect(
+            session_configuration=query_tags_config, **self.DUMMY_CONNECTION_ARGS
+        )
+
+        call_kwargs = mock_client_class.return_value.open_session.call_args[1]
+        assert call_kwargs["session_configuration"] == query_tags_config
+        assert call_kwargs["session_configuration"]["QUERY_TAGS"] == "team:marketing,dashboard:abc123"
+
     @patch("%s.session.ThriftDatabricksClient" % PACKAGE_NAME)
     def test_initial_namespace_passthrough(self, mock_client_class):
         mock_cat = Mock()