query tags telemetry

Signed-off-by: Sreekanth Vadigi <sreekanth.vadigi@databricks.com>

Author: sreekanth-db

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(

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

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):

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."""