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

airbyte_cdk/sources/message/repository.py

@@ -0,0 +1,137 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+import json
+import logging
+from abc import ABC, abstractmethod
+from collections import deque
+from typing import Callable, Deque, Iterable, List, Optional
+
+from airbyte_cdk.models import AirbyteLogMessage, AirbyteMessage, Level, Type
+from airbyte_cdk.sources.utils.types import JsonType
+from airbyte_cdk.utils.airbyte_secrets_utils import filter_secrets
+
+_LOGGER = logging.getLogger("MessageRepository")
+_SUPPORTED_MESSAGE_TYPES = {Type.CONTROL, Type.LOG}
+LogMessage = dict[str, JsonType]
+
+_SEVERITY_BY_LOG_LEVEL = {
+    Level.FATAL: 1,
+    Level.ERROR: 2,
+    Level.WARN: 3,
+    Level.INFO: 4,
+    Level.DEBUG: 5,
+    Level.TRACE: 5,
+}
+
+
+def _is_severe_enough(threshold: Level, level: Level) -> bool:
+    if threshold not in _SEVERITY_BY_LOG_LEVEL:
+        _LOGGER.warning(
+            f"Log level {threshold} for threshold is not supported. This is probably a CDK bug. Please contact Airbyte."
+        )
+        return True
+
+    if level not in _SEVERITY_BY_LOG_LEVEL:
+        _LOGGER.warning(
+            f"Log level {level} is not supported. This is probably a source bug. Please contact the owner of the source or Airbyte."
+        )
+        return True
+
+    return _SEVERITY_BY_LOG_LEVEL[threshold] >= _SEVERITY_BY_LOG_LEVEL[level]
+
+
+class MessageRepository(ABC):
+    @abstractmethod
+    def emit_message(self, message: AirbyteMessage) -> None:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def log_message(self, level: Level, message_provider: Callable[[], LogMessage]) -> None:
+        """
+        Computing messages can be resource consuming. This method is specialized for logging because we want to allow for lazy evaluation if
+        the log level is less severe than what is configured
+        """
+        raise NotImplementedError()
+
+    @abstractmethod
+    def consume_queue(self) -> Iterable[AirbyteMessage]:
+        raise NotImplementedError()
+
+
+class NoopMessageRepository(MessageRepository):
+    def emit_message(self, message: AirbyteMessage) -> None:
+        pass
+
+    def log_message(self, level: Level, message_provider: Callable[[], LogMessage]) -> None:
+        pass
+
+    def consume_queue(self) -> Iterable[AirbyteMessage]:
+        return []
+
+
+class InMemoryMessageRepository(MessageRepository):
+    def __init__(self, log_level: Level = Level.INFO) -> None:
+        self._message_queue: Deque[AirbyteMessage] = deque()
+        self._log_level = log_level
+
+    def emit_message(self, message: AirbyteMessage) -> None:
+        self._message_queue.append(message)
+
+    def log_message(self, level: Level, message_provider: Callable[[], LogMessage]) -> None:
+        if _is_severe_enough(self._log_level, level):
+            self.emit_message(
+                AirbyteMessage(
+                    type=Type.LOG,
+                    log=AirbyteLogMessage(
+                        level=level, message=filter_secrets(json.dumps(message_provider()))
+                    ),
+                )
+            )
+
+    def consume_queue(self) -> Iterable[AirbyteMessage]:
+        while self._message_queue:
+            yield self._message_queue.popleft()
+
+
+class LogAppenderMessageRepositoryDecorator(MessageRepository):
+    def __init__(
+        self,
+        dict_to_append: LogMessage,
+        decorated: MessageRepository,
+        log_level: Level = Level.INFO,
+    ):
+        self._dict_to_append = dict_to_append
+        self._decorated = decorated
+        self._log_level = log_level
+
+    def emit_message(self, message: AirbyteMessage) -> None:
+        self._decorated.emit_message(message)
+
+    def log_message(self, level: Level, message_provider: Callable[[], LogMessage]) -> None:
+        if _is_severe_enough(self._log_level, level):
+            message = message_provider()
+            self._append_second_to_first(message, self._dict_to_append)
+            self._decorated.log_message(level, lambda: message)
+
+    def consume_queue(self) -> Iterable[AirbyteMessage]:
+        return self._decorated.consume_queue()
+
+    def _append_second_to_first(
+        self, first: LogMessage, second: LogMessage, path: Optional[List[str]] = None
+    ) -> LogMessage:
+        if path is None:
+            path = []
+
+        for key in second:
+            if key in first:
+                if isinstance(first[key], dict) and isinstance(second[key], dict):
+                    self._append_second_to_first(first[key], second[key], path + [str(key)])  # type: ignore # type is verified above
+                else:
+                    if first[key] != second[key]:
+                        _LOGGER.warning("Conflict at %s" % ".".join(path + [str(key)]))
+                    first[key] = second[key]
+            else:
+                first[key] = second[key]
+        return first

airbyte_cdk/sources/source.py

@@ -0,0 +1,95 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+
+import logging
+from abc import ABC, abstractmethod
+from typing import Any, Generic, Iterable, List, Mapping, Optional, TypeVar
+
+from airbyte_cdk.connector import BaseConnector, DefaultConnectorMixin, TConfig
+from airbyte_cdk.models import (
+    AirbyteCatalog,
+    AirbyteMessage,
+    AirbyteStateMessage,
+    AirbyteStateMessageSerializer,
+    ConfiguredAirbyteCatalog,
+    ConfiguredAirbyteCatalogSerializer,
+)
+
+TState = TypeVar("TState")
+TCatalog = TypeVar("TCatalog")
+
+
+class ExperimentalClassWarning(DeprecationWarning):
+    pass
+
+
+class BaseSource(BaseConnector[TConfig], ABC, Generic[TConfig, TState, TCatalog]):
+    @abstractmethod
+    def read_state(self, state_path: str) -> TState: ...
+
+    @abstractmethod
+    def read_catalog(self, catalog_path: str) -> TCatalog: ...
+
+    @abstractmethod
+    def read(
+        self,
+        logger: logging.Logger,
+        config: TConfig,
+        catalog: TCatalog,
+        state: Optional[TState] = None,
+    ) -> Iterable[AirbyteMessage]:
+        """
+        Returns a generator of the AirbyteMessages generated by reading the source with the given configuration, catalog, and state.
+        """
+
+    @abstractmethod
+    def discover(self, logger: logging.Logger, config: TConfig) -> AirbyteCatalog:
+        """
+        Returns an AirbyteCatalog representing the available streams and fields in this integration. For example, given valid credentials to a
+        Postgres database, returns an Airbyte catalog where each postgres table is a stream, and each table column is a field.
+        """
+
+
+class Source(
+    DefaultConnectorMixin,
+    BaseSource[Mapping[str, Any], List[AirbyteStateMessage], ConfiguredAirbyteCatalog],
+    ABC,
+):
+    # can be overridden to change an input state.
+    @classmethod
+    def read_state(cls, state_path: str) -> List[AirbyteStateMessage]:
+        """
+        Retrieves the input state of a sync by reading from the specified JSON file. Incoming state can be deserialized into either
+        a JSON object for legacy state input or as a list of AirbyteStateMessages for the per-stream state format. Regardless of the
+        incoming input type, it will always be transformed and output as a list of AirbyteStateMessage(s).
+        :param state_path: The filepath to where the stream states are located
+        :return: The complete stream state based on the connector's previous sync
+        """
+        parsed_state_messages = []
+        if state_path:
+            state_obj = BaseConnector._read_json_file(state_path)
+            if state_obj:
+                for state in state_obj:  # type: ignore  # `isinstance(state_obj, List)` ensures that this is a list
+                    parsed_message = AirbyteStateMessageSerializer.load(state)
+                    if (
+                        not parsed_message.stream
+                        and not parsed_message.data
+                        and not parsed_message.global_
+                    ):
+                        raise ValueError(
+                            "AirbyteStateMessage should contain either a stream, global, or state field"
+                        )
+                    parsed_state_messages.append(parsed_message)
+        return parsed_state_messages
+
+    # can be overridden to change an input catalog
+    @classmethod
+    def read_catalog(cls, catalog_path: str) -> ConfiguredAirbyteCatalog:
+        return ConfiguredAirbyteCatalogSerializer.load(cls._read_json_file(catalog_path))
+
+    @property
+    def name(self) -> str:
+        """Source name"""
+        return self.__class__.__name__

airbyte_cdk/sources/specs/transfer_modes.py

@@ -0,0 +1,26 @@
+#
+# Copyright (c) 2025 Airbyte, Inc., all rights reserved.
+#
+
+from typing import Literal
+
+from pydantic.v1 import AnyUrl, BaseModel, Field
+
+from airbyte_cdk import OneOfOptionConfig
+
+
+class DeliverPermissions(BaseModel):
+    class Config(OneOfOptionConfig):
+        title = "Replicate Permissions ACL"
+        description = "Sends one identity stream and one for more permissions (ACL) streams to the destination. This data can be used in downstream systems to recreate permission restrictions mirroring the original source."
+        discriminator = "delivery_type"
+
+    delivery_type: Literal["use_permissions_transfer"] = Field(
+        "use_permissions_transfer", const=True
+    )
+
+    include_identities_stream: bool = Field(
+        title="Include Identity Stream",
+        description="This data can be used in downstream systems to recreate permission restrictions mirroring the original source",
+        default=True,
+    )

airbyte_cdk/sources/streams/__init__.py

@@ -0,0 +1,8 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+# Initialize Streams Package
+from .core import NO_CURSOR_STATE_KEY, CheckpointMixin, IncrementalMixin, Stream
+
+__all__ = ["NO_CURSOR_STATE_KEY", "IncrementalMixin", "CheckpointMixin", "Stream"]

airbyte_cdk/sources/streams/availability_strategy.py

@@ -0,0 +1,84 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+import logging
+import typing
+from abc import ABC, abstractmethod
+from typing import Any, Mapping, Optional, Tuple
+
+from airbyte_cdk.models import SyncMode
+from airbyte_cdk.sources.streams.core import Stream, StreamData
+
+if typing.TYPE_CHECKING:
+    from airbyte_cdk.sources import Source
+
+
+class AvailabilityStrategy(ABC):
+    """
+    Abstract base class for checking stream availability.
+    """
+
+    @abstractmethod
+    def check_availability(
+        self, stream: Stream, logger: logging.Logger, source: Optional["Source"] = None
+    ) -> Tuple[bool, Optional[str]]:
+        """
+        Checks stream availability.
+
+        :param stream: stream
+        :param logger: source logger
+        :param source: (optional) source
+        :return: A tuple of (boolean, str). If boolean is true, then the stream
+          is available, and no str is required. Otherwise, the stream is unavailable
+          for some reason and the str should describe what went wrong and how to
+          resolve the unavailability, if possible.
+        """
+
+    @staticmethod
+    def get_first_stream_slice(stream: Stream) -> Optional[Mapping[str, Any]]:
+        """
+        Gets the first stream_slice from a given stream's stream_slices.
+        :param stream: stream
+        :raises StopIteration: if there is no first slice to return (the stream_slices generator is empty)
+        :return: first stream slice from 'stream_slices' generator (`None` is a valid stream slice)
+        """
+        # We wrap the return output of stream_slices() because some implementations return types that are iterable,
+        # but not iterators such as lists or tuples
+        slices = iter(
+            stream.stream_slices(
+                cursor_field=stream.cursor_field,  # type: ignore[arg-type]
+                sync_mode=SyncMode.full_refresh,
+            )
+        )
+        return next(slices)
+
+    @staticmethod
+    def get_first_record_for_slice(
+        stream: Stream, stream_slice: Optional[Mapping[str, Any]]
+    ) -> StreamData:
+        """
+        Gets the first record for a stream_slice of a stream.
+
+        :param stream: stream instance from which to read records
+        :param stream_slice: stream_slice parameters for slicing the stream
+        :raises StopIteration: if there is no first record to return (the read_records generator is empty)
+        :return: StreamData containing the first record in the slice
+        """
+        # Store the original value of exit_on_rate_limit
+        original_exit_on_rate_limit = stream.exit_on_rate_limit
+
+        try:
+            # Ensure exit_on_rate_limit is safely set to True if possible
+            stream.exit_on_rate_limit = True
+
+            # We wrap the return output of read_records() because some implementations return types that are iterable,
+            # but not iterators such as lists or tuples
+            records_for_slice = iter(
+                stream.read_records(sync_mode=SyncMode.full_refresh, stream_slice=stream_slice)
+            )
+
+            return next(records_for_slice)
+        finally:
+            # Restore the original exit_on_rate_limit value
+            stream.exit_on_rate_limit = original_exit_on_rate_limit