remove deepcopy

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

Author: varun-edachali-dbx

src/databricks/sql/backend/sea/utils/filters.py

@@ -8,7 +8,6 @@
 
 import io
 import logging
-from copy import deepcopy
 from typing import (
     List,
     Optional,
@@ -62,7 +61,7 @@ def _create_execute_response(result_set: SeaResultSet) -> ExecuteResponse:
         )
 
     @staticmethod
-    def _create_filtered_manifest(result_set: SeaResultSet, new_row_count: int):
+    def _update_manifest(result_set: SeaResultSet, new_row_count: int):
         """
         Create a copy of the manifest with updated row count.
 
@@ -73,7 +72,7 @@ def _create_filtered_manifest(result_set: SeaResultSet, new_row_count: int):
         Returns:
             Updated manifest copy
         """
-        filtered_manifest = deepcopy(result_set.manifest)
+        filtered_manifest = result_set.manifest
         filtered_manifest.total_row_count = new_row_count
         return filtered_manifest
 
@@ -97,9 +96,7 @@ def _create_filtered_result_set(
         from databricks.sql.backend.sea.result_set import SeaResultSet
 
         execute_response = ResultSetFilter._create_execute_response(result_set)
-        filtered_manifest = ResultSetFilter._create_filtered_manifest(
-            result_set, row_count
-        )
+        filtered_manifest = ResultSetFilter._update_manifest(result_set, row_count)
 
         return SeaResultSet(
             connection=result_set.connection,

tests/e2e/test_concurrent_telemetry.py

@@ -6,9 +6,13 @@
 import pytest
 
 from databricks.sql.telemetry.models.enums import StatementType
-from databricks.sql.telemetry.telemetry_client import TelemetryClient, TelemetryClientFactory
+from databricks.sql.telemetry.telemetry_client import (
+    TelemetryClient,
+    TelemetryClientFactory,
+)
 from tests.e2e.test_driver import PySQLPytestTestCase
 
+
 def run_in_threads(target, num_threads, pass_index=False):
     """Helper to run target function in multiple threads."""
     threads = [
@@ -22,7 +26,6 @@ def run_in_threads(target, num_threads, pass_index=False):
 
 
 class TestE2ETelemetry(PySQLPytestTestCase):
-    
     @pytest.fixture(autouse=True)
     def telemetry_setup_teardown(self):
         """
@@ -31,7 +34,7 @@ def telemetry_setup_teardown(self):
         this robust and automatic.
         """
         try:
-            yield 
+            yield
         finally:
             if TelemetryClientFactory._executor:
                 TelemetryClientFactory._executor.shutdown(wait=True)
@@ -68,20 +71,25 @@ def callback_wrapper(self_client, future, sent_count):
                 captured_futures.append(future)
             original_callback(self_client, future, sent_count)
 
-        with patch.object(TelemetryClient, "_send_telemetry", send_telemetry_wrapper), \
-             patch.object(TelemetryClient, "_telemetry_request_callback", callback_wrapper):
+        with patch.object(
+            TelemetryClient, "_send_telemetry", send_telemetry_wrapper
+        ), patch.object(
+            TelemetryClient, "_telemetry_request_callback", callback_wrapper
+        ):
 
             def execute_query_worker(thread_id):
                 """Each thread creates a connection and executes a query."""
 
                 time.sleep(random.uniform(0, 0.05))
-                
-                with self.connection(extra_params={"force_enable_telemetry": True}) as conn:
+
+                with self.connection(
+                    extra_params={"force_enable_telemetry": True}
+                ) as conn:
                     # Capture the session ID from the connection before executing the query
                     session_id_hex = conn.get_session_id_hex()
                     with capture_lock:
                         captured_session_ids.append(session_id_hex)
-                    
+
                     with conn.cursor() as cursor:
                         cursor.execute(f"SELECT {thread_id}")
                         # Capture the statement ID after executing the query
@@ -97,7 +105,10 @@ def execute_query_worker(thread_id):
             start_time = time.time()
             expected_event_count = num_threads
 
-            while len(captured_futures) < expected_event_count and time.time() - start_time < timeout_seconds:
+            while (
+                len(captured_futures) < expected_event_count
+                and time.time() - start_time < timeout_seconds
+            ):
                 time.sleep(0.1)
 
             done, not_done = wait(captured_futures, timeout=timeout_seconds)
@@ -115,30 +126,37 @@ def execute_query_worker(thread_id):
 
             assert not captured_exceptions
             assert len(captured_responses) > 0
-            
+
             total_successful_events = 0
             for response in captured_responses:
                 assert "errors" not in response or not response["errors"]
                 if "numProtoSuccess" in response:
                     total_successful_events += response["numProtoSuccess"]
             assert total_successful_events == num_threads * 2
 
-            assert len(captured_telemetry) == num_threads * 2  # 2 events per thread (initial_telemetry_log, latency_log (execute))
+            assert (
+                len(captured_telemetry) == num_threads * 2
+            )  # 2 events per thread (initial_telemetry_log, latency_log (execute))
             assert len(captured_session_ids) == num_threads  # One session ID per thread
-            assert len(captured_statement_ids) == num_threads  # One statement ID per thread (per query)
+            assert (
+                len(captured_statement_ids) == num_threads
+            )  # One statement ID per thread (per query)
 
             # Separate initial logs from latency logs
             initial_logs = [
-                e for e in captured_telemetry
+                e
+                for e in captured_telemetry
                 if e.entry.sql_driver_log.operation_latency_ms is None
                 and e.entry.sql_driver_log.driver_connection_params is not None
                 and e.entry.sql_driver_log.system_configuration is not None
             ]
             latency_logs = [
-                e for e in captured_telemetry
-                if e.entry.sql_driver_log.operation_latency_ms is not None 
-                and e.entry.sql_driver_log.sql_statement_id is not None 
-                and e.entry.sql_driver_log.sql_operation.statement_type == StatementType.QUERY
+                e
+                for e in captured_telemetry
+                if e.entry.sql_driver_log.operation_latency_ms is not None
+                and e.entry.sql_driver_log.sql_statement_id is not None
+                and e.entry.sql_driver_log.sql_operation.statement_type
+                == StatementType.QUERY
             ]
 
             # Verify counts
@@ -171,4 +189,4 @@ def execute_query_worker(thread_id):
             for event in latency_logs:
                 log = event.entry.sql_driver_log
                 assert log.sql_statement_id in captured_statement_ids
-                assert log.session_id in captured_session_ids
+                assert log.session_id in captured_session_ids

tests/unit/test_telemetry.py

@@ -30,7 +30,7 @@ def mock_telemetry_client():
         auth_provider=auth_provider,
         host_url="test-host.com",
         executor=executor,
-        batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE
+        batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE,
     )
 
 
@@ -215,7 +215,7 @@ def test_client_lifecycle_flow(self):
             session_id_hex=session_id_hex,
             auth_provider=auth_provider,
             host_url="test-host.com",
-            batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE
+            batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE,
         )
 
         client = TelemetryClientFactory.get_telemetry_client(session_id_hex)
@@ -240,7 +240,7 @@ def test_disabled_telemetry_flow(self):
             session_id_hex=session_id_hex,
             auth_provider=None,
             host_url="test-host.com",
-            batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE
+            batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE,
         )
 
         client = TelemetryClientFactory.get_telemetry_client(session_id_hex)
@@ -260,7 +260,7 @@ def test_factory_error_handling(self):
                 session_id_hex=session_id,
                 auth_provider=AccessTokenAuthProvider("token"),
                 host_url="test-host.com",
-                batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE
+                batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE,
             )
 
         # Should fall back to NoopTelemetryClient
@@ -279,7 +279,7 @@ def test_factory_shutdown_flow(self):
                 session_id_hex=session,
                 auth_provider=AccessTokenAuthProvider("token"),
                 host_url="test-host.com",
-                batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE
+                batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE,
             )
 
         # Factory should be initialized
@@ -342,9 +342,7 @@ def _mock_ff_response(self, mock_requests_get, enabled: bool):
         mock_requests_get.return_value = mock_response
 
     @patch("databricks.sql.common.feature_flag.requests.get")
-    def test_telemetry_enabled_when_flag_is_true(
-        self, mock_requests_get, MockSession
-    ):
+    def test_telemetry_enabled_when_flag_is_true(self, mock_requests_get, MockSession):
         """Telemetry should be ON when enable_telemetry=True and server flag is 'true'."""
         self._mock_ff_response(mock_requests_get, enabled=True)
         mock_session_instance = MockSession.return_value
@@ -405,4 +403,4 @@ def test_telemetry_disabled_when_flag_request_fails(
         assert conn.telemetry_enabled is False
         mock_requests_get.assert_called_once()
         client = TelemetryClientFactory.get_telemetry_client("test-session-ff-fail")
-        assert isinstance(client, NoopTelemetryClient)
+        assert isinstance(client, NoopTelemetryClient)

tests/unit/test_telemetry_retry.py

@@ -51,7 +51,7 @@ def get_client(self, session_id, num_retries=3):
             session_id_hex=session_id,
             auth_provider=None,
             host_url="test.databricks.com",
-            batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE
+            batch_size=TelemetryClientFactory.DEFAULT_BATCH_SIZE,
         )
         client = TelemetryClientFactory.get_telemetry_client(session_id)