airbyte-cdk 6.45.1__tar.gz → 6.45.1.post46.dev14423672753__tar.gz

{airbyte_cdk-6.45.1 → airbyte_cdk-6.45.1.post46.dev14423672753}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: airbyte-cdk
-Version: 6.45.1
+Version: 6.45.1.post46.dev14423672753
 Summary: A framework for writing Airbyte Connectors.
 Home-page: https://airbyte.com
 License: MIT
@@ -26,6 +26,7 @@ Requires-Dist: airbyte-protocol-models-dataclasses (>=0.14,<0.15)
 Requires-Dist: anyascii (>=0.3.2,<0.4.0)
 Requires-Dist: avro (>=1.11.2,<1.13.0) ; extra == "file-based"
 Requires-Dist: backoff
+Requires-Dist: boltons (>=25.0.0,<26.0.0)
 Requires-Dist: cachetools
 Requires-Dist: cohere (==4.21) ; extra == "vector-db-based"
 Requires-Dist: cryptography (>=44.0.0,<45.0.0)

{airbyte_cdk-6.45.1 → airbyte_cdk-6.45.1.post46.dev14423672753}/airbyte_cdk/connector_builder/connector_builder_handler.py

@@ -35,8 +35,10 @@ MAX_RECORDS_KEY = "max_records"
 MAX_STREAMS_KEY = "max_streams"
 
 
-@dataclass
+@dataclass(kw_only=True)
 class TestLimits:
+    __test__: bool = False  # Prevent pytest from treating this as a test case, despite its name
+
     max_records: int = field(default=DEFAULT_MAXIMUM_RECORDS)
     max_pages_per_slice: int = field(default=DEFAULT_MAXIMUM_NUMBER_OF_PAGES_PER_SLICE)
     max_slices: int = field(default=DEFAULT_MAXIMUM_NUMBER_OF_SLICES)
@@ -51,7 +53,12 @@ def get_limits(config: Mapping[str, Any]) -> TestLimits:
     max_slices = command_config.get(MAX_SLICES_KEY) or DEFAULT_MAXIMUM_NUMBER_OF_SLICES
     max_records = command_config.get(MAX_RECORDS_KEY) or DEFAULT_MAXIMUM_RECORDS
     max_streams = command_config.get(MAX_STREAMS_KEY) or DEFAULT_MAXIMUM_STREAMS
-    return TestLimits(max_records, max_pages_per_slice, max_slices, max_streams)
+    return TestLimits(
+        max_records=max_records,
+        max_pages_per_slice=max_pages_per_slice,
+        max_slices=max_slices,
+        max_streams=max_streams,
+    )
 
 
 def create_source(config: Mapping[str, Any], limits: TestLimits) -> ManifestDeclarativeSource:
@@ -79,7 +86,9 @@ def read_stream(
 ) -> AirbyteMessage:
     try:
         test_read_handler = TestReader(
-            limits.max_pages_per_slice, limits.max_slices, limits.max_records
+            max_pages_per_slice=limits.max_pages_per_slice,
+            max_slices=limits.max_slices,
+            max_record_limit=limits.max_records,
         )
         # The connector builder only supports a single stream
         stream_name = configured_catalog.streams[0].stream.name

{airbyte_cdk-6.45.1 → airbyte_cdk-6.45.1.post46.dev14423672753}/airbyte_cdk/connector_builder/main.py

@@ -78,9 +78,9 @@ def handle_connector_builder_request(
     if command == "resolve_manifest":
         return resolve_manifest(source)
     elif command == "test_read":
-        assert (
-            catalog is not None
-        ), "`test_read` requires a valid `ConfiguredAirbyteCatalog`, got None."
+        assert catalog is not None, (
+            "`test_read` requires a valid `ConfiguredAirbyteCatalog`, got None."
+        )
         return read_stream(source, config, catalog, state, limits)
     elif command == "full_resolve_manifest":
         return full_resolve_manifest(source, limits)

{airbyte_cdk-6.45.1 → airbyte_cdk-6.45.1.post46.dev14423672753}/airbyte_cdk/connector_builder/test_reader/reader.py

@@ -66,6 +66,8 @@ class TestReader:
 
     """
 
+    __test__: bool = False  # Prevent pytest from treating this as a test case, despite its name
+
     logger = logging.getLogger("airbyte.connector-builder")
 
     def __init__(

{airbyte_cdk-6.45.1 → airbyte_cdk-6.45.1.post46.dev14423672753}/airbyte_cdk/sources/concurrent_source/concurrent_source.py

@@ -49,9 +49,9 @@ class ConcurrentSource:
         too_many_generator = (
             not is_single_threaded and initial_number_of_partitions_to_generate >= num_workers
         )
-        assert (
-            not too_many_generator
-        ), "It is required to have more workers than threads generating partitions"
+        assert not too_many_generator, (
+            "It is required to have more workers than threads generating partitions"
+        )
         threadpool = ThreadPoolManager(
             concurrent.futures.ThreadPoolExecutor(
                 max_workers=num_workers, thread_name_prefix="workerpool"

{airbyte_cdk-6.45.1 → airbyte_cdk-6.45.1.post46.dev14423672753}/airbyte_cdk/sources/file_based/file_based_source.py

@@ -282,9 +282,9 @@ class FileBasedSource(ConcurrentSourceAdapter, ABC):
                     and hasattr(self, "_concurrency_level")
                     and self._concurrency_level is not None
                 ):
-                    assert (
-                        state_manager is not None
-                    ), "No ConnectorStateManager was created, but it is required for incremental syncs. This is unexpected. Please contact Support."
+                    assert state_manager is not None, (
+                        "No ConnectorStateManager was created, but it is required for incremental syncs. This is unexpected. Please contact Support."
+                    )
 
                     cursor = self.cursor_cls(
                         stream_config,

{airbyte_cdk-6.45.1 → airbyte_cdk-6.45.1.post46.dev14423672753}/airbyte_cdk/sources/file_based/file_types/avro_parser.py

@@ -154,7 +154,7 @@ class AvroParser(FileTypeParser):
                 # For example: ^-?\d{1,5}(?:\.\d{1,3})?$ would accept 12345.123 and 123456.12345 would  be rejected
                 return {
                     "type": "string",
-                    "pattern": f"^-?\\d{{{1,max_whole_number_range}}}(?:\\.\\d{1,decimal_range})?$",
+                    "pattern": f"^-?\\d{{{1, max_whole_number_range}}}(?:\\.\\d{1, decimal_range})?$",
                 }
             elif "logicalType" in avro_field:
                 if avro_field["logicalType"] not in AVRO_LOGICAL_TYPE_TO_JSON:

{airbyte_cdk-6.45.1 → airbyte_cdk-6.45.1.post46.dev14423672753}/airbyte_cdk/sources/file_based/stream/concurrent/adapters.py

@@ -284,9 +284,9 @@ class FileBasedStreamPartition(Partition):
     def to_slice(self) -> Optional[Mapping[str, Any]]:
         if self._slice is None:
             return None
-        assert (
-            len(self._slice["files"]) == 1
-        ), f"Expected 1 file per partition but got {len(self._slice['files'])} for stream {self.stream_name()}"
+        assert len(self._slice["files"]) == 1, (
+            f"Expected 1 file per partition but got {len(self._slice['files'])} for stream {self.stream_name()}"
+        )
         file = self._slice["files"][0]
         return {"files": [file]}
 

{airbyte_cdk-6.45.1 → airbyte_cdk-6.45.1.post46.dev14423672753}/airbyte_cdk/sql/shared/sql_processor.py

@@ -326,9 +326,9 @@ class SqlProcessorBase(abc.ABC):
 
         if DEBUG_MODE:
             found_schemas = schemas_list
-            assert (
-                schema_name in found_schemas
-            ), f"Schema {schema_name} was not created. Found: {found_schemas}"
+            assert schema_name in found_schemas, (
+                f"Schema {schema_name} was not created. Found: {found_schemas}"
+            )
 
     def _quote_identifier(self, identifier: str) -> str:
         """Return the given identifier, quoted."""
@@ -617,10 +617,10 @@ class SqlProcessorBase(abc.ABC):
         self._execute_sql(
             f"""
             INSERT INTO {self._fully_qualified(final_table_name)} (
-            {f',{nl}  '.join(columns)}
+            {f",{nl}  ".join(columns)}
             )
             SELECT
-            {f',{nl}  '.join(columns)}
+            {f",{nl}  ".join(columns)}
             FROM {self._fully_qualified(temp_table_name)}
             """,
         )
@@ -645,8 +645,7 @@ class SqlProcessorBase(abc.ABC):
         deletion_name = f"{final_table_name}_deleteme"
         commands = "\n".join(
             [
-                f"ALTER TABLE {self._fully_qualified(final_table_name)} RENAME "
-                f"TO {deletion_name};",
+                f"ALTER TABLE {self._fully_qualified(final_table_name)} RENAME TO {deletion_name};",
                 f"ALTER TABLE {self._fully_qualified(temp_table_name)} RENAME "
                 f"TO {final_table_name};",
                 f"DROP TABLE {self._fully_qualified(deletion_name)};",
@@ -686,10 +685,10 @@ class SqlProcessorBase(abc.ABC):
                 {set_clause}
             WHEN NOT MATCHED THEN INSERT
             (
-                {f',{nl}    '.join(columns)}
+                {f",{nl}    ".join(columns)}
             )
             VALUES (
-                tmp.{f',{nl}    tmp.'.join(columns)}
+                tmp.{f",{nl}    tmp.".join(columns)}
             );
             """,
         )

airbyte_cdk-6.45.1.post46.dev14423672753/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-6.45.1.post46.dev14423672753/airbyte_cdk/test/declarative/models/__init__.py

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

airbyte_cdk-6.45.1.post46.dev14423672753/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 instance, as a Pydantic model.
+
+    This class represents an acceptance test instance, 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-6.45.1.post46.dev14423672753/airbyte_cdk/test/declarative/test_suites/__init__.py

@@ -0,0 +1,25 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+"""Declarative test suites.
+
+Here we have base classes for a robust set of declarative connector test suites.
+"""
+
+from airbyte_cdk.test.declarative.test_suites.connector_base import (
+    ConnectorTestScenario,
+    ConnectorTestSuiteBase,
+    generate_tests,
+)
+from airbyte_cdk.test.declarative.test_suites.declarative_sources import (
+    DeclarativeSourceTestSuite,
+)
+from airbyte_cdk.test.declarative.test_suites.destination_base import DestinationTestSuiteBase
+from airbyte_cdk.test.declarative.test_suites.source_base import SourceTestSuiteBase
+
+__all__ = [
+    "ConnectorTestScenario",
+    "ConnectorTestSuiteBase",
+    "DeclarativeSourceTestSuite",
+    "DestinationTestSuiteBase",
+    "SourceTestSuiteBase",
+    "generate_tests",
+]

airbyte_cdk-6.45.1.post46.dev14423672753/airbyte_cdk/test/declarative/test_suites/connector_base.py

@@ -0,0 +1,223 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+"""Base class for connector test suites."""
+
+from __future__ import annotations
+
+import abc
+import inspect
+import sys
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any, Literal, cast
+
+import pytest
+import yaml
+from boltons.typeutils import classproperty
+
+from airbyte_cdk import Connector
+from airbyte_cdk.models import (
+    AirbyteMessage,
+    Type,
+)
+from airbyte_cdk.test import entrypoint_wrapper
+from airbyte_cdk.test.declarative.models import (
+    ConnectorTestScenario,
+)
+from airbyte_cdk.test.declarative.utils.job_runner import IConnector, run_test_job
+
+ACCEPTANCE_TEST_CONFIG = "acceptance-test-config.yml"
+MANIFEST_YAML = "manifest.yaml"
+
+
+class JavaClass(str):
+    """A string that represents a Java class."""
+
+
+class DockerImage(str):
+    """A string that represents a Docker image."""
+
+
+class RunnableConnector(abc.ABC):
+    """A connector that can be run in a test scenario."""
+
+    @abc.abstractmethod
+    def launch(cls, args: list[str] | None) -> None: ...
+
+
+def generate_tests(metafunc: pytest.Metafunc) -> None:
+    """
+    A helper for pytest_generate_tests hook.
+
+    If a test method (in a class subclassed from our base class)
+    declares an argument 'instance', this function retrieves the
+    'scenarios' attribute from the test class and parametrizes that
+    test with the values from 'scenarios'.
+
+    ## Usage
+
+    ```python
+    from airbyte_cdk.test.declarative.test_suites.connector_base import (
+        generate_tests,
+        ConnectorTestSuiteBase,
+    )
+
+    def pytest_generate_tests(metafunc):
+        generate_tests(metafunc)
+
+    class TestMyConnector(ConnectorTestSuiteBase):
+        ...
+
+    ```
+    """
+    # Check if the test function requires an 'instance' argument
+    if "instance" in metafunc.fixturenames:
+        # Retrieve the test class
+        test_class = metafunc.cls
+        if test_class is None:
+            raise ValueError("Expected a class here.")
+        # Get the 'scenarios' attribute from the class
+        scenarios_attr = getattr(test_class, "get_scenarios", None)
+        if scenarios_attr is None:
+            raise ValueError(
+                f"Test class {test_class} does not have a 'scenarios' attribute. "
+                "Please define the 'scenarios' attribute in the test class."
+            )
+
+        scenarios = test_class.get_scenarios()
+        ids = [str(scenario) for scenario in scenarios]
+        metafunc.parametrize("instance", scenarios, ids=ids)
+
+
+class ConnectorTestSuiteBase(abc.ABC):
+    """Base class for connector test suites."""
+
+    connector: type[IConnector] | Callable[[], IConnector] | None = None
+    """The connector class or a factory function that returns an instance of IConnector."""
+
+    @classmethod
+    def get_test_class_dir(cls) -> Path:
+        """Get the file path that contains the class."""
+        module = sys.modules[cls.__module__]
+        # Get the directory containing the test file
+        return Path(inspect.getfile(module)).parent
+
+    @classmethod
+    def create_connector(
+        cls,
+        scenario: ConnectorTestScenario,
+    ) -> IConnector:
+        """Instantiate the connector class."""
+        connector = cls.connector  # type: ignore
+        if connector:
+            if callable(connector) or isinstance(connector, type):
+                # If the connector is a class or factory function, instantiate it:
+                return cast(IConnector, connector())  # type: ignore [redundant-cast]
+
+        # Otherwise, we can't instantiate the connector. Fail with a clear error message.
+        raise NotImplementedError(
+            "No connector class or connector factory function provided. "
+            "Please provide a class or factory function in `cls.connector`, or "
+            "override `cls.create_connector()` to define a custom initialization process."
+        )
+
+    def run_test_scenario(
+        self,
+        verb: Literal["read", "check", "discover"],
+        test_scenario: ConnectorTestScenario,
+        *,
+        catalog: dict[str, Any] | None = None,
+    ) -> entrypoint_wrapper.EntrypointOutput:
+        """Run a test job from provided CLI args and return the result."""
+        return run_test_job(
+            self.create_connector(test_scenario),
+            verb,
+            test_instance=test_scenario,
+            catalog=catalog,
+        )
+
+    # Test Definitions
+
+    def test_check(
+        self,
+        instance: ConnectorTestScenario,
+    ) -> None:
+        """Run `connection` acceptance tests."""
+        result = self.run_test_scenario(
+            "check",
+            test_scenario=instance,
+        )
+        conn_status_messages: list[AirbyteMessage] = [
+            msg for msg in result._messages if msg.type == Type.CONNECTION_STATUS
+        ]  # noqa: SLF001  # Non-public API
+        assert len(conn_status_messages) == 1, (
+            f"Expected exactly one CONNECTION_STATUS message. Got: {result._messages}"
+        )
+
+    @classmethod
+    def get_connector_root_dir(cls) -> Path:
+        """Get the root directory of the connector."""
+        for parent in cls.get_test_class_dir().parents:
+            if (parent / MANIFEST_YAML).exists():
+                return parent
+            if (parent / ACCEPTANCE_TEST_CONFIG).exists():
+                return parent
+            if parent.name == "airbyte_cdk":
+                break
+        # If we reach here, we didn't find the manifest file in any parent directory
+        # Check if the manifest file exists in the current directory
+        for parent in Path.cwd().parents:
+            if (parent / MANIFEST_YAML).exists():
+                return parent
+            if (parent / ACCEPTANCE_TEST_CONFIG).exists():
+                return parent
+            if parent.name == "airbyte_cdk":
+                break
+
+        raise FileNotFoundError(
+            "Could not find connector root directory relative to "
+            f"'{str(cls.get_test_class_dir())}' or '{str(Path.cwd())}'."
+        )
+
+    @classproperty
+    def acceptance_test_config_path(cls) -> Path:
+        """Get the path to the acceptance test config file."""
+        result = cls.get_connector_root_dir() / ACCEPTANCE_TEST_CONFIG
+        if result.exists():
+            return result
+
+        raise FileNotFoundError(f"Acceptance test config file not found at: {str(result)}")
+
+    @classmethod
+    def get_scenarios(
+        cls,
+    ) -> list[ConnectorTestScenario]:
+        """Get acceptance tests for a given category.
+
+        This has to be a separate function because pytest does not allow
+        parametrization of fixtures with arguments from the test class itself.
+        """
+        category = "connection"
+        all_tests_config = yaml.safe_load(cls.acceptance_test_config_path.read_text())
+        if "acceptance_tests" not in all_tests_config:
+            raise ValueError(
+                f"Acceptance tests config not found in {cls.acceptance_test_config_path}."
+                f" Found only: {str(all_tests_config)}."
+            )
+        if category not in all_tests_config["acceptance_tests"]:
+            return []
+        if "tests" not in all_tests_config["acceptance_tests"][category]:
+            raise ValueError(f"No tests found for category {category}")
+
+        tests_scenarios = [
+            ConnectorTestScenario.model_validate(test)
+            for test in all_tests_config["acceptance_tests"][category]["tests"]
+            if "iam_role" not in test["config_path"]
+        ]
+        connector_root = cls.get_connector_root_dir().absolute()
+        for test in tests_scenarios:
+            if test.config_path:
+                test.config_path = connector_root / test.config_path
+            if test.configured_catalog_path:
+                test.configured_catalog_path = connector_root / test.configured_catalog_path
+
+        return tests_scenarios

airbyte_cdk-6.45.1.post46.dev14423672753/airbyte_cdk/test/declarative/test_suites/declarative_sources.py

@@ -0,0 +1,74 @@
+import os
+from hashlib import md5
+from pathlib import Path
+from typing import Any, cast
+
+import yaml
+from boltons.typeutils import classproperty
+
+from airbyte_cdk.sources.declarative.concurrent_declarative_source import (
+    ConcurrentDeclarativeSource,
+)
+from airbyte_cdk.test.declarative.models import ConnectorTestScenario
+from airbyte_cdk.test.declarative.test_suites.connector_base import MANIFEST_YAML
+from airbyte_cdk.test.declarative.test_suites.source_base import (
+    SourceTestSuiteBase,
+)
+from airbyte_cdk.test.declarative.utils.job_runner import IConnector
+
+
+def md5_checksum(file_path: Path) -> str:
+    with open(file_path, "rb") as file:
+        return md5(file.read()).hexdigest()
+
+
+class DeclarativeSourceTestSuite(SourceTestSuiteBase):
+    @classproperty
+    def manifest_yaml_path(cls) -> Path:
+        """Get the path to the manifest.yaml file."""
+        result = cls.get_connector_root_dir() / MANIFEST_YAML
+        if result.exists():
+            return result
+
+        raise FileNotFoundError(
+            f"Manifest YAML file not found at {result}. "
+            "Please ensure that the test suite is run in the correct directory.",
+        )
+
+    @classproperty
+    def components_py_path(cls) -> Path | None:
+        """Get the path to the components.py file."""
+        result = cls.get_connector_root_dir() / "components.py"
+        if result.exists():
+            return result
+
+        return None
+
+    @classmethod
+    def create_connector(
+        cls,
+        scenario: ConnectorTestScenario,
+    ) -> IConnector:
+        """Create a connector instance for the test suite."""
+        config: dict[str, Any] = scenario.get_config_dict()
+        # catalog = scenario.get_catalog()
+        # state = scenario.get_state()
+        # source_config = scenario.get_source_config()
+
+        manifest_dict = yaml.safe_load(cls.manifest_yaml_path.read_text())
+        if cls.components_py_path and cls.components_py_path.exists():
+            os.environ["AIRBYTE_ENABLE_UNSAFE_CODE"] = "true"
+            config["__injected_components_py"] = cls.components_py_path.read_text()
+            config["__injected_components_py_checksums"] = {
+                "md5": md5_checksum(cls.components_py_path),
+            }
+
+        return cast(
+            IConnector,
+            ConcurrentDeclarativeSource(
+                config=config,
+                catalog=None,
+                state=None,
+                source_config=manifest_dict,
+            ),
+        )

airbyte_cdk-6.45.1.post46.dev14423672753/airbyte_cdk/test/declarative/test_suites/destination_base.py

@@ -0,0 +1,12 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+"""Base class for destination test suites."""
+
+from airbyte_cdk.test.declarative.test_suites.connector_base import ConnectorTestSuiteBase
+
+
+class DestinationTestSuiteBase(ConnectorTestSuiteBase):
+    """Base class for destination test suites.
+
+    This class provides a base set of functionality for testing destination connectors, and it
+    inherits all generic connector tests from the `ConnectorTestSuiteBase` class.
+    """