Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@
from __future__ import annotations

import logging
import traceback
from typing import TYPE_CHECKING, Optional, Tuple

from airbyte_cdk import AirbyteTracedException
Expand Down Expand Up @@ -39,12 +38,18 @@ def check_availability( # type: ignore[override] # Signature doesn't match bas
"""
Perform a connection check for the stream (verify that we can list files from the stream).

Returns (True, None) if successful, otherwise (False, <error message>).
Returns `(True, None)` if successful, otherwise `(False, <error message>)`.

The returned reason is the actionable error string from the underlying
`CheckAvailabilityError`. Connector-raised `AirbyteTracedException`s are
re-raised so that their actionable `message` propagates to the platform
unchanged. The full traceback is still logged for debugging.
"""
try:
self._check_list_files(stream)
except CheckAvailabilityError:
return False, "".join(traceback.format_exc())
except CheckAvailabilityError as exc:
logger.exception("Stream availability check failed")
return False, str(exc)

return True, None

Expand Down Expand Up @@ -83,25 +88,33 @@ def check_availability_and_parsability(
# If the parser is set to not check parsability, we still want to check that we can open the file.
handle = stream.stream_reader.open_file(file, parser.file_read_mode, None, logger)
handle.close()
except AirbyteTracedException as ate:
raise ate
except CheckAvailabilityError:
return False, "".join(traceback.format_exc())
except AirbyteTracedException:
raise
except CheckAvailabilityError as exc:
logger.exception("Stream availability check failed")
return False, str(exc)

return True, None

def _check_list_files(self, stream: AbstractFileBasedStream) -> RemoteFile:
"""
Check that we can list files from the stream.

Returns the first file if successful, otherwise raises a CheckAvailabilityError.
Returns the first file if successful, otherwise raises a `CheckAvailabilityError`.

`AirbyteTracedException`s raised by the underlying stream reader are
re-raised unchanged so that connector-specific actionable messages
(e.g. invalid credentials, missing permissions) reach the platform
without being wrapped in a generic `ERROR_LISTING_FILES` reason.
"""
try:
file = next(iter(stream.get_files()))
except StopIteration:
raise CheckAvailabilityError(FileBasedSourceError.EMPTY_STREAM, stream=stream.name)
except CustomFileBasedException as exc:
raise CheckAvailabilityError(str(exc), stream=stream.name) from exc
except AirbyteTracedException:
raise
except Exception as exc:
raise CheckAvailabilityError(
FileBasedSourceError.ERROR_LISTING_FILES, stream=stream.name
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,9 @@
from datetime import datetime
from unittest.mock import Mock, PropertyMock

import pytest

from airbyte_cdk import AirbyteTracedException, FailureType
from airbyte_cdk.sources.file_based.availability_strategy.default_file_based_availability_strategy import (
DefaultFileBasedAvailabilityStrategy,
)
Expand Down Expand Up @@ -61,7 +64,9 @@ def test_given_file_extension_does_not_match_when_check_availability_and_parsabi

def test_not_available_given_no_files(self) -> None:
"""
If no files are returned, then the stream is not available.
If no files are returned, then the stream is not available, and the
reason is the actionable `EMPTY_STREAM` message rather than a Python
traceback.
"""
self._stream.get_files.return_value = []

Expand All @@ -71,6 +76,62 @@ def test_not_available_given_no_files(self) -> None:

assert not is_available
assert "No files were identified in the stream" in reason
assert "Traceback" not in reason
assert "raise CheckAvailabilityError" not in reason

def test_check_availability_returns_actionable_reason_when_no_files(self) -> None:
"""
`check_availability` (used by file-transfer / permissions-transfer modes)
must also return the actionable reason rather than a traceback string.
"""
self._stream.get_files.return_value = []

is_available, reason = self._strategy.check_availability(self._stream, Mock(), Mock())

assert not is_available
assert "No files were identified in the stream" in reason
assert "Traceback" not in reason
assert "raise CheckAvailabilityError" not in reason

def test_airbyte_traced_exception_from_stream_reader_propagates(self) -> None:
"""
When the underlying stream reader raises an `AirbyteTracedException`
(e.g. invalid credentials), the actionable message must propagate
unchanged instead of being wrapped in a generic `ERROR_LISTING_FILES`
reason.
"""
self._stream.get_files.side_effect = AirbyteTracedException(
internal_message="raw provider error",
message="Could not authenticate with Google drive. Please check your credentials.",
failure_type=FailureType.config_error,
)

with pytest.raises(AirbyteTracedException) as exc_info:
self._strategy.check_availability_and_parsability(self._stream, Mock(), Mock())

assert (
exc_info.value.message
== "Could not authenticate with Google drive. Please check your credentials."
)
assert exc_info.value.failure_type == FailureType.config_error

def test_airbyte_traced_exception_propagates_in_check_availability(self) -> None:
"""
Same propagation guarantee as above for `check_availability`.
"""
self._stream.get_files.side_effect = AirbyteTracedException(
internal_message="raw provider error",
message="Could not authenticate with Google drive. Please check your credentials.",
failure_type=FailureType.config_error,
)

with pytest.raises(AirbyteTracedException) as exc_info:
self._strategy.check_availability(self._stream, Mock(), Mock())

assert (
exc_info.value.message
== "Could not authenticate with Google drive. Please check your credentials."
)

def test_parse_records_is_not_called_with_parser_max_n_files_for_parsability_set(self) -> None:
"""
Expand Down
Loading