airbyte-cdk 6.55.0__py3-none-any.whl → 6.55.1__py3-none-any.whl

airbyte_cdk/cli/airbyte_cdk/_connector.py

@@ -101,7 +101,7 @@ def connector_cli_group() -> None:
     pass
 
 
-@connector_cli_group.command()
+@connector_cli_group.command("test")
 @click.argument(
     "connector",
     required=False,
@@ -114,10 +114,18 @@ def connector_cli_group() -> None:
     default=False,
     help="Only collect tests, do not run them.",
 )
-def test(
+@click.option(
+    "--pytest-arg",
+    "pytest_args",  # ← map --pytest-arg into pytest_args
+    type=str,
+    multiple=True,
+    help="Additional argument(s) to pass to pytest. Can be specified multiple times.",
+)
+def connector_test(
     connector: str | Path | None = None,
     *,
     collect_only: bool = False,
+    pytest_args: list[str] | None = None,
 ) -> None:
     """Run connector tests.
 
@@ -130,19 +138,36 @@ def test(
     directory. If the current working directory is not a connector directory (e.g. starting
     with 'source-') and no connector name or path is provided, the process will fail.
     """
+    click.echo("Connector test command executed.")
+    connector_name, connector_directory = resolve_connector_name_and_directory(connector)
+
+    pytest_args = pytest_args or []
+    if collect_only:
+        pytest_args.append("--collect-only")
+
+    run_connector_tests(
+        connector_name=connector_name,
+        connector_directory=connector_directory,
+        extra_pytest_args=pytest_args,
+    )
+
+
+def run_connector_tests(
+    connector_name: str,
+    connector_directory: Path,
+    extra_pytest_args: list[str],
+) -> None:
     if pytest is None:
         raise ImportError(
             "pytest is not installed. Please install pytest to run the connector tests."
         )
-    click.echo("Connector test command executed.")
-    connector_name, connector_directory = resolve_connector_name_and_directory(connector)
 
     connector_test_suite = create_connector_test_suite(
         connector_name=connector_name if not connector_directory else None,
         connector_directory=connector_directory,
     )
 
-    pytest_args: list[str] = []
+    pytest_args: list[str] = ["-p", "airbyte_cdk.test.standard_tests.pytest_hooks"]
     if connector_directory:
         pytest_args.append(f"--rootdir={connector_directory}")
         os.chdir(str(connector_directory))
@@ -158,8 +183,8 @@ def test(
     test_file_path.parent.mkdir(parents=True, exist_ok=True)
     test_file_path.write_text(file_text)
 
-    if collect_only:
-        pytest_args.append("--collect-only")
+    if extra_pytest_args:
+        pytest_args.extend(extra_pytest_args)
 
     pytest_args.append(str(test_file_path))
 
@@ -170,7 +195,6 @@ def test(
 
     click.echo(f"Running tests from connector directory: {connector_directory}...")
     click.echo(f"Test file: {test_file_path}")
-    click.echo(f"Collect only: {collect_only}")
     click.echo(f"Pytest args: {pytest_args}")
     click.echo("Invoking Pytest...")
     exit_code = pytest.main(

airbyte_cdk/cli/airbyte_cdk/_image.py

@@ -10,6 +10,7 @@ from pathlib import Path
 
 import rich_click as click
 
+from airbyte_cdk.cli.airbyte_cdk._connector import run_connector_tests
 from airbyte_cdk.models.connector_metadata import MetadataFile
 from airbyte_cdk.utils.connector_paths import resolve_connector_name_and_directory
 from airbyte_cdk.utils.docker import (
@@ -88,6 +89,81 @@ def build(
         sys.exit(1)
 
 
+@image_cli_group.command("test")
+@click.argument(
+    "connector",
+    required=False,
+    type=str,
+    metavar="[CONNECTOR]",
+)
+@click.option(
+    "--image",
+    help="Image to test, instead of building a new one.",
+)
+def image_test(  # "image test" command
+    connector: str | None = None,
+    *,
+    image: str | None = None,
+) -> None:
+    """Test a connector Docker image.
+
+    [CONNECTOR] can be a connector name (e.g. 'source-pokeapi'), a path to a connector directory, or omitted to use the current working directory.
+    If a string containing '/' is provided, it is treated as a path. Otherwise, it is treated as a connector name.
+
+    If an image is provided, it will be used for testing instead of building a new one.
+
+    Note: You should run `airbyte-cdk secrets fetch` before running this command to ensure
+    that the secrets are available for the connector tests.
+    """
+    if not verify_docker_installation():
+        click.echo(
+            "Docker is not installed or not running. Please install Docker and try again.", err=True
+        )
+        sys.exit(1)
+
+    connector_name, connector_directory = resolve_connector_name_and_directory(connector)
+
+    # Select only tests with the 'image_tests' mark
+    pytest_args = ["-m", "image_tests"]
+    if not image:
+        metadata_file_path: Path = connector_directory / "metadata.yaml"
+        try:
+            metadata = MetadataFile.from_file(metadata_file_path)
+        except (FileNotFoundError, ValueError) as e:
+            click.echo(
+                f"Error loading metadata file '{metadata_file_path}': {e!s}",
+                err=True,
+            )
+            sys.exit(1)
+
+        tag = "dev-latest"
+        image = f"{metadata.data.dockerRepository}:{tag}"
+        click.echo(f"Building Image for Connector: {image}")
+        try:
+            image = build_connector_image(
+                connector_directory=connector_directory,
+                connector_name=connector_name,
+                metadata=metadata,
+                tag=tag,
+                no_verify=True,
+            )
+        except ConnectorImageBuildError as e:
+            click.echo(
+                f"Error building connector image: {e!s}",
+                err=True,
+            )
+            sys.exit(1)
+
+    pytest_args.extend(["--connector-image", image])
+
+    click.echo(f"Testing Connector Image: {image}")
+    run_connector_tests(
+        connector_name=connector_name,
+        connector_directory=connector_directory,
+        extra_pytest_args=pytest_args,
+    )
+
+
 __all__ = [
     "image_cli_group",
 ]

airbyte_cdk/cli/airbyte_cdk/_secrets.py

@@ -99,12 +99,12 @@ def secrets_cli_group() -> None:
     help="Print GitHub CI mask for secrets.",
     type=bool,
     is_flag=True,
-    default=False,
+    default=None,
 )
 def fetch(
     connector: str | Path | None = None,
     gcp_project_id: str = GCP_PROJECT_ID,
-    print_ci_secrets_masks: bool = False,
+    print_ci_secrets_masks: bool | None = None,
 ) -> None:
     """Fetch secrets for a connector from Google Secret Manager.
 
@@ -181,22 +181,23 @@ def fetch(
         if secret_count == 0:
             raise exceptions[0]
 
-    if not print_ci_secrets_masks:
-        return
-
-    if not os.environ.get("CI", None):
+    if print_ci_secrets_masks and "CI" not in os.environ:
         click.echo(
             "The `--print-ci-secrets-masks` option is only available in CI environments. "
             "The `CI` env var is either not set or not set to a truthy value. "
             "Skipping printing secret masks.",
             err=True,
         )
-        return
-
-    # Else print the CI mask
-    _print_ci_secrets_masks(
-        secrets_dir=secrets_dir,
-    )
+        print_ci_secrets_masks = False
+    elif print_ci_secrets_masks is None:
+        # If not explicitly set, we check if we are in a CI environment
+        # and set to True if so.
+        print_ci_secrets_masks = os.environ.get("CI", "") != ""
+
+    if print_ci_secrets_masks:
+        _print_ci_secrets_masks(
+            secrets_dir=secrets_dir,
+        )
 
 
 @secrets_cli_group.command("list")

airbyte_cdk/models/airbyte_protocol_serializers.py

@@ -4,9 +4,11 @@ from typing import Any, Dict
 from serpyco_rs import CustomType, Serializer
 
 from .airbyte_protocol import (  # type: ignore[attr-defined] # all classes are imported to airbyte_protocol via *
+    AirbyteCatalog,
     AirbyteMessage,
     AirbyteStateBlob,
     AirbyteStateMessage,
+    AirbyteStream,
     AirbyteStreamState,
     ConfiguredAirbyteCatalog,
     ConfiguredAirbyteStream,
@@ -30,6 +32,8 @@ def custom_type_resolver(t: type) -> CustomType[AirbyteStateBlob, Dict[str, Any]
     return AirbyteStateBlobType() if t is AirbyteStateBlob else None
 
 
+AirbyteCatalogSerializer = Serializer(AirbyteCatalog, omit_none=True)
+AirbyteStreamSerializer = Serializer(AirbyteStream, omit_none=True)
 AirbyteStreamStateSerializer = Serializer(
     AirbyteStreamState, omit_none=True, custom_type_resolver=custom_type_resolver
 )

airbyte_cdk/models/connector_metadata.py

@@ -34,6 +34,15 @@ class ConnectorBuildOptions(BaseModel):
     )
 
 
+class SuggestedStreams(BaseModel):
+    """Suggested streams from metadata.yaml."""
+
+    streams: list[str] = Field(
+        default=[],
+        description="List of suggested streams for the connector",
+    )
+
+
 class ConnectorMetadata(BaseModel):
     """Connector metadata from metadata.yaml."""
 
@@ -47,6 +56,11 @@ class ConnectorMetadata(BaseModel):
         description="List of tags for the connector",
     )
 
+    suggestedStreams: SuggestedStreams | None = Field(
+        default=None,
+        description="Suggested streams for the connector",
+    )
+
     @property
     def language(self) -> ConnectorLanguage:
         """Get the connector language."""

airbyte_cdk/test/entrypoint_wrapper.py

@@ -19,9 +19,11 @@ import logging
 import re
 import tempfile
 import traceback
+from collections import deque
+from collections.abc import Generator, Mapping
 from io import StringIO
 from pathlib import Path
-from typing import Any, List, Mapping, Optional, Union
+from typing import Any, List, Literal, Optional, Union, final, overload
 
 import orjson
 from pydantic import ValidationError as V2ValidationError
@@ -36,6 +38,7 @@ from airbyte_cdk.models import (
     AirbyteMessageSerializer,
     AirbyteStateMessage,
     AirbyteStateMessageSerializer,
+    AirbyteStreamState,
     AirbyteStreamStatus,
     ConfiguredAirbyteCatalog,
     ConfiguredAirbyteCatalogSerializer,
@@ -48,13 +51,41 @@ from airbyte_cdk.test.models.scenario import ExpectedOutcome
 
 
 class EntrypointOutput:
-    def __init__(self, messages: List[str], uncaught_exception: Optional[BaseException] = None):
-        try:
-            self._messages = [self._parse_message(message) for message in messages]
-        except V2ValidationError as exception:
-            raise ValueError("All messages are expected to be AirbyteMessage") from exception
+    """A class to encapsulate the output of an Airbyte connector's execution.
+
+    This class can be initialized with a list of messages or a file containing messages.
+    It provides methods to access different types of messages produced during the execution
+    of an Airbyte connector, including both successful messages and error messages.
+
+    When working with records and state messages, it provides both a list and an iterator
+    implementation. Lists are easier to work with, but generators are better suited to handle
+    large volumes of messages without overflowing the available memory.
+    """
+
+    def __init__(
+        self,
+        messages: list[str] | None = None,
+        uncaught_exception: Optional[BaseException] = None,
+        *,
+        message_file: Path | None = None,
+    ) -> None:
+        if messages is None and message_file is None:
+            raise ValueError("Either messages or message_file must be provided")
+        if messages is not None and message_file is not None:
+            raise ValueError("Only one of messages or message_file can be provided")
+
+        self._messages: list[AirbyteMessage] | None = None
+        self._message_file: Path | None = message_file
+        if messages:
+            try:
+                self._messages = [self._parse_message(message) for message in messages]
+            except V2ValidationError as exception:
+                raise ValueError("All messages are expected to be AirbyteMessage") from exception
 
         if uncaught_exception:
+            if self._messages is None:
+                self._messages = []
+
             self._messages.append(
                 assemble_uncaught_exception(
                     type(uncaught_exception), uncaught_exception
@@ -72,39 +103,76 @@ class EntrypointOutput:
             )
 
     @property
-    def records_and_state_messages(self) -> List[AirbyteMessage]:
-        return self._get_message_by_types([Type.RECORD, Type.STATE])
+    def records_and_state_messages(
+        self,
+    ) -> list[AirbyteMessage]:
+        return self.get_message_by_types(
+            message_types=[Type.RECORD, Type.STATE],
+            safe_iterator=False,
+        )
+
+    def records_and_state_messages_iterator(
+        self,
+    ) -> Generator[AirbyteMessage, None, None]:
+        """Returns a generator that yields record and state messages one by one.
+
+        Use this instead of `records_and_state_messages` when the volume of messages could be large
+        enough to overload available memory.
+        """
+        return self.get_message_by_types(
+            message_types=[Type.RECORD, Type.STATE],
+            safe_iterator=True,
+        )
 
     @property
     def records(self) -> List[AirbyteMessage]:
-        return self._get_message_by_types([Type.RECORD])
+        return self.get_message_by_types([Type.RECORD])
+
+    @property
+    def records_iterator(self) -> Generator[AirbyteMessage, None, None]:
+        """Returns a generator that yields record messages one by one.
+
+        Use this instead of `records` when the volume of records could be large
+        enough to overload available memory.
+        """
+        return self.get_message_by_types([Type.RECORD], safe_iterator=True)
 
     @property
     def state_messages(self) -> List[AirbyteMessage]:
-        return self._get_message_by_types([Type.STATE])
+        return self.get_message_by_types([Type.STATE])
 
     @property
     def spec_messages(self) -> List[AirbyteMessage]:
-        return self._get_message_by_types([Type.SPEC])
+        return self.get_message_by_types([Type.SPEC])
 
     @property
     def connection_status_messages(self) -> List[AirbyteMessage]:
-        return self._get_message_by_types([Type.CONNECTION_STATUS])
+        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])
-        if not state_messages:
-            raise ValueError("Can't provide most recent state as there are no state messages")
-        return state_messages[-1].state.stream  # type: ignore[union-attr] # state has `stream`
+    def most_recent_state(self) -> AirbyteStreamState | None:
+        state_message_iterator = self.get_message_by_types(
+            [Type.STATE],
+            safe_iterator=True,
+        )
+        # Use a deque with maxlen=1 to efficiently get the last state message
+        double_ended_queue = deque(state_message_iterator, maxlen=1)
+        try:
+            final_state_message: AirbyteMessage = double_ended_queue.pop()
+        except IndexError:
+            raise ValueError(
+                "Can't provide most recent state as there are no state messages."
+            ) from None
+
+        return final_state_message.state.stream  # type: ignore[union-attr] # state has `stream`
 
     @property
     def logs(self) -> List[AirbyteMessage]:
-        return self._get_message_by_types([Type.LOG])
+        return self.get_message_by_types([Type.LOG])
 
     @property
     def trace_messages(self) -> List[AirbyteMessage]:
-        return self._get_message_by_types([Type.TRACE])
+        return self.get_message_by_types([Type.TRACE])
 
     @property
     def analytics_messages(self) -> List[AirbyteMessage]:
@@ -116,7 +184,7 @@ class EntrypointOutput:
 
     @property
     def catalog(self) -> AirbyteMessage:
-        catalog = self._get_message_by_types([Type.CATALOG])
+        catalog = self.get_message_by_types([Type.CATALOG])
         if len(catalog) != 1:
             raise ValueError(f"Expected exactly one catalog but got {len(catalog)}")
         return catalog[0]
@@ -131,13 +199,80 @@ class EntrypointOutput:
         )
         return list(status_messages)
 
-    def _get_message_by_types(self, message_types: List[Type]) -> List[AirbyteMessage]:
-        return [message for message in self._messages if message.type in message_types]
+    def get_message_iterator(self) -> Generator[AirbyteMessage, None, None]:
+        """Creates a generator which yields messages one by one.
+
+        This will iterate over all messages in the output file (if provided) or the messages
+        provided during initialization. File results are provided first, followed by any
+        messages that were passed in directly.
+        """
+        if self._message_file:
+            try:
+                with open(self._message_file, "r", encoding="utf-8") as file:
+                    for line in file:
+                        if not line.strip():
+                            # Skip empty lines
+                            continue
+
+                        yield self._parse_message(line.strip())
+            except FileNotFoundError:
+                raise ValueError(f"Message file {self._message_file} not found")
+
+        if self._messages is not None:
+            yield from self._messages
+
+    # Overloads to provide proper type hints for different usages of `get_message_by_types`.
+
+    @overload
+    def get_message_by_types(
+        self,
+        message_types: list[Type],
+    ) -> list[AirbyteMessage]: ...
+
+    @overload
+    def get_message_by_types(
+        self,
+        message_types: list[Type],
+        *,
+        safe_iterator: Literal[False],
+    ) -> list[AirbyteMessage]: ...
+
+    @overload
+    def get_message_by_types(
+        self,
+        message_types: list[Type],
+        *,
+        safe_iterator: Literal[True],
+    ) -> Generator[AirbyteMessage, None, None]: ...
+
+    def get_message_by_types(
+        self,
+        message_types: list[Type],
+        *,
+        safe_iterator: bool = False,
+    ) -> list[AirbyteMessage] | Generator[AirbyteMessage, None, None]:
+        """Get messages of specific types.
+
+        If `safe_iterator` is True, returns a generator that yields messages one by one.
+        If `safe_iterator` is False, returns a list of messages.
+
+        Use `safe_iterator=True` when the volume of messages could overload the available
+        memory.
+        """
+        message_generator = self.get_message_iterator()
+
+        if safe_iterator:
+            return (message for message in message_generator if message.type in message_types)
+
+        return [message for message in message_generator if message.type in message_types]
 
     def _get_trace_message_by_trace_type(self, trace_type: TraceType) -> List[AirbyteMessage]:
         return [
             message
-            for message in self._get_message_by_types([Type.TRACE])
+            for message in self.get_message_by_types(
+                [Type.TRACE],
+                safe_iterator=True,
+            )
             if message.trace.type == trace_type  # type: ignore[union-attr] # trace has `type`
         ]
 
@@ -184,7 +319,7 @@ def _run_command(
     parsed_args = AirbyteEntrypoint.parse_args(args)
 
     source_entrypoint = AirbyteEntrypoint(source)
-    messages = []
+    messages: list[str] = []
     uncaught_exception = None
     try:
         for message in source_entrypoint.run(parsed_args):
@@ -199,8 +334,10 @@ def _run_command(
     captured_logs = log_capture_buffer.getvalue().split("\n")[:-1]
 
     parent_logger.removeHandler(stream_handler)
-
-    return EntrypointOutput(messages + captured_logs, uncaught_exception=uncaught_exception)
+    return EntrypointOutput(
+        messages=messages + captured_logs,
+        uncaught_exception=uncaught_exception,
+    )
 
 
 def discover(

airbyte_cdk/test/models/scenario.py

@@ -9,14 +9,20 @@ up iteration cycles.
 
 from __future__ import annotations
 
+import json
+import tempfile
+from contextlib import contextmanager, suppress
 from pathlib import Path  # noqa: TC003  # Pydantic needs this (don't move to 'if typing' block)
-from typing import Any, Literal, cast
+from typing import TYPE_CHECKING, Any, Literal, cast
 
 import yaml
 from pydantic import BaseModel, ConfigDict
 
 from airbyte_cdk.test.models.outcome import ExpectedOutcome
 
+if TYPE_CHECKING:
+    from collections.abc import Generator
+
 
 class ConnectorTestScenario(BaseModel):
     """Acceptance test scenario, as a Pydantic model.
@@ -41,13 +47,13 @@ class ConnectorTestScenario(BaseModel):
     config_path: Path | None = None
     config_dict: dict[str, Any] | None = None
 
-    id: str | None = None
+    _id: str | None = None  # Used to override the default ID generation
 
     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
+    status: Literal["succeed", "failed", "exception"] | None = None
 
     def get_config_dict(
         self,
@@ -93,16 +99,49 @@ class ConnectorTestScenario(BaseModel):
         return ExpectedOutcome.from_status_str(self.status)
 
     @property
-    def instance_name(self) -> str:
-        return self.config_path.stem if self.config_path else "Unnamed Scenario"
+    def id(self) -> str:
+        """Return a unique identifier for the test scenario.
+
+        This is used by PyTest to identify the test scenario.
+        """
+        if self._id:
+            return self._id
 
-    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 self.config_path.stem
+
+        return str(hash(self))
 
-        return f"'{hash(self)}' Test Scenario"
+    def __str__(self) -> str:
+        return f"'{self.id}' Test Scenario"
+
+    @contextmanager
+    def with_temp_config_file(
+        self,
+        connector_root: Path,
+    ) -> Generator[Path, None, None]:
+        """Yield a temporary JSON file path containing the config dict and delete it on exit."""
+        config = self.get_config_dict(
+            empty_if_missing=True,
+            connector_root=connector_root,
+        )
+        with tempfile.NamedTemporaryFile(
+            prefix="config-",
+            suffix=".json",
+            mode="w",
+            delete=False,  # Don't fail if cannot delete the file on exit
+            encoding="utf-8",
+        ) as temp_file:
+            temp_file.write(json.dumps(config))
+            temp_file.flush()
+            # Allow the file to be read by other processes
+            temp_path = Path(temp_file.name)
+            temp_path.chmod(temp_path.stat().st_mode | 0o444)
+            yield temp_path
+
+        # attempt cleanup, ignore errors
+        with suppress(OSError):
+            temp_path.unlink()
 
     def without_expected_outcome(self) -> ConnectorTestScenario:
         """Return a copy of the scenario that does not expect failure or success.

airbyte_cdk/test/standard_tests/__init__.py

@@ -27,10 +27,8 @@ Available test suites base classes:
 
 '''
 
-from airbyte_cdk.test.standard_tests.connector_base import (
-    ConnectorTestScenario,
-    ConnectorTestSuiteBase,
-)
+from airbyte_cdk.test.models.scenario import ConnectorTestScenario
+from airbyte_cdk.test.standard_tests.connector_base import ConnectorTestSuiteBase
 from airbyte_cdk.test.standard_tests.declarative_sources import (
     DeclarativeSourceTestSuite,
 )