airbyte-cdk 0.0.0.dev0__py3-none-any.whl

airbyte_cdk/sql/types.py

@@ -0,0 +1,160 @@
+# noqa: A005  # Allow shadowing the built-in 'types' module
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+
+"""Type conversion methods for SQL Caches."""
+
+from __future__ import annotations
+
+from typing import Any, cast
+
+import sqlalchemy
+
+# Compare to documentation here: https://docs.airbyte.com/understanding-airbyte/supported-data-types
+CONVERSION_MAP = {
+    "string": sqlalchemy.types.VARCHAR,
+    "integer": sqlalchemy.types.BIGINT,
+    "number": sqlalchemy.types.DECIMAL(38, 9),
+    "boolean": sqlalchemy.types.BOOLEAN,
+    "date": sqlalchemy.types.DATE,
+    "timestamp_with_timezone": sqlalchemy.types.TIMESTAMP,
+    "timestamp_without_timezone": sqlalchemy.types.TIMESTAMP,
+    "time_with_timezone": sqlalchemy.types.TIME,
+    "time_without_timezone": sqlalchemy.types.TIME,
+    # Technically 'object' and 'array' as JSON Schema types, not airbyte types.
+    # We include them here for completeness.
+    "object": sqlalchemy.types.JSON,
+    "array": sqlalchemy.types.JSON,
+    "vector_array": sqlalchemy.types.ARRAY,
+}
+
+
+class SQLTypeConversionError(Exception):
+    """An exception to be raised when a type conversion fails."""
+
+
+def _get_airbyte_type(  # noqa: PLR0911  # Too many return statements
+    json_schema_property_def: dict[str, str | dict[str, Any] | list[Any]],
+) -> tuple[str, str | None]:
+    """Get the airbyte type and subtype from a JSON schema property definition.
+
+    Subtype is only used for array types. Otherwise, subtype will return None.
+    """
+    airbyte_type = cast(str, json_schema_property_def.get("airbyte_type", None))
+    if airbyte_type:
+        return airbyte_type, None
+
+    json_schema_type = json_schema_property_def.get("type", None)
+    json_schema_format = json_schema_property_def.get("format", None)
+
+    # if json_schema_type is an array of two strings with one of them being null, pick the other one
+    # this strategy is often used by connectors to indicate a field might not be set all the time
+    if isinstance(json_schema_type, list):
+        non_null_types = [t for t in json_schema_type if t != "null"]
+        if len(non_null_types) == 1:
+            json_schema_type = non_null_types[0]
+
+    if json_schema_type == "string":
+        if json_schema_format == "date":
+            return "date", None
+
+        if json_schema_format == "date-time":
+            return "timestamp_with_timezone", None
+
+        if json_schema_format == "time":
+            return "time_without_timezone", None
+
+    if isinstance(json_schema_type, str) and json_schema_type in {
+        "string",
+        "number",
+        "boolean",
+        "integer",
+    }:
+        return json_schema_type, None
+
+    if json_schema_type == "object":
+        return "object", None
+
+    if json_schema_type == "array":
+        items_def = json_schema_property_def.get("items", None)
+        if isinstance(items_def, dict):
+            try:
+                subtype, _ = _get_airbyte_type(items_def)
+            except SQLTypeConversionError:
+                # We have enough information, so we can ignore parsing errors on subtype.
+                subtype = None
+
+            return "array", subtype
+
+        return "array", None
+
+    if json_schema_type == "vector_array":
+        return "vector_array", "Float"
+
+    err_msg = f"Could not determine airbyte type from JSON schema type: {json_schema_property_def}"
+    raise SQLTypeConversionError(err_msg)
+
+
+class SQLTypeConverter:
+    """A base class to perform type conversions."""
+
+    def __init__(
+        self,
+        conversion_map: dict[str, Any] | None = None,
+    ) -> None:
+        """Initialize the type converter."""
+        self.conversion_map = conversion_map or CONVERSION_MAP
+
+    @classmethod
+    def get_string_type(cls) -> sqlalchemy.types.TypeEngine[str]:
+        """Get the type to use for string data."""
+        return sqlalchemy.types.VARCHAR()
+
+    @classmethod
+    def get_failover_type(cls) -> sqlalchemy.types.TypeEngine[str]:
+        """Get the 'last resort' type to use if no other type is found."""
+        return cls.get_string_type()
+
+    @classmethod
+    def get_json_type(cls) -> sqlalchemy.types.TypeEngine[Any]:
+        """Get the type to use for nested JSON data."""
+        return sqlalchemy.types.JSON()
+
+    def to_sql_type(  # noqa: PLR0911  # Too many return statements
+        self,
+        json_schema_property_def: dict[str, str | dict[str, Any] | list[Any]],
+    ) -> Any:
+        """Convert a value to a SQL type."""
+        try:
+            airbyte_type, _ = _get_airbyte_type(json_schema_property_def)
+            # to-do - is there a better way to check the following
+            if airbyte_type == "vector_array":
+                return sqlalchemy.types.ARRAY(sqlalchemy.types.Float())
+            sql_type = self.conversion_map[airbyte_type]
+        except SQLTypeConversionError:
+            print(f"Could not determine airbyte type from JSON schema: {json_schema_property_def}")
+        except KeyError:
+            print(f"Could not find SQL type for airbyte type: {airbyte_type}")
+        else:
+            # No exceptions were raised, so we can return the SQL type.
+            if isinstance(sql_type, type):
+                # This is a class. Call its constructor.
+                sql_type = sql_type()
+
+            return sql_type
+
+        json_schema_type = json_schema_property_def.get("type", None)
+        json_schema_format = json_schema_property_def.get("format", None)
+
+        if json_schema_type == "string" and json_schema_format == "date":
+            return sqlalchemy.types.DATE()
+
+        if json_schema_type == "string" and json_schema_format == "date-time":
+            return sqlalchemy.types.TIMESTAMP()
+
+        if json_schema_type == "array":
+            return sqlalchemy.types.JSON()
+
+        if json_schema_type == "object":
+            return sqlalchemy.types.JSON()
+
+        return self.get_failover_type()

airbyte_cdk/test/__init__.py

@@ -0,0 +1,7 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+"""
+This package is provided as tooling to help sources test their implementation. It is not expected to be used as production code.
+"""

airbyte_cdk/test/catalog_builder.py

@@ -0,0 +1,81 @@
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+
+from typing import Any, Dict, List, Union, overload
+
+from airbyte_cdk.models import (
+    ConfiguredAirbyteCatalog,
+    ConfiguredAirbyteStream,
+    ConfiguredAirbyteStreamSerializer,
+    SyncMode,
+)
+
+
+class ConfiguredAirbyteStreamBuilder:
+    def __init__(self) -> None:
+        self._stream: Dict[str, Any] = {
+            "stream": {
+                "name": "any name",
+                "json_schema": {},
+                "supported_sync_modes": ["full_refresh", "incremental"],
+                "source_defined_primary_key": [["id"]],
+            },
+            "primary_key": [["id"]],
+            "sync_mode": "full_refresh",
+            "destination_sync_mode": "overwrite",
+        }
+
+    def with_name(self, name: str) -> "ConfiguredAirbyteStreamBuilder":
+        self._stream["stream"]["name"] = name  # type: ignore  # we assume that self._stream["stream"] is a Dict[str, Any]
+        return self
+
+    def with_sync_mode(self, sync_mode: SyncMode) -> "ConfiguredAirbyteStreamBuilder":
+        self._stream["sync_mode"] = sync_mode.name
+        return self
+
+    def with_primary_key(self, pk: List[List[str]]) -> "ConfiguredAirbyteStreamBuilder":
+        self._stream["primary_key"] = pk
+        self._stream["stream"]["source_defined_primary_key"] = pk  # type: ignore  # we assume that self._stream["stream"] is a Dict[str, Any]
+        return self
+
+    def with_json_schema(self, json_schema: Dict[str, Any]) -> "ConfiguredAirbyteStreamBuilder":
+        self._stream["stream"]["json_schema"] = json_schema
+        return self
+
+    def build(self) -> ConfiguredAirbyteStream:
+        return ConfiguredAirbyteStreamSerializer.load(self._stream)
+
+
+class CatalogBuilder:
+    def __init__(self) -> None:
+        self._streams: List[ConfiguredAirbyteStreamBuilder] = []
+
+    @overload
+    def with_stream(self, name: ConfiguredAirbyteStreamBuilder) -> "CatalogBuilder": ...
+
+    @overload
+    def with_stream(self, name: str, sync_mode: SyncMode) -> "CatalogBuilder": ...
+
+    def with_stream(
+        self,
+        name: Union[str, ConfiguredAirbyteStreamBuilder],
+        sync_mode: Union[SyncMode, None] = None,
+    ) -> "CatalogBuilder":
+        # As we are introducing a fully fledge ConfiguredAirbyteStreamBuilder, we would like to deprecate the previous interface
+        # with_stream(str, SyncMode)
+
+        # to avoid a breaking change, `name` needs to stay in the API but this can be either a name or a builder
+        name_or_builder = name
+        builder = (
+            name_or_builder
+            if isinstance(name_or_builder, ConfiguredAirbyteStreamBuilder)
+            else ConfiguredAirbyteStreamBuilder()
+            .with_name(name_or_builder)
+            .with_sync_mode(sync_mode)
+        )
+        self._streams.append(builder)
+        return self
+
+    def build(self) -> ConfiguredAirbyteCatalog:
+        return ConfiguredAirbyteCatalog(
+            streams=list(map(lambda builder: builder.build(), self._streams))
+        )

airbyte_cdk/test/entrypoint_wrapper.py

@@ -0,0 +1,250 @@
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+
+"""
+The AirbyteEntrypoint is important because it is a service layer that orchestrate how we execute commands from the
+[common interface](https://docs.airbyte.com/understanding-airbyte/airbyte-protocol#common-interface) through the source Python
+implementation. There is some logic about which message we send to the platform and when which is relevant for integration testing. Other
+than that, there are integrations point that are annoying to integrate with using Python code:
+* Sources communicate with the platform using stdout. The implication is that the source could just print every message instead of
+    returning things to source.<method> or to using the message repository. WARNING: As part of integration testing, we will not support
+    messages that are simply printed. The reason is that capturing stdout relies on overriding sys.stdout (see
+    https://docs.python.org/3/library/contextlib.html#contextlib.redirect_stdout) which clashes with how pytest captures logs and brings
+    considerations for multithreaded applications. If code you work with uses `print` statements, please migrate to
+    source.message_repository to emit those messages
+* The entrypoint interface relies on file being written on the file system
+"""
+
+import json
+import logging
+import re
+import tempfile
+import traceback
+from io import StringIO
+from pathlib import Path
+from typing import Any, List, Mapping, Optional, Union
+
+import orjson
+from pydantic import ValidationError as V2ValidationError
+from serpyco_rs import SchemaValidationError
+
+from airbyte_cdk.entrypoint import AirbyteEntrypoint
+from airbyte_cdk.exception_handler import assemble_uncaught_exception
+from airbyte_cdk.logger import AirbyteLogFormatter
+from airbyte_cdk.models import (
+    AirbyteLogMessage,
+    AirbyteMessage,
+    AirbyteMessageSerializer,
+    AirbyteStateMessage,
+    AirbyteStateMessageSerializer,
+    AirbyteStreamStatus,
+    ConfiguredAirbyteCatalog,
+    ConfiguredAirbyteCatalogSerializer,
+    Level,
+    TraceType,
+    Type,
+)
+from airbyte_cdk.sources import Source
+
+
+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
+
+        if uncaught_exception:
+            self._messages.append(
+                assemble_uncaught_exception(
+                    type(uncaught_exception), uncaught_exception
+                ).as_airbyte_message()
+            )
+
+    @staticmethod
+    def _parse_message(message: str) -> AirbyteMessage:
+        try:
+            return AirbyteMessageSerializer.load(orjson.loads(message))
+        except (orjson.JSONDecodeError, SchemaValidationError):
+            # The platform assumes that logs that are not of AirbyteMessage format are log messages
+            return AirbyteMessage(
+                type=Type.LOG, log=AirbyteLogMessage(level=Level.INFO, message=message)
+            )
+
+    @property
+    def records_and_state_messages(self) -> List[AirbyteMessage]:
+        return self._get_message_by_types([Type.RECORD, Type.STATE])
+
+    @property
+    def records(self) -> List[AirbyteMessage]:
+        return self._get_message_by_types([Type.RECORD])
+
+    @property
+    def state_messages(self) -> List[AirbyteMessage]:
+        return self._get_message_by_types([Type.STATE])
+
+    @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`
+
+    @property
+    def logs(self) -> List[AirbyteMessage]:
+        return self._get_message_by_types([Type.LOG])
+
+    @property
+    def trace_messages(self) -> List[AirbyteMessage]:
+        return self._get_message_by_types([Type.TRACE])
+
+    @property
+    def analytics_messages(self) -> List[AirbyteMessage]:
+        return self._get_trace_message_by_trace_type(TraceType.ANALYTICS)
+
+    @property
+    def errors(self) -> List[AirbyteMessage]:
+        return self._get_trace_message_by_trace_type(TraceType.ERROR)
+
+    @property
+    def catalog(self) -> AirbyteMessage:
+        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]
+
+    def get_stream_statuses(self, stream_name: str) -> List[AirbyteStreamStatus]:
+        status_messages = map(
+            lambda message: message.trace.stream_status.status,  # type: ignore
+            filter(
+                lambda message: message.trace.stream_status.stream_descriptor.name == stream_name,  # type: ignore # callable; trace has `stream_status`
+                self._get_trace_message_by_trace_type(TraceType.STREAM_STATUS),
+            ),
+        )
+        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_trace_message_by_trace_type(self, trace_type: TraceType) -> List[AirbyteMessage]:
+        return [
+            message
+            for message in self._get_message_by_types([Type.TRACE])
+            if message.trace.type == trace_type  # type: ignore[union-attr] # trace has `type`
+        ]
+
+    def is_in_logs(self, pattern: str) -> bool:
+        """Check if any log message case-insensitive matches the pattern."""
+        return any(
+            re.search(
+                pattern,
+                entry.log.message,  # type: ignore[union-attr] # log has `message`
+                flags=re.IGNORECASE,
+            )
+            for entry in self.logs
+        )
+
+    def is_not_in_logs(self, pattern: str) -> bool:
+        """Check if no log message matches the case-insensitive pattern."""
+        return not self.is_in_logs(pattern)
+
+
+def _run_command(
+    source: Source, args: List[str], expecting_exception: bool = False
+) -> EntrypointOutput:
+    log_capture_buffer = StringIO()
+    stream_handler = logging.StreamHandler(log_capture_buffer)
+    stream_handler.setLevel(logging.INFO)
+    stream_handler.setFormatter(AirbyteLogFormatter())
+    parent_logger = logging.getLogger("")
+    parent_logger.addHandler(stream_handler)
+
+    parsed_args = AirbyteEntrypoint.parse_args(args)
+
+    source_entrypoint = AirbyteEntrypoint(source)
+    messages = []
+    uncaught_exception = None
+    try:
+        for message in source_entrypoint.run(parsed_args):
+            messages.append(message)
+    except Exception as exception:
+        if not expecting_exception:
+            print("Printing unexpected error from entrypoint_wrapper")
+            print("".join(traceback.format_exception(None, exception, exception.__traceback__)))
+        uncaught_exception = exception
+
+    captured_logs = log_capture_buffer.getvalue().split("\n")[:-1]
+
+    parent_logger.removeHandler(stream_handler)
+
+    return EntrypointOutput(messages + captured_logs, uncaught_exception)
+
+
+def discover(
+    source: Source,
+    config: Mapping[str, Any],
+    expecting_exception: bool = False,
+) -> EntrypointOutput:
+    """
+    config must be json serializable
+    :param expecting_exception: By default if there is an uncaught exception, the exception will be printed out. If this is expected, please
+        provide expecting_exception=True so that the test output logs are cleaner
+    """
+
+    with tempfile.TemporaryDirectory() as tmp_directory:
+        tmp_directory_path = Path(tmp_directory)
+        config_file = make_file(tmp_directory_path / "config.json", config)
+
+        return _run_command(
+            source, ["discover", "--config", config_file, "--debug"], expecting_exception
+        )
+
+
+def read(
+    source: Source,
+    config: Mapping[str, Any],
+    catalog: ConfiguredAirbyteCatalog,
+    state: Optional[List[AirbyteStateMessage]] = None,
+    expecting_exception: bool = False,
+) -> EntrypointOutput:
+    """
+    config and state must be json serializable
+
+    :param expecting_exception: By default if there is an uncaught exception, the exception will be printed out. If this is expected, please
+        provide expecting_exception=True so that the test output logs are cleaner
+    """
+    with tempfile.TemporaryDirectory() as tmp_directory:
+        tmp_directory_path = Path(tmp_directory)
+        config_file = make_file(tmp_directory_path / "config.json", config)
+        catalog_file = make_file(
+            tmp_directory_path / "catalog.json",
+            orjson.dumps(ConfiguredAirbyteCatalogSerializer.dump(catalog)).decode(),
+        )
+        args = [
+            "read",
+            "--config",
+            config_file,
+            "--catalog",
+            catalog_file,
+        ]
+        if state is not None:
+            args.extend(
+                [
+                    "--state",
+                    make_file(
+                        tmp_directory_path / "state.json",
+                        f"[{','.join([orjson.dumps(AirbyteStateMessageSerializer.dump(stream_state)).decode() for stream_state in state])}]",
+                    ),
+                ]
+            )
+
+        return _run_command(source, args, expecting_exception)
+
+
+def make_file(
+    path: Path, file_contents: Optional[Union[str, Mapping[str, Any], List[Mapping[str, Any]]]]
+) -> str:
+    if isinstance(file_contents, str):
+        path.write_text(file_contents)
+    else:
+        path.write_text(json.dumps(file_contents))
+    return str(path)

airbyte_cdk/test/mock_http/__init__.py

@@ -0,0 +1,6 @@
+from airbyte_cdk.test.mock_http.matcher import HttpRequestMatcher
+from airbyte_cdk.test.mock_http.mocker import HttpMocker
+from airbyte_cdk.test.mock_http.request import HttpRequest
+from airbyte_cdk.test.mock_http.response import HttpResponse
+
+__all__ = ["HttpMocker", "HttpRequest", "HttpRequestMatcher", "HttpResponse"]

airbyte_cdk/test/mock_http/matcher.py

@@ -0,0 +1,41 @@
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+from typing import Any
+
+from airbyte_cdk.test.mock_http.request import HttpRequest
+
+
+class HttpRequestMatcher:
+    def __init__(self, request: HttpRequest, minimum_number_of_expected_match: int):
+        self._request_to_match = request
+        self._minimum_number_of_expected_match = minimum_number_of_expected_match
+        self._actual_number_of_matches = 0
+
+    def matches(self, request: HttpRequest) -> bool:
+        hit = request.matches(self._request_to_match)
+        if hit:
+            self._actual_number_of_matches += 1
+        return hit
+
+    def has_expected_match_count(self) -> bool:
+        return self._actual_number_of_matches >= self._minimum_number_of_expected_match
+
+    @property
+    def actual_number_of_matches(self) -> int:
+        return self._actual_number_of_matches
+
+    @property
+    def request(self) -> HttpRequest:
+        return self._request_to_match
+
+    def __str__(self) -> str:
+        return (
+            f"HttpRequestMatcher("
+            f"request_to_match={self._request_to_match}, "
+            f"minimum_number_of_expected_match={self._minimum_number_of_expected_match}, "
+            f"actual_number_of_matches={self._actual_number_of_matches})"
+        )
+
+    def __eq__(self, other: Any) -> bool:
+        if isinstance(other, HttpRequestMatcher):
+            return self._request_to_match == other._request_to_match
+        return False