airbyte-cdk 0.72.0__py3-none-any.whl → 6.17.1.dev0__py3-none-any.whl

airbyte_cdk/test/mock_http/request.py

@@ -24,7 +24,9 @@ class HttpRequest:
         if not self._parsed_url.query and query_params:
             self._parsed_url = urlparse(f"{url}?{self._encode_qs(query_params)}")
         elif self._parsed_url.query and query_params:
-            raise ValueError("If query params are provided as part of the url, `query_params` should be empty")
+            raise ValueError(
+                "If query params are provided as part of the url, `query_params` should be empty"
+            )
 
         self._headers = headers or {}
         self._body = body
@@ -62,7 +64,9 @@ class HttpRequest:
         return False
 
     @staticmethod
-    def _to_mapping(body: Optional[Union[str, bytes, Mapping[str, Any]]]) -> Optional[Mapping[str, Any]]:
+    def _to_mapping(
+        body: Optional[Union[str, bytes, Mapping[str, Any]]],
+    ) -> Optional[Mapping[str, Any]]:
         if isinstance(body, Mapping):
             return body
         elif isinstance(body, bytes):
@@ -84,4 +88,16 @@ class HttpRequest:
         return f"{self._parsed_url} with headers {self._headers} and body {self._body!r})"
 
     def __repr__(self) -> str:
-        return f"HttpRequest(request={self._parsed_url}, headers={self._headers}, body={self._body!r})"
+        return (
+            f"HttpRequest(request={self._parsed_url}, headers={self._headers}, body={self._body!r})"
+        )
+
+    def __eq__(self, other: Any) -> bool:
+        if isinstance(other, HttpRequest):
+            return (
+                self._parsed_url == other._parsed_url
+                and self._query_params == other._query_params
+                and self._headers == other._headers
+                and self._body == other._body
+            )
+        return False

airbyte_cdk/test/mock_http/response.py

@@ -5,7 +5,9 @@ from typing import Mapping
 
 
 class HttpResponse:
-    def __init__(self, body: str, status_code: int = 200, headers: Mapping[str, str] = MappingProxyType({})):
+    def __init__(
+        self, body: str, status_code: int = 200, headers: Mapping[str, str] = MappingProxyType({})
+    ):
         self._body = body
         self._status_code = status_code
         self._headers = headers

airbyte_cdk/test/mock_http/response_builder.py

@@ -6,7 +6,8 @@ from abc import ABC, abstractmethod
 from pathlib import Path as FilePath
 from typing import Any, Dict, List, Optional, Union
 
-from airbyte_cdk.test.mock_http import HttpResponse
+from airbyte_cdk.test.mock_http.response import HttpResponse
+from airbyte_cdk.test.utils.data import get_unit_test_folder
 
 
 def _extract(path: List[str], response_template: Dict[str, Any]) -> Any:
@@ -90,7 +91,12 @@ class FieldUpdatePaginationStrategy(PaginationStrategy):
 
 
 class RecordBuilder:
-    def __init__(self, template: Dict[str, Any], id_path: Optional[Path], cursor_path: Optional[Union[FieldPath, NestedPath]]):
+    def __init__(
+        self,
+        template: Dict[str, Any],
+        id_path: Optional[Path],
+        cursor_path: Optional[Union[FieldPath, NestedPath]],
+    ):
         self._record = template
         self._id_path = id_path
         self._cursor_path = cursor_path
@@ -108,9 +114,13 @@ class RecordBuilder:
     def _validate_field(self, field_name: str, path: Optional[Path]) -> None:
         try:
             if path and not path.extract(self._record):
-                raise ValueError(f"{field_name} `{path}` was provided but it is not part of the template `{self._record}`")
+                raise ValueError(
+                    f"{field_name} `{path}` was provided but it is not part of the template `{self._record}`"
+                )
         except (IndexError, KeyError) as exception:
-            raise ValueError(f"{field_name} `{path}` was provided but it is not part of the template `{self._record}`") from exception
+            raise ValueError(
+                f"{field_name} `{path}` was provided but it is not part of the template `{self._record}`"
+            ) from exception
 
     def with_id(self, identifier: Any) -> "RecordBuilder":
         self._set_field("id", self._id_path, identifier)
@@ -141,7 +151,7 @@ class HttpResponseBuilder:
         self,
         template: Dict[str, Any],
         records_path: Union[FieldPath, NestedPath],
-        pagination_strategy: Optional[PaginationStrategy]
+        pagination_strategy: Optional[PaginationStrategy],
     ):
         self._response = template
         self._records: List[RecordBuilder] = []
@@ -172,16 +182,18 @@ class HttpResponseBuilder:
 
 
 def _get_unit_test_folder(execution_folder: str) -> FilePath:
-    path = FilePath(execution_folder)
-    while path.name != "unit_tests":
-        if path.name == path.root or path.name == path.drive:
-            raise ValueError(f"Could not find `unit_tests` folder as a parent of {execution_folder}")
-        path = path.parent
-    return path
+    # FIXME: This function should be removed after the next CDK release to avoid breaking amazon-seller-partner test code.
+    return get_unit_test_folder(execution_folder)
 
 
 def find_template(resource: str, execution_folder: str) -> Dict[str, Any]:
-    response_template_filepath = str(_get_unit_test_folder(execution_folder) / "resource" / "http" / "response" / f"{resource}.json")
+    response_template_filepath = str(
+        get_unit_test_folder(execution_folder)
+        / "resource"
+        / "http"
+        / "response"
+        / f"{resource}.json"
+    )
     with open(response_template_filepath, "r") as template_file:
         return json.load(template_file)  # type: ignore  # we assume the dev correctly set up the resource file
 
@@ -198,16 +210,20 @@ def create_record_builder(
     try:
         record_template = records_path.extract(response_template)[0]
         if not record_template:
-            raise ValueError(f"Could not extract any record from template at path `{records_path}`. "
-                             f"Please fix the template to provide a record sample or fix `records_path`.")
+            raise ValueError(
+                f"Could not extract any record from template at path `{records_path}`. "
+                f"Please fix the template to provide a record sample or fix `records_path`."
+            )
         return RecordBuilder(record_template, record_id_path, record_cursor_path)
     except (IndexError, KeyError):
-        raise ValueError(f"Error while extracting records at path `{records_path}` from response template `{response_template}`")
+        raise ValueError(
+            f"Error while extracting records at path `{records_path}` from response template `{response_template}`"
+        )
 
 
 def create_response_builder(
     response_template: Dict[str, Any],
     records_path: Union[FieldPath, NestedPath],
-    pagination_strategy: Optional[PaginationStrategy] = None
+    pagination_strategy: Optional[PaginationStrategy] = None,
 ) -> HttpResponseBuilder:
     return HttpResponseBuilder(response_template, records_path, pagination_strategy)

airbyte_cdk/test/state_builder.py

@@ -2,7 +2,13 @@
 
 from typing import Any, List
 
-from airbyte_protocol.models import AirbyteStateMessage
+from airbyte_cdk.models import (
+    AirbyteStateBlob,
+    AirbyteStateMessage,
+    AirbyteStateType,
+    AirbyteStreamState,
+    StreamDescriptor,
+)
 
 
 class StateBuilder:
@@ -10,15 +16,17 @@ class StateBuilder:
         self._state: List[AirbyteStateMessage] = []
 
     def with_stream_state(self, stream_name: str, state: Any) -> "StateBuilder":
-        self._state.append(AirbyteStateMessage.parse_obj({
-            "type": "STREAM",
-            "stream": {
-                "stream_state": state,
-                "stream_descriptor": {
-                    "name": stream_name
-                }
-            }
-        }))
+        self._state.append(
+            AirbyteStateMessage(
+                type=AirbyteStateType.STREAM,
+                stream=AirbyteStreamState(
+                    stream_state=state
+                    if isinstance(state, AirbyteStateBlob)
+                    else AirbyteStateBlob(state),
+                    stream_descriptor=StreamDescriptor(**{"name": stream_name}),
+                ),
+            )
+        )
         return self
 
     def build(self) -> List[AirbyteStateMessage]:

airbyte_cdk/test/utils/__init__.py

@@ -0,0 +1 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.

airbyte_cdk/test/utils/data.py

@@ -0,0 +1,24 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+
+from pydantic import FilePath
+
+
+def get_unit_test_folder(execution_folder: str) -> FilePath:
+    path = FilePath(execution_folder)
+    while path.name != "unit_tests":
+        if path.name == path.root or path.name == path.drive:
+            raise ValueError(
+                f"Could not find `unit_tests` folder as a parent of {execution_folder}"
+            )
+        path = path.parent
+    return path
+
+
+def read_resource_file_contents(resource: str, test_location: str) -> str:
+    """Read the contents of a test data file from the test resource folder."""
+    file_path = str(
+        get_unit_test_folder(test_location) / "resource" / "http" / "response" / f"{resource}"
+    )
+    with open(file_path) as f:
+        response = f.read()
+    return response

airbyte_cdk/test/utils/http_mocking.py

@@ -0,0 +1,16 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+
+import re
+from typing import Any, Mapping
+
+from requests_mock import Mocker
+
+
+def register_mock_responses(
+    mocker: Mocker, http_calls: list[Mapping[str, Mapping[str, Any]]]
+) -> None:
+    """Register a list of HTTP request-response pairs."""
+    for call in http_calls:
+        request, response = call["request"], call["response"]
+        matcher = re.compile(request["url"]) if request["is_regex"] else request["url"]
+        mocker.register_uri(request["method"], matcher, **response)

airbyte_cdk/test/utils/manifest_only_fixtures.py

@@ -0,0 +1,60 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+
+
+import importlib.util
+from pathlib import Path
+from types import ModuleType
+from typing import Optional
+
+import pytest
+
+# The following fixtures are used to load a manifest-only connector's components module and manifest file.
+# They can be accessed from any test file in the connector's unit_tests directory by importing them as follows:
+
+# from airbyte_cdk.test.utils.manifest_only_fixtures import components_module, connector_dir, manifest_path
+
+# individual components can then be referenced as: components_module.<CustomComponentClass>
+
+
+@pytest.fixture(scope="session")
+def connector_dir(request: pytest.FixtureRequest) -> Path:
+    """Return the connector's root directory."""
+
+    current_dir = Path(request.config.invocation_params.dir)
+
+    # If the tests are run locally from the connector's unit_tests directory, return the parent (connector) directory
+    if current_dir.name == "unit_tests":
+        return current_dir.parent
+    # In CI, the tests are run from the connector directory itself
+    return current_dir
+
+
+@pytest.fixture(scope="session")
+def components_module(connector_dir: Path) -> Optional[ModuleType]:
+    """Load and return the components module from the connector directory.
+
+    This assumes the components module is located at <connector_dir>/components.py.
+    """
+    components_path = connector_dir / "components.py"
+    if not components_path.exists():
+        return None
+
+    components_spec = importlib.util.spec_from_file_location("components", components_path)
+    if components_spec is None:
+        return None
+
+    components_module = importlib.util.module_from_spec(components_spec)
+    if components_spec.loader is None:
+        return None
+
+    components_spec.loader.exec_module(components_module)
+    return components_module
+
+
+@pytest.fixture(scope="session")
+def manifest_path(connector_dir: Path) -> Path:
+    """Return the path to the connector's manifest file."""
+    path = connector_dir / "manifest.yaml"
+    if not path.exists():
+        raise FileNotFoundError(f"Manifest file not found at {path}")
+    return path

airbyte_cdk/test/utils/reading.py

@@ -0,0 +1,26 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+
+from typing import Any, List, Mapping, Optional
+
+from airbyte_cdk import AbstractSource
+from airbyte_cdk.models import AirbyteStateMessage, ConfiguredAirbyteCatalog, SyncMode
+from airbyte_cdk.test.catalog_builder import CatalogBuilder
+from airbyte_cdk.test.entrypoint_wrapper import EntrypointOutput, read
+
+
+def catalog(stream_name: str, sync_mode: SyncMode) -> ConfiguredAirbyteCatalog:
+    """Create a catalog with a single stream."""
+    return CatalogBuilder().with_stream(stream_name, sync_mode).build()
+
+
+def read_records(
+    source: AbstractSource,
+    config: Mapping[str, Any],
+    stream_name: str,
+    sync_mode: SyncMode,
+    state: Optional[List[AirbyteStateMessage]] = None,
+    expecting_exception: bool = False,
+) -> EntrypointOutput:
+    """Read records from a stream."""
+    _catalog = catalog(stream_name, sync_mode)
+    return read(source, config, _catalog, state, expecting_exception)

airbyte_cdk/utils/__init__.py

@@ -3,7 +3,8 @@
 #
 
 from .is_cloud_environment import is_cloud_environment
+from .print_buffer import PrintBuffer
 from .schema_inferrer import SchemaInferrer
 from .traced_exception import AirbyteTracedException
 
-__all__ = ["AirbyteTracedException", "SchemaInferrer", "is_cloud_environment"]
+__all__ = ["AirbyteTracedException", "SchemaInferrer", "is_cloud_environment", "PrintBuffer"]

airbyte_cdk/utils/airbyte_secrets_utils.py

@@ -4,7 +4,7 @@
 
 from typing import Any, List, Mapping
 
-import dpath.util
+import dpath
 
 
 def get_secret_paths(spec: Mapping[str, Any]) -> List[List[str]]:
@@ -36,7 +36,9 @@ def get_secret_paths(spec: Mapping[str, Any]) -> List[List[str]]:
     return paths
 
 
-def get_secrets(connection_specification: Mapping[str, Any], config: Mapping[str, Any]) -> List[Any]:
+def get_secrets(
+    connection_specification: Mapping[str, Any], config: Mapping[str, Any]
+) -> List[Any]:
     """
     Get a list of secret values from the source config based on the source specification
     :type connection_specification: the connection_specification field of an AirbyteSpecification i.e the JSONSchema definition
@@ -45,7 +47,7 @@ def get_secrets(connection_specification: Mapping[str, Any], config: Mapping[str
     result = []
     for path in secret_paths:
         try:
-            result.append(dpath.util.get(config, path))
+            result.append(dpath.get(config, path))  # type: ignore  # dpath expect MutableMapping but doesn't need it
         except KeyError:
             # Since we try to get paths to all known secrets in the spec, in the case of oneOfs, some secret fields may not be present
             # In that case, a KeyError is thrown. This is expected behavior.

airbyte_cdk/utils/analytics_message.py

@@ -3,7 +3,13 @@
 import time
 from typing import Any, Optional
 
-from airbyte_cdk.models import AirbyteAnalyticsTraceMessage, AirbyteMessage, AirbyteTraceMessage, TraceType, Type
+from airbyte_cdk.models import (
+    AirbyteAnalyticsTraceMessage,
+    AirbyteMessage,
+    AirbyteTraceMessage,
+    TraceType,
+    Type,
+)
 
 
 def create_analytics_message(type: str, value: Optional[Any]) -> AirbyteMessage:
@@ -12,6 +18,8 @@ def create_analytics_message(type: str, value: Optional[Any]) -> AirbyteMessage:
         trace=AirbyteTraceMessage(
             type=TraceType.ANALYTICS,
             emitted_at=time.time() * 1000,
-            analytics=AirbyteAnalyticsTraceMessage(type=type, value=str(value) if value is not None else None),
+            analytics=AirbyteAnalyticsTraceMessage(
+                type=type, value=str(value) if value is not None else None
+            ),
         ),
     )

airbyte_cdk/utils/datetime_format_inferrer.py

@@ -29,7 +29,10 @@ class DatetimeFormatInferrer:
             "%Y-%m",
             "%d-%m-%Y",
         ]
-        self._timestamp_heuristic_ranges = [range(1_000_000_000, 2_000_000_000), range(1_000_000_000_000, 2_000_000_000_000)]
+        self._timestamp_heuristic_ranges = [
+            range(1_000_000_000, 2_000_000_000),
+            range(1_000_000_000_000, 2_000_000_000_000),
+        ]
 
     def _can_be_datetime(self, value: Any) -> bool:
         """Checks if the value can be a datetime.

airbyte_cdk/utils/event_timing.py

@@ -7,7 +7,7 @@ import logging
 import time
 from contextlib import contextmanager
 from dataclasses import dataclass, field
-from typing import Optional
+from typing import Any, Generator, Literal, Optional
 
 logger = logging.getLogger("airbyte")
 
@@ -18,13 +18,13 @@ class EventTimer:
        Event nesting follows a LIFO pattern, so finish will apply to the last started event.
     """
 
-    def __init__(self, name):
+    def __init__(self, name: str) -> None:
         self.name = name
-        self.events = {}
+        self.events: dict[str, Any] = {}
         self.count = 0
-        self.stack = []
+        self.stack: list[Any] = []
 
-    def start_event(self, name):
+    def start_event(self, name: str) -> None:
         """
         Start a new event and push it to the stack.
         """
@@ -32,7 +32,7 @@ class EventTimer:
         self.count += 1
         self.stack.insert(0, self.events[name])
 
-    def finish_event(self):
+    def finish_event(self) -> None:
         """
         Finish the current event and pop it from the stack.
         """
@@ -43,7 +43,7 @@ class EventTimer:
         else:
             logger.warning(f"{self.name} finish_event called without start_event")
 
-    def report(self, order_by="name"):
+    def report(self, order_by: Literal["name", "duration"] = "name") -> str:
         """
         :param order_by: 'name' or 'duration'
         """
@@ -69,15 +69,15 @@ class Event:
             return (self.end - self.start) / 1e9
         return float("+inf")
 
-    def __str__(self):
+    def __str__(self) -> str:
         return f"{self.name} {datetime.timedelta(seconds=self.duration)}"
 
-    def finish(self):
+    def finish(self) -> None:
         self.end = time.perf_counter_ns()
 
 
 @contextmanager
-def create_timer(name):
+def create_timer(name: str) -> Generator[EventTimer, Any, None]:
     """
     Creates a new EventTimer as a context manager to improve code readability.
     """

airbyte_cdk/utils/mapping_helpers.py

@@ -6,7 +6,9 @@
 from typing import Any, List, Mapping, Optional, Set, Union
 
 
-def combine_mappings(mappings: List[Optional[Union[Mapping[str, Any], str]]]) -> Union[Mapping[str, Any], str]:
+def combine_mappings(
+    mappings: List[Optional[Union[Mapping[str, Any], str]]],
+) -> Union[Mapping[str, Any], str]:
     """
     Combine multiple mappings into a single mapping. If any of the mappings are a string, return
     that string. Raise errors in the following cases:

airbyte_cdk/utils/message_utils.py

@@ -1,17 +1,26 @@
 # Copyright (c) 2024 Airbyte, Inc., all rights reserved.
 
+from airbyte_cdk.models import AirbyteMessage, Type
 from airbyte_cdk.sources.connector_state_manager import HashableStreamDescriptor
-from airbyte_protocol.models import AirbyteMessage, Type
 
 
 def get_stream_descriptor(message: AirbyteMessage) -> HashableStreamDescriptor:
-    if message.type == Type.RECORD:
-        return HashableStreamDescriptor(name=message.record.stream, namespace=message.record.namespace)
-    elif message.type == Type.STATE:
-        if not message.state.stream or not message.state.stream.stream_descriptor:
-            raise ValueError("State message was not in per-stream state format, which is required for record counts.")
-        return HashableStreamDescriptor(
-            name=message.state.stream.stream_descriptor.name, namespace=message.state.stream.stream_descriptor.namespace
-        )
-    else:
-        raise NotImplementedError(f"get_stream_descriptor is not implemented for message type '{message.type}'.")
+    match message.type:
+        case Type.RECORD:
+            return HashableStreamDescriptor(
+                name=message.record.stream,  # type: ignore[union-attr] # record has `stream`
+                namespace=message.record.namespace,  # type: ignore[union-attr] # record has `namespace`
+            )
+        case Type.STATE:
+            if not message.state.stream or not message.state.stream.stream_descriptor:  # type: ignore[union-attr] # state has `stream`
+                raise ValueError(
+                    "State message was not in per-stream state format, which is required for record counts."
+                )
+            return HashableStreamDescriptor(
+                name=message.state.stream.stream_descriptor.name,  # type: ignore[union-attr] # state has `stream`
+                namespace=message.state.stream.stream_descriptor.namespace,  # type: ignore[union-attr] # state has `stream`
+            )
+        case _:
+            raise NotImplementedError(
+                f"get_stream_descriptor is not implemented for message type '{message.type}'."
+            )

airbyte_cdk/utils/print_buffer.py

@@ -0,0 +1,75 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+
+import sys
+import time
+from io import StringIO
+from threading import RLock
+from types import TracebackType
+from typing import Optional
+
+
+class PrintBuffer:
+    """
+    A class to buffer print statements and flush them at a specified interval.
+
+    The PrintBuffer class is designed to capture and buffer output that would
+    normally be printed to the standard output (stdout). This can be useful for
+    scenarios where you want to minimize the number of I/O operations by grouping
+    multiple print statements together and flushing them as a single operation.
+
+    Attributes:
+        buffer (StringIO): A buffer to store the messages before flushing.
+        flush_interval (float): The time interval (in seconds) after which the buffer is flushed.
+        last_flush_time (float): The last time the buffer was flushed.
+        lock (RLock): A reentrant lock to ensure thread-safe operations.
+
+    Methods:
+        write(message: str) -> None:
+            Writes a message to the buffer and flushes if the interval has passed.
+
+        flush() -> None:
+            Flushes the buffer content to the standard output.
+
+        __enter__() -> "PrintBuffer":
+            Enters the runtime context related to this object, redirecting stdout and stderr.
+
+        __exit__(exc_type, exc_val, exc_tb) -> None:
+            Exits the runtime context and restores the original stdout and stderr.
+    """
+
+    def __init__(self, flush_interval: float = 0.1):
+        self.buffer = StringIO()
+        self.flush_interval = flush_interval
+        self.last_flush_time = time.monotonic()
+        self.lock = RLock()
+
+    def write(self, message: str) -> None:
+        with self.lock:
+            self.buffer.write(message)
+            current_time = time.monotonic()
+            if (current_time - self.last_flush_time) >= self.flush_interval:
+                self.flush()
+                self.last_flush_time = current_time
+
+    def flush(self) -> None:
+        with self.lock:
+            combined_message = self.buffer.getvalue()
+            sys.__stdout__.write(combined_message)  # type: ignore[union-attr]
+            self.buffer = StringIO()
+
+    def __enter__(self) -> "PrintBuffer":
+        self.old_stdout, self.old_stderr = sys.stdout, sys.stderr
+        # Used to disable buffering during the pytest session, because it is not compatible with capsys
+        if "pytest" not in str(type(sys.stdout)).lower():
+            sys.stdout = self
+            sys.stderr = self
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[BaseException],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> None:
+        self.flush()
+        sys.stdout, sys.stderr = self.old_stdout, self.old_stderr