Skip to content

Commit

Permalink
🐛 Python CDK: fix StopIteration error for check_availability (#20429
Browse files Browse the repository at this point in the history
)
  • Loading branch information
bazarnov authored Dec 13, 2022
1 parent 81341d0 commit 4e9b014
Show file tree
Hide file tree
Showing 4 changed files with 38 additions and 5 deletions.
3 changes: 3 additions & 0 deletions airbyte-cdk/python/CHANGELOG.md
Original file line number Diff line number Diff line change
@@ -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

Expand Down
11 changes: 7 additions & 4 deletions airbyte-cdk/python/airbyte_cdk/sources/utils/stream_helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]]:
Expand Down
2 changes: 1 addition & 1 deletion airbyte-cdk/python/setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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

0 comments on commit 4e9b014

Please sign in to comment.