do an actual fix

Author: pedroslopez

airbyte_cdk/manifest_server/command_processor/processor.py

@@ -7,7 +7,6 @@
 )
 from fastapi import HTTPException
 
-import airbyte_cdk.sources.declarative.stream_slicers.declarative_partition_generator
 from airbyte_cdk.connector_builder.models import StreamRead
 from airbyte_cdk.connector_builder.test_reader import TestReader
 from airbyte_cdk.entrypoint import AirbyteEntrypoint
@@ -42,13 +41,6 @@ def test_read(
         """
         Test the read method of the source.
         """
-
-        # HACK: reset total_record_counter
-        # DeclarativePartition defines total_record_counter as a global variable, which keeps around the record count
-        # across multiple test_read calls, even if the source is different. This is a hack to reset the counter for
-        # each test_read call.
-        airbyte_cdk.sources.declarative.stream_slicers.declarative_partition_generator.total_record_counter = 0
-
         test_read_handler = TestReader(
             max_pages_per_slice=page_limit,
             max_slices=slice_limit,

airbyte_cdk/sources/declarative/stream_slicers/declarative_partition_generator.py

@@ -14,10 +14,21 @@
 from airbyte_cdk.sources.types import Record, StreamSlice
 from airbyte_cdk.utils.slice_hasher import SliceHasher
 
+
 # For Connector Builder test read operations, we track the total number of records
-# read for the stream at the global level so that we can stop reading early if we
-# exceed the record limit
-total_record_counter = 0
+# read for the stream so that we can stop reading early if we exceed the record limit.
+class RecordCounter:
+    def __init__(self):
+        self.total_record_counter = 0
+
+    def increment(self):
+        self.total_record_counter += 1
+
+    def get_total_records(self) -> int:
+        return self.total_record_counter
+
+    def reset(self):
+        self.total_record_counter = 0
 
 
 class SchemaLoaderCachingDecorator(SchemaLoader):
@@ -51,6 +62,7 @@ def __init__(
         self._retriever = retriever
         self._message_repository = message_repository
         self._max_records_limit = max_records_limit
+        self._record_counter = RecordCounter()
 
     def create(self, stream_slice: StreamSlice) -> Partition:
         return DeclarativePartition(
@@ -60,6 +72,7 @@ def create(self, stream_slice: StreamSlice) -> Partition:
             message_repository=self._message_repository,
             max_records_limit=self._max_records_limit,
             stream_slice=stream_slice,
+            record_counter=self._record_counter,
         )
 
 
@@ -72,6 +85,7 @@ def __init__(
         message_repository: MessageRepository,
         max_records_limit: Optional[int],
         stream_slice: StreamSlice,
+        record_counter: RecordCounter,
     ):
         self._stream_name = stream_name
         self._schema_loader = schema_loader
@@ -80,17 +94,17 @@ def __init__(
         self._max_records_limit = max_records_limit
         self._stream_slice = stream_slice
         self._hash = SliceHasher.hash(self._stream_name, self._stream_slice)
+        self._record_counter = record_counter
 
     def read(self) -> Iterable[Record]:
         if self._max_records_limit is not None:
-            global total_record_counter
-            if total_record_counter >= self._max_records_limit:
+            if self._record_counter.get_total_records() >= self._max_records_limit:
                 return
         for stream_data in self._retriever.read_records(
             self._schema_loader.get_json_schema(), self._stream_slice
         ):
             if self._max_records_limit is not None:
-                if total_record_counter >= self._max_records_limit:
+                if self._record_counter.get_total_records() >= self._max_records_limit:
                     break
 
             if isinstance(stream_data, Mapping):
@@ -108,7 +122,7 @@ def read(self) -> Iterable[Record]:
                 self._message_repository.emit_message(stream_data)
 
             if self._max_records_limit is not None:
-                total_record_counter += 1
+                self._record_counter.increment()
 
     def to_slice(self) -> Optional[Mapping[str, Any]]:
         return self._stream_slice

unit_tests/manifest_server/command_processor/test_processor.py

@@ -308,57 +308,3 @@ def test_discover_with_trace_error(self, command_processor, sample_config):
             # Verify exception is raised
             with pytest.raises(HTTPException):
                 command_processor.discover(sample_config)
-
-    def test_test_read_resets_global_record_counter(
-        self, command_processor, sample_config, sample_catalog
-    ):
-        """Test that test_read resets the global total_record_counter between calls."""
-        from airbyte_cdk.sources.declarative import stream_slicers
-
-        # Mock the TestReader
-        with patch(
-            "airbyte_cdk.manifest_server.command_processor.processor.TestReader"
-        ) as mock_test_reader_class:
-            mock_test_reader_instance = Mock()
-            mock_test_reader_class.return_value = mock_test_reader_instance
-            mock_stream_read = Mock()
-            mock_test_reader_instance.run_test_read.return_value = mock_stream_read
-
-            # Set initial counter value to simulate previous test_read execution
-            stream_slicers.declarative_partition_generator.total_record_counter = 100
-
-            # Execute test_read
-            result1 = command_processor.test_read(
-                config=sample_config,
-                catalog=sample_catalog,
-                state=[],
-                record_limit=50,
-                page_limit=3,
-                slice_limit=7,
-            )
-
-            # Verify counter was reset to 0 before the test_read
-            assert stream_slicers.declarative_partition_generator.total_record_counter == 0
-
-            # Set counter again to simulate state from first call
-            stream_slicers.declarative_partition_generator.total_record_counter = 200
-
-            # Execute another test_read
-            result2 = command_processor.test_read(
-                config=sample_config,
-                catalog=sample_catalog,
-                state=[],
-                record_limit=25,
-                page_limit=2,
-                slice_limit=5,
-            )
-
-            # Verify counter was reset again
-            assert stream_slicers.declarative_partition_generator.total_record_counter == 0
-
-            # Verify both calls returned the expected results
-            assert result1 == mock_stream_read
-            assert result2 == mock_stream_read
-
-            # Verify TestReader was called twice
-            assert mock_test_reader_class.call_count == 2

unit_tests/sources/declarative/stream_slicers/test_declarative_partition_generator.py

@@ -4,8 +4,6 @@
 from unittest import TestCase
 from unittest.mock import Mock
 
-# This allows for the global total_record_counter to be reset between tests
-import airbyte_cdk.sources.declarative.stream_slicers.declarative_partition_generator as declarative_partition_generator
 from airbyte_cdk.models import AirbyteLogMessage, AirbyteMessage, Level, Type
 from airbyte_cdk.sources.declarative.retrievers import Retriever
 from airbyte_cdk.sources.declarative.schema import InlineSchemaLoader
@@ -35,7 +33,7 @@ class StreamSlicerPartitionGeneratorTest(TestCase):
     def test_given_multiple_slices_partition_generator_uses_the_same_retriever(self) -> None:
         retriever = self._mock_retriever([])
         message_repository = Mock(spec=MessageRepository)
-        partition_factory = declarative_partition_generator.DeclarativePartitionFactory(
+        partition_factory = DeclarativePartitionFactory(
             _STREAM_NAME,
             _SCHEMA_LOADER,
             retriever,
@@ -50,7 +48,7 @@ def test_given_multiple_slices_partition_generator_uses_the_same_retriever(self)
     def test_given_a_mapping_when_read_then_yield_record(self) -> None:
         retriever = self._mock_retriever([_A_RECORD])
         message_repository = Mock(spec=MessageRepository)
-        partition_factory = declarative_partition_generator.DeclarativePartitionFactory(
+        partition_factory = DeclarativePartitionFactory(
             _STREAM_NAME,
             _SCHEMA_LOADER,
             retriever,
@@ -68,7 +66,7 @@ def test_given_a_mapping_when_read_then_yield_record(self) -> None:
     def test_given_not_a_record_when_read_then_send_to_message_repository(self) -> None:
         retriever = self._mock_retriever([_AIRBYTE_LOG_MESSAGE])
         message_repository = Mock(spec=MessageRepository)
-        partition_factory = declarative_partition_generator.DeclarativePartitionFactory(
+        partition_factory = DeclarativePartitionFactory(
             _STREAM_NAME,
             _SCHEMA_LOADER,
             retriever,
@@ -80,8 +78,6 @@ def test_given_not_a_record_when_read_then_send_to_message_repository(self) -> N
         message_repository.emit_message.assert_called_once_with(_AIRBYTE_LOG_MESSAGE)
 
     def test_max_records_reached_stops_reading(self) -> None:
-        declarative_partition_generator.total_record_counter = 0
-
         expected_records = [
             Record(data={"id": 1, "name": "Max"}, stream_name="stream_name"),
             Record(data={"id": 1, "name": "Oscar"}, stream_name="stream_name"),
@@ -97,7 +93,7 @@ def test_max_records_reached_stops_reading(self) -> None:
 
         retriever = self._mock_retriever(mock_records)
         message_repository = Mock(spec=MessageRepository)
-        partition_factory = declarative_partition_generator.DeclarativePartitionFactory(
+        partition_factory = DeclarativePartitionFactory(
             _STREAM_NAME,
             _SCHEMA_LOADER,
             retriever,
@@ -113,8 +109,6 @@ def test_max_records_reached_stops_reading(self) -> None:
         assert actual_records == expected_records
 
     def test_max_records_reached_on_previous_partition(self) -> None:
-        declarative_partition_generator.total_record_counter = 0
-
         expected_records = [
             Record(data={"id": 1, "name": "Max"}, stream_name="stream_name"),
             Record(data={"id": 1, "name": "Oscar"}, stream_name="stream_name"),
@@ -128,7 +122,7 @@ def test_max_records_reached_on_previous_partition(self) -> None:
 
         retriever = self._mock_retriever(mock_records)
         message_repository = Mock(spec=MessageRepository)
-        partition_factory = declarative_partition_generator.DeclarativePartitionFactory(
+        partition_factory = DeclarativePartitionFactory(
             _STREAM_NAME,
             _SCHEMA_LOADER,
             retriever,