statement type, unit test fix

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

Author: saishreeeee

src/databricks/sql/client.py

@@ -62,6 +62,7 @@
     HostDetails,
 )
 from databricks.sql.telemetry.latency_logger import log_latency
+from databricks.sql.telemetry.models.enums import StatementType
 
 logger = logging.getLogger(__name__)
 
@@ -827,7 +828,7 @@ def _handle_staging_remove(
                 session_id_hex=self.connection.get_session_id_hex(),
             )
 
-    @log_latency()
+    @log_latency(StatementType.SQL)
     def execute(
         self,
         operation: str,
@@ -918,7 +919,7 @@ def execute(
 
         return self
 
-    @log_latency()
+    @log_latency(StatementType.SQL)
     def execute_async(
         self,
         operation: str,
@@ -1044,7 +1045,7 @@ def executemany(self, operation, seq_of_parameters):
             self.execute(operation, parameters)
         return self
 
-    @log_latency()
+    @log_latency(StatementType.METADATA)
     def catalogs(self) -> "Cursor":
         """
         Get all available catalogs.
@@ -1068,7 +1069,7 @@ def catalogs(self) -> "Cursor":
         )
         return self
 
-    @log_latency()
+    @log_latency(StatementType.METADATA)
     def schemas(
         self, catalog_name: Optional[str] = None, schema_name: Optional[str] = None
     ) -> "Cursor":
@@ -1097,7 +1098,7 @@ def schemas(
         )
         return self
 
-    @log_latency()
+    @log_latency(StatementType.METADATA)
     def tables(
         self,
         catalog_name: Optional[str] = None,
@@ -1133,7 +1134,7 @@ def tables(
         )
         return self
 
-    @log_latency()
+    @log_latency(StatementType.METADATA)
     def columns(
         self,
         catalog_name: Optional[str] = None,
@@ -1444,6 +1445,7 @@ def _convert_arrow_table(self, table):
     def rownumber(self):
         return self._next_row_index
 
+    @log_latency()
     def fetchmany_arrow(self, size: int) -> "pyarrow.Table":
         """
         Fetch the next set of rows of a query result, returning a PyArrow table.
@@ -1486,6 +1488,7 @@ def merge_columnar(self, result1, result2):
         ]
         return ColumnTable(merged_result, result1.column_names)
 
+    @log_latency()
     def fetchmany_columnar(self, size: int):
         """
         Fetch the next set of rows of a query result, returning a Columnar Table.
@@ -1511,6 +1514,7 @@ def fetchmany_columnar(self, size: int):
 
         return results
 
+    @log_latency()
     def fetchall_arrow(self) -> "pyarrow.Table":
         """Fetch all (remaining) rows of a query result, returning them as a PyArrow table."""
         results = self.results.remaining_rows()
@@ -1537,6 +1541,7 @@ def fetchall_arrow(self) -> "pyarrow.Table":
             return pyarrow.Table.from_pydict(data)
         return results
 
+    @log_latency()
     def fetchall_columnar(self):
         """Fetch all (remaining) rows of a query result, returning them as a Columnar table."""
         results = self.results.remaining_rows()

src/databricks/sql/telemetry/latency_logger.py

@@ -45,9 +45,6 @@ def get_session_id_hex(self):
     def get_statement_id(self):
         pass
 
-    def get_statement_type(self):
-        pass
-
     def get_is_compressed(self):
         pass
 
@@ -95,10 +92,6 @@ def get_retry_count(self) -> int:
             return len(self.thrift_backend.retry_policy.history)
         return 0
 
-    def get_statement_type(self) -> StatementType:
-        # TODO: Implement this
-        return StatementType.SQL
-
 
 class ResultSetExtractor(TelemetryExtractor):
     """
@@ -128,10 +121,6 @@ def get_execution_result(self) -> ExecutionResultFormat:
             return ExecutionResultFormat.INLINE_ARROW
         return ExecutionResultFormat.FORMAT_UNSPECIFIED
 
-    def get_statement_type(self) -> StatementType:
-        # TODO: Implement this
-        return StatementType.SQL
-
     def get_retry_count(self) -> int:
         if (
             hasattr(self.thrift_backend, "retry_policy")
@@ -166,7 +155,7 @@ def get_extractor(obj):
         raise NotImplementedError(f"No extractor found for {obj.__class__.__name__}")
 
 
-def log_latency():
+def log_latency(statement_type: StatementType = StatementType.NONE):
     """
     Decorator for logging execution latency and telemetry information.
 
@@ -180,17 +169,15 @@ def log_latency():
     - Creates a SqlExecutionEvent with execution details
     - Sends the telemetry data asynchronously via TelemetryClient
 
+    Args:
+        statement_type (StatementType): The type of SQL statement being executed.
+
     Usage:
-        @log_latency()
+        @log_latency(StatementType.SQL)
         def execute(self, query):
             # Method implementation
             pass
 
-        @log_latency()
-        def fetchall(self):
-            # Method implementation
-            pass
-
     Returns:
         function: A decorator that wraps methods to add latency logging.
 
@@ -216,7 +203,7 @@ def wrapper(self, *args, **kwargs):
                 statement_id = extractor.get_statement_id()
 
                 sql_exec_event = SqlExecutionEvent(
-                    statement_type=extractor.get_statement_type(),
+                    statement_type=statement_type,
                     is_compressed=extractor.get_is_compressed(),
                     execution_result=extractor.get_execution_result(),
                     retry_count=extractor.get_retry_count(),

tests/unit/test_client.py

@@ -8,6 +8,19 @@
 from datetime import datetime, date
 from uuid import UUID
 
+def noop_log_latency_decorator(*args, **kwargs):
+    """
+    This is a no-op decorator. It is used to patch the log_latency decorator
+    during tests, so that the tests for client logic are not affected by the
+    telemetry logging logic. It accepts any arguments and returns a decorator
+    that returns the original function unmodified.
+    """
+    def decorator(func):
+        return func
+    return decorator
+
+patch('databricks.sql.telemetry.latency_logger.log_latency', new=noop_log_latency_decorator).start()
+
 from databricks.sql.thrift_api.TCLIService.ttypes import (
     TOpenSessionResp,
     TExecuteStatementResp,
@@ -38,11 +51,6 @@ def new(cls):
         cls.apply_property_to_mock(ThriftBackendMock, staging_allowed_local_path=None)
         MockTExecuteStatementResp = MagicMock(spec=TExecuteStatementResp())
 
-        # Mock retry_policy with history attribute
-        mock_retry_policy = Mock()
-        mock_retry_policy.history = []
-        cls.apply_property_to_mock(ThriftBackendMock, retry_policy=mock_retry_policy)
-
         cls.apply_property_to_mock(
             MockTExecuteStatementResp,
             description=None,
@@ -74,15 +82,6 @@ def apply_property_to_mock(self, mock_obj, **kwargs):
             prop = PropertyMock(**kwargs)
             setattr(type(mock_obj), key, prop)
 
-    @classmethod
-    def mock_thrift_backend_with_retry_policy(cls): # Required for log_latency() decorator
-        """Create a simple thrift_backend mock with retry_policy for basic tests."""
-        mock_thrift_backend = Mock()
-        mock_retry_policy = Mock()
-        mock_retry_policy.history = []
-        mock_thrift_backend.retry_policy = mock_retry_policy
-        return mock_thrift_backend
-
 
 class ClientTestSuite(unittest.TestCase):
     """
@@ -332,7 +331,7 @@ def test_executing_multiple_commands_uses_the_most_recent_command(
         mock_result_sets[1].fetchall.assert_called_once_with()
 
     def test_closed_cursor_doesnt_allow_operations(self):
-        cursor = client.Cursor(Mock(), ThriftBackendMockFactory.mock_thrift_backend_with_retry_policy())
+        cursor = client.Cursor(Mock(), Mock())
         cursor.close()
 
         with self.assertRaises(Error) as e:
@@ -394,7 +393,7 @@ def test_get_schemas_parameters_passed_to_thrift_backend(self, mock_thrift_backe
         for req_args in req_args_combinations:
             req_args = {k: v for k, v in req_args.items() if v != "NOT_SET"}
             with self.subTest(req_args=req_args):
-                mock_thrift_backend = ThriftBackendMockFactory.mock_thrift_backend_with_retry_policy()
+                mock_thrift_backend = Mock()
 
                 cursor = client.Cursor(Mock(), mock_thrift_backend)
                 cursor.schemas(**req_args)
@@ -417,7 +416,7 @@ def test_get_tables_parameters_passed_to_thrift_backend(self, mock_thrift_backen
         for req_args in req_args_combinations:
             req_args = {k: v for k, v in req_args.items() if v != "NOT_SET"}
             with self.subTest(req_args=req_args):
-                mock_thrift_backend = ThriftBackendMockFactory.mock_thrift_backend_with_retry_policy()
+                mock_thrift_backend = Mock()
 
                 cursor = client.Cursor(Mock(), mock_thrift_backend)
                 cursor.tables(**req_args)
@@ -440,7 +439,7 @@ def test_get_columns_parameters_passed_to_thrift_backend(self, mock_thrift_backe
         for req_args in req_args_combinations:
             req_args = {k: v for k, v in req_args.items() if v != "NOT_SET"}
             with self.subTest(req_args=req_args):
-                mock_thrift_backend = ThriftBackendMockFactory.mock_thrift_backend_with_retry_policy()
+                mock_thrift_backend = Mock()
 
                 cursor = client.Cursor(Mock(), mock_thrift_backend)
                 cursor.columns(**req_args)