added the get_attribute functions to the classes

Signed-off-by: Sai Shree Pradhan <saishree.pradhan@databricks.com>

Author: saishreeeee

.github/workflows/code-quality-checks.yml

@@ -112,7 +112,7 @@ jobs:
       #              run test suite
       #----------------------------------------------
       -   name: Run tests
-          run: poetry run python -m pytest tests/unit -v -s
+          run: poetry run python -m pytest tests/unit -v -s 
   check-linting:
     runs-on: ubuntu-latest
     strategy:

src/databricks/sql/client.py

@@ -31,6 +31,8 @@
     transform_paramstyle,
     ColumnTable,
     ColumnQueue,
+    ArrowQueue,
+    CloudFetchQueue,
 )
 from databricks.sql.parameters.native import (
     DbsqlParameterBase,
@@ -62,6 +64,7 @@
     HostDetails,
 )
 from databricks.sql.telemetry.latency_logger import log_latency
+from databricks.sql.telemetry.models.enums import ExecutionResultFormat, StatementType
 
 logger = logging.getLogger(__name__)
 
@@ -1355,6 +1358,35 @@ def setoutputsize(self, size, column=None):
         """Does nothing by default"""
         pass
 
+    def get_statement_id(self) -> Optional[str]:
+        return self.query_id
+
+    def get_session_id_hex(self) -> Optional[str]:
+        return self.connection.get_session_id_hex()
+
+    def get_is_compressed(self) -> bool:
+        return self.connection.lz4_compression
+
+    def get_execution_result(self) -> ExecutionResultFormat:
+        if self.active_result_set is None:
+            return ExecutionResultFormat.FORMAT_UNSPECIFIED
+
+        if isinstance(self.active_result_set.results, ColumnQueue):
+            return ExecutionResultFormat.COLUMNAR_INLINE
+        elif isinstance(self.active_result_set.results, CloudFetchQueue):
+            return ExecutionResultFormat.EXTERNAL_LINKS
+        elif isinstance(self.active_result_set.results, ArrowQueue):
+            return ExecutionResultFormat.INLINE_ARROW
+        return ExecutionResultFormat.FORMAT_UNSPECIFIED
+
+    def get_retry_count(self) -> int:
+        # return len(self.thrift_backend.retry_policy.history)
+        return 0
+
+    def get_statement_type(self, func_name: str) -> StatementType:
+        # TODO: Implement this
+        return StatementType.SQL
+
 
 class ResultSet:
     def __init__(
@@ -1654,3 +1686,31 @@ def map_col_type(type_):
             (column.name, map_col_type(column.datatype), None, None, None, None, None)
             for column in table_schema_message.columns
         ]
+
+    def get_statement_id(self) -> Optional[str]:
+        if self.command_id:
+            return str(UUID(bytes=self.command_id.operationId.guid))
+        return None
+
+    def get_session_id_hex(self) -> Optional[str]:
+        return self.connection.get_session_id_hex()
+
+    def get_is_compressed(self) -> bool:
+        return self.lz4_compressed
+
+    def get_execution_result(self) -> ExecutionResultFormat:
+        if isinstance(self.results, ColumnQueue):
+            return ExecutionResultFormat.COLUMNAR_INLINE
+        elif isinstance(self.results, CloudFetchQueue):
+            return ExecutionResultFormat.EXTERNAL_LINKS
+        elif isinstance(self.results, ArrowQueue):
+            return ExecutionResultFormat.INLINE_ARROW
+        return ExecutionResultFormat.FORMAT_UNSPECIFIED
+
+    def get_statement_type(self, func_name: str) -> StatementType:
+        # TODO: Implement this
+        return StatementType.SQL
+
+    def get_retry_count(self) -> int:
+        # return len(self.thrift_backend.retry_policy.history)
+        return 0

src/databricks/sql/telemetry/latency_logger.py

@@ -1,91 +1,9 @@
 import time
 import functools
-from typing import Optional
-from uuid import UUID
 from databricks.sql.telemetry.telemetry_client import TelemetryClientFactory
 from databricks.sql.telemetry.models.event import (
     SqlExecutionEvent,
 )
-from databricks.sql.telemetry.models.enums import ExecutionResultFormat, StatementType
-from databricks.sql.utils import ColumnQueue
-
-# Helper to get statement_id/query_id from instance if available
-def _get_statement_id(instance) -> Optional[str]:
-    """
-    Get statement ID from an instance using various methods:
-    1. For Cursor: Use query_id property which returns UUID from active_op_handle
-    2. For ResultSet: Use command_id which contains operationId
-
-    Note: ThriftBackend itself doesn't have a statement ID since one backend
-    can handle multiple concurrent operations/cursors.
-    """
-    if hasattr(instance, "query_id"):
-        return instance.query_id
-
-    if hasattr(instance, "command_id") and instance.command_id:
-        return str(UUID(bytes=instance.command_id.operationId.guid))
-
-    return None
-
-
-def _get_session_id_hex(instance) -> Optional[str]:
-    if hasattr(instance, "connection") and instance.connection:
-        return instance.connection.get_session_id_hex()
-    if hasattr(instance, "get_session_id_hex"):
-        return instance.get_session_id_hex()
-    return None
-
-
-def _get_statement_type(func_name: str) -> StatementType:  # TODO: implement this
-    return StatementType.SQL
-
-
-def _get_is_compressed(instance) -> bool:
-    """
-    Get compression status from instance:
-    1. Direct lz4_compression attribute (Connection)
-    2. Through connection attribute (Cursor/ResultSet)
-    3. Through thrift_backend attribute (Cursor)
-    """
-    if hasattr(instance, "lz4_compression"):
-        return instance.lz4_compression
-    if hasattr(instance, "connection") and instance.connection:
-        return instance.connection.lz4_compression
-    if hasattr(instance, "thrift_backend") and instance.thrift_backend:
-        return instance.thrift_backend.lz4_compressed
-    return False
-
-
-def _get_execution_result(instance) -> ExecutionResultFormat:
-    """
-    Get execution result format from instance:
-    1. For ResultSet: Check if using cloud fetch (external_links) or arrow/columnar format
-    2. For Cursor: Check through active_result_set
-    3. For ThriftBackend: Check result format from server
-    """
-    if hasattr(instance, "_use_cloud_fetch") and instance._use_cloud_fetch:
-        return ExecutionResultFormat.EXTERNAL_LINKS
-
-    if hasattr(instance, "active_result_set") and instance.active_result_set:
-        if isinstance(instance.active_result_set.results, ColumnQueue):
-            return ExecutionResultFormat.COLUMNAR_INLINE
-        return ExecutionResultFormat.INLINE_ARROW
-
-    if hasattr(instance, "thrift_backend") and instance.thrift_backend:
-        if hasattr(instance.thrift_backend, "_use_arrow_native_complex_types"):
-            return ExecutionResultFormat.INLINE_ARROW
-
-    return ExecutionResultFormat.FORMAT_UNSPECIFIED
-
-
-def _get_retry_count(instance) -> int:
-    """
-    Get retry count from instance by checking retry_policy.history length.
-    The retry_policy is only accessible through thrift_backend.
-    """
-    # TODO: implement this
-
-    return 0
 
 
 def log_latency():
@@ -101,30 +19,24 @@ def wrapper(self, *args, **kwargs):
                 end_time = time.perf_counter()
                 duration_ms = int((end_time - start_time) * 1000)
 
-                session_id_hex = _get_session_id_hex(self)
-
-                if session_id_hex:
-                    statement_id = _get_statement_id(self)
-                    statement_type = _get_statement_type(func.__name__)
-                    is_compressed = _get_is_compressed(self)
-                    execution_result = _get_execution_result(self)
-                    retry_count = _get_retry_count(self)
-
-                    sql_exec_event = SqlExecutionEvent(
-                        statement_type=statement_type,
-                        is_compressed=is_compressed,
-                        execution_result=execution_result,
-                        retry_count=retry_count,
-                    )
-
-                    telemetry_client = TelemetryClientFactory.get_telemetry_client(
-                        session_id_hex
-                    )
-                    telemetry_client.export_latency_log(
-                        latency_ms=duration_ms,
-                        sql_execution_event=sql_exec_event,
-                        sql_statement_id=statement_id,
-                    )
+                session_id_hex = self.get_session_id_hex()
+                statement_id = self.get_statement_id()
+
+                sql_exec_event = SqlExecutionEvent(
+                    statement_type=self.get_statement_type(func.__name__),
+                    is_compressed=self.get_is_compressed(),
+                    execution_result=self.get_execution_result(),
+                    retry_count=self.get_retry_count(),
+                )
+
+                telemetry_client = TelemetryClientFactory.get_telemetry_client(
+                    session_id_hex
+                )
+                telemetry_client.export_latency_log(
+                    latency_ms=duration_ms,
+                    sql_execution_event=sql_exec_event,
+                    sql_statement_id=statement_id,
+                )
 
         return wrapper
 

src/databricks/sql/thrift_backend.py

@@ -41,7 +41,7 @@
     convert_column_based_set_to_arrow_table,
 )
 from databricks.sql.types import SSLOptions
-from typing import Optional
+
 
 logger = logging.getLogger(__name__)
 
@@ -584,9 +584,6 @@ def open_session(self, session_configuration, catalog, schema):
             self._transport.close()
             raise
 
-    def get_session_id_hex(self) -> Optional[str]:
-        return self._session_id_hex
-
     def close_session(self, session_handle) -> None:
         req = ttypes.TCloseSessionReq(sessionHandle=session_handle)
         try: