🐛 Python CDK: fix StopIteration error for check_availability (#20429

)
airbytehq · Dec 13, 2022 · 4e9b014 · 4e9b014
1 parent 81341d0
commit 4e9b014
Show file tree

Hide file tree

Showing 4 changed files with 38 additions and 5 deletions.
diff --git a/airbyte-cdk/python/CHANGELOG.md b/airbyte-cdk/python/CHANGELOG.md
@@ -1,5 +1,8 @@
 # Changelog
 
+## 0.13.3
+Fixed `StopIteration` exception for empty streams while `check_availability` runs.
+
 ## 0.13.2
 Low-code: Enable low-code CDK users to specify schema inline in the manifest
 

diff --git a/airbyte-cdk/python/airbyte_cdk/sources/utils/stream_helpers.py b/airbyte-cdk/python/airbyte_cdk/sources/utils/stream_helpers.py
@@ -17,10 +17,13 @@ def get_first_record(self, stream: Stream) -> StreamData:
         :param stream: stream
         :return: StreamData containing the first record in the stream
         """
-        # Some streams need a stream slice to read records (e.g. if they have a SubstreamSlicer)
-        stream_slice = self.get_stream_slice(stream)
-        records = stream.read_records(sync_mode=SyncMode.full_refresh, stream_slice=stream_slice)
-        next(records)
+        try:
+            # Some streams need a stream slice to read records (e.g. if they have a SubstreamSlicer)
+            stream_slice = self.get_stream_slice(stream)
+            records = stream.read_records(sync_mode=SyncMode.full_refresh, stream_slice=stream_slice)
+            return next(records)
+        except StopIteration:
+            return {}
 
     @staticmethod
     def get_stream_slice(stream: Stream) -> Optional[Mapping[str, Any]]:

diff --git a/airbyte-cdk/python/setup.py b/airbyte-cdk/python/setup.py
@@ -15,7 +15,7 @@
 
 setup(
     name="airbyte-cdk",
-    version="0.13.2",
+    version="0.13.3",
     description="A framework for writing Airbyte Connectors.",
     long_description=README,
     long_description_content_type="text/markdown",

diff --git a/airbyte-cdk/python/unit_tests/sources/streams/http/test_availability_strategy.py b/airbyte-cdk/python/unit_tests/sources/streams/http/test_availability_strategy.py
@@ -139,3 +139,30 @@ def test_send_handles_retries_when_checking_availability(mocker, caplog):
     assert mock_send.call_count == 3
     for message in ["Caught retryable error", "Response Code: 429", "Response Code: 503"]:
         assert message in caplog.text
+
+
+def test_http_availability_strategy_on_empty_stream(mocker):
+    mocker.patch.multiple(HttpStream, __abstractmethods__=set())
+    mocker.patch.multiple(Stream, __abstractmethods__=set())
+
+    class MockEmptyStream(mocker.MagicMock, HttpStream):
+        page_size = None
+        get_json_schema = mocker.MagicMock()
+
+        def __init__(self, *args, **kvargs):
+            mocker.MagicMock.__init__(self)
+            self.read_records = mocker.MagicMock()
+
+    empty_stream = MockEmptyStream()
+    assert isinstance(empty_stream, HttpStream)
+
+    assert isinstance(empty_stream.availability_strategy, HttpAvailabilityStrategy)
+
+    # Generator should have no values to generate
+    empty_stream.read_records.return_value = iter([])
+
+    logger = logging.getLogger("airbyte.test-source")
+    stream_is_available, _ = empty_stream.check_availability(logger)
+
+    assert stream_is_available
+    assert empty_stream.read_records.called