formatting

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

Author: saishreeeee

src/databricks/sql/client.py

@@ -1650,4 +1650,3 @@ def map_col_type(type_):
             (column.name, map_col_type(column.datatype), None, None, None, None, None)
             for column in table_schema_message.columns
         ]
-

src/databricks/sql/telemetry/latency_logger.py

@@ -13,16 +13,27 @@
 class TelemetryExtractor:
     def __init__(self, obj):
         self._obj = obj
-    
+
     def __getattr__(self, name):
         return getattr(self._obj, name)
-    
-    def get_session_id_hex(self): pass
-    def get_statement_id(self): pass
-    def get_statement_type(self): pass
-    def get_is_compressed(self): pass
-    def get_execution_result(self): pass
-    def get_retry_count(self): pass
+
+    def get_session_id_hex(self):
+        pass
+
+    def get_statement_id(self):
+        pass
+
+    def get_statement_type(self):
+        pass
+
+    def get_is_compressed(self):
+        pass
+
+    def get_execution_result(self):
+        pass
+
+    def get_retry_count(self):
+        pass
 
 
 class CursorExtractor(TelemetryExtractor):
@@ -58,8 +69,8 @@ def get_retry_count(self) -> int:
     def get_statement_type(self: str) -> StatementType:
         # TODO: Implement this
         return StatementType.SQL
-    
-    
+
+
 class ResultSetExtractor(TelemetryExtractor):
     def get_statement_id(self) -> Optional[str]:
         if self.command_id:
@@ -95,9 +106,9 @@ def get_retry_count(self) -> int:
 
 
 def get_extractor(obj):
-    if obj.__class__.__name__ == 'Cursor':
+    if obj.__class__.__name__ == "Cursor":
         return CursorExtractor(obj)
-    elif obj.__class__.__name__ == 'ResultSet':
+    elif obj.__class__.__name__ == "ResultSet":
         return ResultSetExtractor(obj)
     else:
         return TelemetryExtractor(obj)

src/databricks/sql/telemetry/telemetry_client.py

@@ -140,9 +140,7 @@ def export_initial_telemetry_log(self, driver_connection_params, user_agent):
     def export_failure_log(self, error_name, error_message):
         pass
 
-    def export_latency_log(
-        self, latency_ms, sql_execution_event, sql_statement_id=None
-    ):
+    def export_latency_log(self, latency_ms, sql_execution_event, sql_statement_id):
         pass
 
     def close(self):
@@ -158,7 +156,6 @@ class TelemetryClient(BaseTelemetryClient):
     # Telemetry endpoint paths
     TELEMETRY_AUTHENTICATED_PATH = "/telemetry-ext"
     TELEMETRY_UNAUTHENTICATED_PATH = "/telemetry-unauth"
-    DEFAULT_BATCH_SIZE = 10
 
     def __init__(
         self,
@@ -170,7 +167,7 @@ def __init__(
     ):
         logger.debug("Initializing TelemetryClient for connection: %s", session_id_hex)
         self._telemetry_enabled = telemetry_enabled
-        self._batch_size = self.DEFAULT_BATCH_SIZE  # TODO: Decide on batch size
+        self._batch_size = 10  # TODO: Decide on batch size
         self._session_id_hex = session_id_hex
         self._auth_provider = auth_provider
         self._user_agent = None

tests/unit/test_telemetry.py

@@ -229,16 +229,17 @@ def test_export_event(self, telemetry_client_setup):
         client = telemetry_client_setup["client"]
         client._flush = MagicMock()
 
-        for i in range(TelemetryClient.DEFAULT_BATCH_SIZE-1):
+        for i in range(5):
             client._export_event(f"event-{i}")
 
         client._flush.assert_not_called()
-        assert len(client._events_batch) == TelemetryClient.DEFAULT_BATCH_SIZE - 1
+        assert len(client._events_batch) == 5
 
-        # Add one more event to reach batch size (this will trigger flush)
-        client._export_event(f"event-{TelemetryClient.DEFAULT_BATCH_SIZE - 1}")
+        for i in range(5):
+            client._export_event(f"event-{i}")
 
         client._flush.assert_called_once()
+        assert len(client._events_batch) == 10
 
     @patch("requests.post")
     def test_send_telemetry_authenticated(self, mock_post, telemetry_client_setup):