airbyte-cdk 6.45.4.post49.dev14495925594__py3-none-any.whl → 6.45.4.post72.dev14497997772__py3-none-any.whl

airbyte_cdk/sources/file_based/stream/default_file_based_stream.py

@@ -11,7 +11,7 @@ from functools import cache
 from os import path
 from typing import Any, Dict, Iterable, List, Mapping, MutableMapping, Optional, Set, Tuple, Union
 
-from airbyte_cdk.models import AirbyteLogMessage, AirbyteMessage, AirbyteStream, FailureType, Level
+from airbyte_cdk.models import AirbyteLogMessage, AirbyteMessage, FailureType, Level
 from airbyte_cdk.models import Type as MessageType
 from airbyte_cdk.sources.file_based.config.file_based_stream_config import PrimaryKeyType
 from airbyte_cdk.sources.file_based.exceptions import (
@@ -56,7 +56,6 @@ class DefaultFileBasedStream(AbstractFileBasedStream, IncrementalMixin):
     airbyte_columns = [ab_last_mod_col, ab_file_name_col]
     use_file_transfer = False
     preserve_directory_structure = True
-    _file_transfer = FileTransfer()
 
     def __init__(self, **kwargs: Any):
         if self.FILE_TRANSFER_KW in kwargs:
@@ -94,6 +93,21 @@ class DefaultFileBasedStream(AbstractFileBasedStream, IncrementalMixin):
             self.config
         )
 
+    def _filter_schema_invalid_properties(
+        self, configured_catalog_json_schema: Dict[str, Any]
+    ) -> Dict[str, Any]:
+        if self.use_file_transfer:
+            return {
+                "type": "object",
+                "properties": {
+                    "file_path": {"type": "string"},
+                    "file_size": {"type": "string"},
+                    self.ab_file_name_col: {"type": "string"},
+                },
+            }
+        else:
+            return super()._filter_schema_invalid_properties(configured_catalog_json_schema)
+
     def _duplicated_files_names(
         self, slices: List[dict[str, List[RemoteFile]]]
     ) -> List[dict[str, List[str]]]:
@@ -131,6 +145,14 @@ class DefaultFileBasedStream(AbstractFileBasedStream, IncrementalMixin):
         record[self.ab_file_name_col] = file.uri
         return record
 
+    def transform_record_for_file_transfer(
+        self, record: dict[str, Any], file: RemoteFile
+    ) -> dict[str, Any]:
+        # timstamp() returns a float representing the number of seconds since the unix epoch
+        record[self.modified] = int(file.last_modified.timestamp()) * 1000
+        record[self.source_file_url] = file.uri
+        return record
+
     def read_records_from_slice(self, stream_slice: StreamSlice) -> Iterable[AirbyteMessage]:
         """
         Yield all records from all remote files in `list_files_for_this_sync`.
@@ -151,13 +173,19 @@ class DefaultFileBasedStream(AbstractFileBasedStream, IncrementalMixin):
 
             try:
                 if self.use_file_transfer:
-                    for file_record_data, file_reference in self._file_transfer.upload(
-                        file=file, stream_reader=self.stream_reader, logger=self.logger
+                    self.logger.info(f"{self.name}: {file} file-based syncing")
+                    # todo: complete here the code to not rely on local parser
+                    file_transfer = FileTransfer()
+                    for record in file_transfer.get_file(
+                        self.config, file, self.stream_reader, self.logger
                     ):
+                        line_no += 1
+                        if not self.record_passes_validation_policy(record):
+                            n_skipped += 1
+                            continue
+                        record = self.transform_record_for_file_transfer(record, file)
                         yield stream_data_to_airbyte_message(
-                            self.name,
-                            file_record_data.dict(exclude_none=True),
-                            file_reference=file_reference,
+                            self.name, record, is_file_transfer_message=True
                         )
                 else:
                     for record in parser.parse_records(
@@ -231,8 +259,6 @@ class DefaultFileBasedStream(AbstractFileBasedStream, IncrementalMixin):
 
     @cache
     def get_json_schema(self) -> JsonSchema:
-        if self.use_file_transfer:
-            return file_transfer_schema
         extra_fields = {
             self.ab_last_mod_col: {"type": "string"},
             self.ab_file_name_col: {"type": "string"},
@@ -256,7 +282,9 @@ class DefaultFileBasedStream(AbstractFileBasedStream, IncrementalMixin):
             return {"type": "object", "properties": {**extra_fields, **schema["properties"]}}
 
     def _get_raw_json_schema(self) -> JsonSchema:
-        if self.config.input_schema:
+        if self.use_file_transfer:
+            return file_transfer_schema
+        elif self.config.input_schema:
             return self.config.get_input_schema()  # type: ignore
         elif self.config.schemaless:
             return schemaless_schema
@@ -313,11 +341,6 @@ class DefaultFileBasedStream(AbstractFileBasedStream, IncrementalMixin):
             self.config.globs or [], self.config.legacy_prefix, self.logger
         )
 
-    def as_airbyte_stream(self) -> AirbyteStream:
-        file_stream = super().as_airbyte_stream()
-        file_stream.is_file_based = self.use_file_transfer
-        return file_stream
-
     def infer_schema(self, files: List[RemoteFile]) -> Mapping[str, Any]:
         loop = asyncio.get_event_loop()
         schema = loop.run_until_complete(self._infer_schema(files))

airbyte_cdk/sources/file_based/stream/permissions_file_based_stream.py

@@ -61,7 +61,9 @@ class PermissionsFileBasedStream(DefaultFileBasedStream):
                 permissions_record = self.transform_record(
                     permissions_record, file, file_datetime_string
                 )
-                yield stream_data_to_airbyte_message(self.name, permissions_record)
+                yield stream_data_to_airbyte_message(
+                    self.name, permissions_record, is_file_transfer_message=False
+                )
             except Exception as e:
                 self.logger.error(f"Failed to retrieve permissions for file {file.uri}: {str(e)}")
                 yield AirbyteMessage(

airbyte_cdk/sources/streams/concurrent/default_stream.py

@@ -29,7 +29,6 @@ class DefaultStream(AbstractStream):
         logger: Logger,
         cursor: Cursor,
         namespace: Optional[str] = None,
-        supports_file_transfer: bool = False,
     ) -> None:
         self._stream_partition_generator = partition_generator
         self._name = name
@@ -40,7 +39,6 @@ class DefaultStream(AbstractStream):
         self._logger = logger
         self._cursor = cursor
         self._namespace = namespace
-        self._supports_file_transfer = supports_file_transfer
 
     def generate_partitions(self) -> Iterable[Partition]:
         yield from self._stream_partition_generator.generate()
@@ -70,7 +68,6 @@ class DefaultStream(AbstractStream):
             json_schema=dict(self._json_schema),
             supported_sync_modes=[SyncMode.full_refresh],
             is_resumable=False,
-            is_file_based=self._supports_file_transfer,
         )
 
         if self._namespace:

airbyte_cdk/sources/types.py

@@ -6,7 +6,6 @@ from __future__ import annotations
 
 from typing import Any, ItemsView, Iterator, KeysView, List, Mapping, Optional, ValuesView
 
-from airbyte_cdk.models import AirbyteRecordMessageFileReference
 from airbyte_cdk.utils.slice_hasher import SliceHasher
 
 # A FieldPointer designates a path to a field inside a mapping. For example, retrieving ["k1", "k1.2"] in the object {"k1" :{"k1.2":
@@ -24,12 +23,12 @@ class Record(Mapping[str, Any]):
         data: Mapping[str, Any],
         stream_name: str,
         associated_slice: Optional[StreamSlice] = None,
-        file_reference: Optional[AirbyteRecordMessageFileReference] = None,
+        is_file_transfer_message: bool = False,
     ):
         self._data = data
         self._associated_slice = associated_slice
         self.stream_name = stream_name
-        self._file_reference = file_reference
+        self.is_file_transfer_message = is_file_transfer_message
 
     @property
     def data(self) -> Mapping[str, Any]:
@@ -39,14 +38,6 @@ class Record(Mapping[str, Any]):
     def associated_slice(self) -> Optional[StreamSlice]:
         return self._associated_slice
 
-    @property
-    def file_reference(self) -> AirbyteRecordMessageFileReference:
-        return self._file_reference
-
-    @file_reference.setter
-    def file_reference(self, value: AirbyteRecordMessageFileReference) -> None:
-        self._file_reference = value
-
     def __repr__(self) -> str:
         return repr(self._data)
 

airbyte_cdk/sources/utils/record_helper.py

@@ -9,10 +9,10 @@ from airbyte_cdk.models import (
     AirbyteLogMessage,
     AirbyteMessage,
     AirbyteRecordMessage,
-    AirbyteRecordMessageFileReference,
     AirbyteTraceMessage,
 )
 from airbyte_cdk.models import Type as MessageType
+from airbyte_cdk.models.file_transfer_record_message import AirbyteFileTransferRecordMessage
 from airbyte_cdk.sources.streams.core import StreamData
 from airbyte_cdk.sources.utils.transform import TransformConfig, TypeTransformer
 
@@ -22,7 +22,7 @@ def stream_data_to_airbyte_message(
     data_or_message: StreamData,
     transformer: TypeTransformer = TypeTransformer(TransformConfig.NoTransform),
     schema: Optional[Mapping[str, Any]] = None,
-    file_reference: Optional[AirbyteRecordMessageFileReference] = None,
+    is_file_transfer_message: bool = False,
 ) -> AirbyteMessage:
     if schema is None:
         schema = {}
@@ -36,12 +36,12 @@ def stream_data_to_airbyte_message(
             # taken unless configured. See
             # docs/connector-development/cdk-python/schemas.md for details.
             transformer.transform(data, schema)
-            message = AirbyteRecordMessage(
-                stream=stream_name,
-                data=data,
-                emitted_at=now_millis,
-                file_reference=file_reference,
-            )
+            if is_file_transfer_message:
+                message = AirbyteFileTransferRecordMessage(
+                    stream=stream_name, file=data, emitted_at=now_millis, data={}
+                )
+            else:
+                message = AirbyteRecordMessage(stream=stream_name, data=data, emitted_at=now_millis)
             return AirbyteMessage(type=MessageType.RECORD, record=message)
         case AirbyteTraceMessage():
             return AirbyteMessage(type=MessageType.TRACE, trace=data_or_message)

airbyte_cdk/test/declarative/__init__.py

@@ -0,0 +1,6 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+"""Declarative tests framework.
+
+This module provides fixtures and utilities for testing Airbyte sources and destinations
+in a declarative way.
+"""

airbyte_cdk/test/declarative/models/__init__.py

@@ -0,0 +1,7 @@
+from airbyte_cdk.test.declarative.models.scenario import (
+    ConnectorTestScenario,
+)
+
+__all__ = [
+    "ConnectorTestScenario",
+]

airbyte_cdk/test/declarative/models/scenario.py

@@ -0,0 +1,74 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+"""Run acceptance tests in PyTest.
+
+These tests leverage the same `acceptance-test-config.yml` configuration files as the
+acceptance tests in CAT, but they run in PyTest instead of CAT. This allows us to run
+the acceptance tests in the same local environment as we are developing in, speeding
+up iteration cycles.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any, Literal, cast
+
+import yaml
+from pydantic import BaseModel
+
+
+class ConnectorTestScenario(BaseModel):
+    """Acceptance test scenario, as a Pydantic model.
+
+    This class represents an acceptance test scenario, which is a single test case
+    that can be run against a connector. It is used to deserialize and validate the
+    acceptance test configuration file.
+    """
+
+    class AcceptanceTestExpectRecords(BaseModel):
+        path: Path
+        exact_order: bool = False
+
+    class AcceptanceTestFileTypes(BaseModel):
+        skip_test: bool
+        bypass_reason: str
+
+    config_path: Path | None = None
+    config_dict: dict[str, Any] | None = None
+
+    id: str | None = None
+
+    configured_catalog_path: Path | None = None
+    timeout_seconds: int | None = None
+    expect_records: AcceptanceTestExpectRecords | None = None
+    file_types: AcceptanceTestFileTypes | None = None
+    status: Literal["succeed", "failed"] | None = None
+
+    def get_config_dict(self) -> dict[str, Any]:
+        """Return the config dictionary.
+
+        If a config dictionary has already been loaded, return it. Otherwise, load
+        the config file and return the dictionary.
+        """
+        if self.config_dict:
+            return self.config_dict
+
+        if self.config_path:
+            return cast(dict[str, Any], yaml.safe_load(self.config_path.read_text()))
+
+        raise ValueError("No config dictionary or path provided.")
+
+    @property
+    def expect_exception(self) -> bool:
+        return self.status and self.status == "failed" or False
+
+    @property
+    def instance_name(self) -> str:
+        return self.config_path.stem if self.config_path else "Unnamed Scenario"
+
+    def __str__(self) -> str:
+        if self.id:
+            return f"'{self.id}' Test Scenario"
+        if self.config_path:
+            return f"'{self.config_path.name}' Test Scenario"
+
+        return f"'{hash(self)}' Test Scenario"

airbyte_cdk/test/declarative/utils/job_runner.py

@@ -0,0 +1,159 @@
+# Copyright (c) 2025 Airbyte, Inc., all rights reserved.
+"""Job runner for Airbyte Standard Tests."""
+
+import logging
+import tempfile
+import uuid
+from dataclasses import asdict
+from pathlib import Path
+from typing import Any, Callable, Literal
+
+import orjson
+from typing_extensions import Protocol, runtime_checkable
+
+from airbyte_cdk.models import (
+    ConfiguredAirbyteCatalog,
+    Status,
+)
+from airbyte_cdk.test import entrypoint_wrapper
+from airbyte_cdk.test.declarative.models import (
+    ConnectorTestScenario,
+)
+
+
+def _errors_to_str(
+    entrypoint_output: entrypoint_wrapper.EntrypointOutput,
+) -> str:
+    """Convert errors from entrypoint output to a string."""
+    if not entrypoint_output.errors:
+        # If there are no errors, return an empty string.
+        return ""
+
+    return "\n" + "\n".join(
+        [
+            str(error.trace.error).replace(
+                "\\n",
+                "\n",
+            )
+            for error in entrypoint_output.errors
+            if error.trace
+        ],
+    )
+
+
+@runtime_checkable
+class IConnector(Protocol):
+    """A connector that can be run in a test scenario.
+
+    Note: We currently use 'spec' to determine if we have a connector object.
+    In the future, it would be preferred to leverage a 'launch' method instead,
+    directly on the connector (which doesn't yet exist).
+    """
+
+    def spec(self, logger: logging.Logger) -> Any:
+        """Connectors should have a `spec` method."""
+
+
+def run_test_job(
+    connector: IConnector | type[IConnector] | Callable[[], IConnector],
+    verb: Literal["read", "check", "discover"],
+    test_scenario: ConnectorTestScenario,
+    *,
+    catalog: ConfiguredAirbyteCatalog | dict[str, Any] | None = None,
+) -> entrypoint_wrapper.EntrypointOutput:
+    """Run a test scenario from provided CLI args and return the result."""
+    if not connector:
+        raise ValueError("Connector is required")
+
+    if catalog and isinstance(catalog, ConfiguredAirbyteCatalog):
+        # Convert the catalog to a dict if it's already a ConfiguredAirbyteCatalog.
+        catalog = asdict(catalog)
+
+    connector_obj: IConnector
+    if isinstance(connector, type) or callable(connector):
+        # If the connector is a class or a factory lambda, instantiate it.
+        connector_obj = connector()
+    elif isinstance(connector, IConnector):
+        connector_obj = connector
+    else:
+        raise ValueError(
+            f"Invalid connector input: {type(connector)}",
+        )
+
+    args: list[str] = [verb]
+    if test_scenario.config_path:
+        args += ["--config", str(test_scenario.config_path)]
+    elif test_scenario.config_dict:
+        config_path = (
+            Path(tempfile.gettempdir()) / "airbyte-test" / f"temp_config_{uuid.uuid4().hex}.json"
+        )
+        config_path.parent.mkdir(parents=True, exist_ok=True)
+        config_path.write_text(orjson.dumps(test_scenario.config_dict).decode())
+        args += ["--config", str(config_path)]
+
+    catalog_path: Path | None = None
+    if verb not in ["discover", "check"]:
+        # We need a catalog for read.
+        if catalog:
+            # Write the catalog to a temp json file and pass the path to the file as an argument.
+            catalog_path = (
+                Path(tempfile.gettempdir())
+                / "airbyte-test"
+                / f"temp_catalog_{uuid.uuid4().hex}.json"
+            )
+            catalog_path.parent.mkdir(parents=True, exist_ok=True)
+            catalog_path.write_text(orjson.dumps(catalog).decode())
+        elif test_scenario.configured_catalog_path:
+            catalog_path = Path(test_scenario.configured_catalog_path)
+
+        if catalog_path:
+            args += ["--catalog", str(catalog_path)]
+
+    # This is a bit of a hack because the source needs the catalog early.
+    # Because it *also* can fail, we have to redundantly wrap it in a try/except block.
+
+    result: entrypoint_wrapper.EntrypointOutput = entrypoint_wrapper._run_command(  # noqa: SLF001  # Non-public API
+        source=connector_obj,  # type: ignore [arg-type]
+        args=args,
+        expecting_exception=test_scenario.expect_exception,
+    )
+    if result.errors and not test_scenario.expect_exception:
+        raise AssertionError(
+            f"Expected no errors but got {len(result.errors)}: \n" + _errors_to_str(result)
+        )
+
+    if verb == "check":
+        # Check is expected to fail gracefully without an exception.
+        # Instead, we assert that we have a CONNECTION_STATUS message with
+        # a failure status.
+        assert len(result.connection_status_messages) == 1, (
+            "Expected exactly one CONNECTION_STATUS message. Got "
+            f"{len(result.connection_status_messages)}:\n"
+            + "\n".join([str(msg) for msg in result.connection_status_messages])
+            + _errors_to_str(result)
+        )
+        if test_scenario.expect_exception:
+            conn_status = result.connection_status_messages[0].connectionStatus
+            assert conn_status, (
+                "Expected CONNECTION_STATUS message to be present. Got: \n"
+                + "\n".join([str(msg) for msg in result.connection_status_messages])
+            )
+            assert conn_status.status == Status.FAILED, (
+                "Expected CONNECTION_STATUS message to be FAILED. Got: \n"
+                + "\n".join([str(msg) for msg in result.connection_status_messages])
+            )
+
+        return result
+
+    # For all other verbs, we assert check that an exception is raised (or not).
+    if test_scenario.expect_exception:
+        if not result.errors:
+            raise AssertionError("Expected exception but got none.")
+
+        return result
+
+    assert not result.errors, (
+        f"Expected no errors but got {len(result.errors)}: \n" + _errors_to_str(result)
+    )
+
+    return result

airbyte_cdk/test/entrypoint_wrapper.py

@@ -82,6 +82,10 @@ class EntrypointOutput:
     def state_messages(self) -> List[AirbyteMessage]:
         return self._get_message_by_types([Type.STATE])
 
+    @property
+    def connection_status_messages(self) -> List[AirbyteMessage]:
+        return self._get_message_by_types([Type.CONNECTION_STATUS])
+
     @property
     def most_recent_state(self) -> Any:
         state_messages = self._get_message_by_types([Type.STATE])

airbyte_cdk/test/mock_http/response_builder.py

@@ -198,14 +198,6 @@ def find_template(resource: str, execution_folder: str) -> Dict[str, Any]:
         return json.load(template_file)  # type: ignore  # we assume the dev correctly set up the resource file
 
 
-def find_binary_response(resource: str, execution_folder: str) -> bytes:
-    response_filepath = str(
-        get_unit_test_folder(execution_folder) / "resource" / "http" / "response" / f"{resource}"
-    )
-    with open(response_filepath, "rb") as response_file:
-        return response_file.read()  # type: ignore  # we assume the dev correctly set up the resource file
-
-
 def create_record_builder(
     response_template: Dict[str, Any],
     records_path: Union[FieldPath, NestedPath],

airbyte_cdk/test/standard_tests/__init__.py

@@ -0,0 +1,46 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+'''FAST Airbyte Standard Tests
+
+This module provides a set of base classes for declarative connector test suites.
+The goal of this module is to provide a robust and extensible framework for testing Airbyte
+connectors.
+
+Example usage:
+
+```python
+# `test_airbyte_standards.py`
+from airbyte_cdk.test import standard_tests
+
+pytest_plugins = [
+    "airbyte_cdk.test.standard_tests.pytest_hooks",
+]
+
+
+class TestSuiteSourcePokeAPI(standard_tests.DeclarativeSourceTestSuite):
+    """Test suite for the source."""
+```
+
+Available test suites base classes:
+- `DeclarativeSourceTestSuite`: A test suite for declarative sources.
+- `SourceTestSuiteBase`: A test suite for sources.
+- `DestinationTestSuiteBase`: A test suite for destinations.
+
+'''
+
+from airbyte_cdk.test.standard_tests.connector_base import (
+    ConnectorTestScenario,
+    ConnectorTestSuiteBase,
+)
+from airbyte_cdk.test.standard_tests.declarative_sources import (
+    DeclarativeSourceTestSuite,
+)
+from airbyte_cdk.test.standard_tests.destination_base import DestinationTestSuiteBase
+from airbyte_cdk.test.standard_tests.source_base import SourceTestSuiteBase
+
+__all__ = [
+    "ConnectorTestScenario",
+    "ConnectorTestSuiteBase",
+    "DeclarativeSourceTestSuite",
+    "DestinationTestSuiteBase",
+    "SourceTestSuiteBase",
+]