From 02c4b5f4ad7ce494cdd6950de967c609ba6384b8 Mon Sep 17 00:00:00 2001
From: Sreekanth Vadigi <sreekanth.vadigi@databricks.com>
Date: Fri, 28 Nov 2025 16:43:55 +0530
Subject: [PATCH] query tags telemetry

Signed-off-by: Sreekanth Vadigi <sreekanth.vadigi@databricks.com>
---
 src/databricks/sql/client.py                 |  2 ++
 src/databricks/sql/telemetry/models/event.py |  2 ++
 src/databricks/sql/utils.py                  | 16 ++++++++++++++++
 tests/unit/test_telemetry.py                 |  2 ++
 4 files changed, 22 insertions(+)

diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py
index a7f802dcd..70d4067e3 100755
--- a/src/databricks/sql/client.py
+++ b/src/databricks/sql/client.py
@@ -35,6 +35,7 @@
     ColumnTable,
     ColumnQueue,
     build_client_context,
+    get_session_config_value,
 )
 from databricks.sql.parameters.native import (
     DbsqlParameterBase,
@@ -386,6 +387,7 @@ def read(self) -> Optional[OAuthToken]:
             support_many_parameters=True,  # Native parameters supported
             enable_complex_datatype_support=_use_arrow_native_complex_types,
             allowed_volume_ingestion_paths=self.staging_allowed_local_path,
+            query_tags=get_session_config_value(session_configuration, 'query_tags'),
         )
 
         self._telemetry_client.export_initial_telemetry_log(
diff --git a/src/databricks/sql/telemetry/models/event.py b/src/databricks/sql/telemetry/models/event.py
index 2e6f63a6f..4d5a45038 100644
--- a/src/databricks/sql/telemetry/models/event.py
+++ b/src/databricks/sql/telemetry/models/event.py
@@ -57,6 +57,7 @@ class DriverConnectionParameters(JsonSerializableMixin):
         support_many_parameters (bool): Whether many parameters are supported
         enable_complex_datatype_support (bool): Whether complex datatypes are supported
         allowed_volume_ingestion_paths (str): Allowed paths for volume ingestion
+        query_tags (str): Query tags for tracking and attribution
     """
 
     http_path: str
@@ -84,6 +85,7 @@ class DriverConnectionParameters(JsonSerializableMixin):
     support_many_parameters: Optional[bool] = None
     enable_complex_datatype_support: Optional[bool] = None
     allowed_volume_ingestion_paths: Optional[str] = None
+    query_tags: Optional[str] = None
 
 
 @dataclass
diff --git a/src/databricks/sql/utils.py b/src/databricks/sql/utils.py
index b46784b10..1c49487e7 100644
--- a/src/databricks/sql/utils.py
+++ b/src/databricks/sql/utils.py
@@ -38,6 +38,22 @@
 logger = logging.getLogger(__name__)
 
 
+def get_session_config_value(
+    session_configuration: Optional[Dict[str, Any]],
+    key: str
+) -> Optional[str]:
+    """ Get a session configuration value with case-insensitive key matching """
+    if not session_configuration:
+        return None
+
+    key_upper = key.upper()
+    for k, v in session_configuration.items():
+        if k.upper() == key_upper:
+            return str(v) if v is not None else None
+
+    return None
+
+
 class ResultSetQueue(ABC):
     @abstractmethod
     def next_n_rows(self, num_rows: int):
diff --git a/tests/unit/test_telemetry.py b/tests/unit/test_telemetry.py
index 96a2f87d8..af9f7efc1 100644
--- a/tests/unit/test_telemetry.py
+++ b/tests/unit/test_telemetry.py
@@ -520,6 +520,7 @@ def test_driver_connection_parameters_all_fields(self):
             support_many_parameters=True,
             enable_complex_datatype_support=True,
             allowed_volume_ingestion_paths="/Volumes/catalog/schema/volume",
+            query_tags="team:engineering,project:telemetry",
         )
         
         # Serialize to JSON and parse back
@@ -552,6 +553,7 @@ def test_driver_connection_parameters_all_fields(self):
         assert json_dict["support_many_parameters"] is True
         assert json_dict["enable_complex_datatype_support"] is True
         assert json_dict["allowed_volume_ingestion_paths"] == "/Volumes/catalog/schema/volume"
+        assert json_dict["query_tags"] == "team:engineering,project:telemetry"
 
     def test_driver_connection_parameters_minimal_fields(self):
         """Test DriverConnectionParameters with only required fields."""