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

airbyte_cdk/sources/concurrent_source/thread_pool_manager.py

@@ -0,0 +1,115 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+import logging
+import threading
+from concurrent.futures import Future, ThreadPoolExecutor
+from typing import Any, Callable, List, Optional
+
+
+class ThreadPoolManager:
+    """
+    Wrapper to abstract away the threadpool and the logic to wait for pending tasks to be completed.
+    """
+
+    DEFAULT_MAX_QUEUE_SIZE = 10_000
+
+    def __init__(
+        self,
+        threadpool: ThreadPoolExecutor,
+        logger: logging.Logger,
+        max_concurrent_tasks: int = DEFAULT_MAX_QUEUE_SIZE,
+    ):
+        """
+        :param threadpool: The threadpool to use
+        :param logger: The logger to use
+        :param max_concurrent_tasks: The maximum number of tasks that can be pending at the same time
+        """
+        self._threadpool = threadpool
+        self._logger = logger
+        self._max_concurrent_tasks = max_concurrent_tasks
+        self._futures: List[Future[Any]] = []
+        self._lock = threading.Lock()
+        self._most_recently_seen_exception: Optional[Exception] = None
+
+        self._logging_threshold = max_concurrent_tasks * 2
+
+    def prune_to_validate_has_reached_futures_limit(self) -> bool:
+        self._prune_futures(self._futures)
+        if len(self._futures) > self._logging_threshold:
+            self._logger.warning(
+                f"ThreadPoolManager: The list of futures is getting bigger than expected ({len(self._futures)})"
+            )
+        return len(self._futures) >= self._max_concurrent_tasks
+
+    def submit(self, function: Callable[..., Any], *args: Any) -> None:
+        self._futures.append(self._threadpool.submit(function, *args))
+
+    def _prune_futures(self, futures: List[Future[Any]]) -> None:
+        """
+        Take a list in input and remove the futures that are completed. If a future has an exception, it'll raise and kill the stream
+        operation.
+
+        We are using a lock here as without it, the algorithm would not be thread safe
+        """
+        with self._lock:
+            if len(futures) < self._max_concurrent_tasks:
+                return
+
+            for index in reversed(range(len(futures))):
+                future = futures[index]
+
+                if future.done():
+                    # Only call future.exception() if the future is known to be done because it will block until the future is done.
+                    # See https://docs.python.org/3/library/concurrent.futures.html#concurrent.futures.Future.exception
+                    optional_exception = future.exception()
+                    if optional_exception:
+                        # Exception handling should be done in the main thread. Hence, we only store the exception and expect the main
+                        # thread to call raise_if_exception
+                        # We do not expect this error to happen. The futures created during concurrent syncs should catch the exception and
+                        # push it to the queue. If this exception occurs, please review the futures and how they handle exceptions.
+                        self._most_recently_seen_exception = RuntimeError(
+                            f"Failed processing a future: {optional_exception}. Please contact the Airbyte team."
+                        )
+                    futures.pop(index)
+
+    def _shutdown(self) -> None:
+        # Without a way to stop the threads that have already started, this will not stop the Python application. We are fine today with
+        # this imperfect approach because we only do this in case of `self._most_recently_seen_exception` which we don't expect to happen
+        self._threadpool.shutdown(wait=False, cancel_futures=True)
+
+    def is_done(self) -> bool:
+        return all([f.done() for f in self._futures])
+
+    def check_for_errors_and_shutdown(self) -> None:
+        """
+        Check if any of the futures have an exception, and raise it if so. If all futures are done, shutdown the threadpool.
+        If the futures are not done, raise an exception.
+        :return:
+        """
+        if self._most_recently_seen_exception:
+            self._logger.exception(
+                "An unknown exception has occurred while reading concurrently",
+                exc_info=self._most_recently_seen_exception,
+            )
+            self._stop_and_raise_exception(self._most_recently_seen_exception)
+
+        exceptions_from_futures = [
+            f for f in [future.exception() for future in self._futures] if f is not None
+        ]
+        if exceptions_from_futures:
+            exception = RuntimeError(f"Failed reading with errors: {exceptions_from_futures}")
+            self._stop_and_raise_exception(exception)
+        else:
+            futures_not_done = [f for f in self._futures if not f.done()]
+            if futures_not_done:
+                exception = RuntimeError(
+                    f"Failed reading with futures not done: {futures_not_done}"
+                )
+                self._stop_and_raise_exception(exception)
+            else:
+                self._shutdown()
+
+    def _stop_and_raise_exception(self, exception: BaseException) -> None:
+        self._shutdown()
+        raise exception

airbyte_cdk/sources/config.py

@@ -0,0 +1,27 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+from typing import Any, Dict
+
+from pydantic.v1 import BaseModel
+
+from airbyte_cdk.sources.utils.schema_helpers import expand_refs, rename_key
+
+
+class BaseConfig(BaseModel):
+    """Base class for connector spec, adds the following behaviour:
+
+    - resolve $ref and replace it with definition
+    - replace all occurrences of anyOf with oneOf
+    - drop description
+    """
+
+    @classmethod
+    def schema(cls, *args: Any, **kwargs: Any) -> Dict[str, Any]:
+        """We're overriding the schema classmethod to enable some post-processing"""
+        schema = super().schema(*args, **kwargs)
+        rename_key(schema, old_key="anyOf", new_key="oneOf")  # UI supports only oneOf
+        expand_refs(schema)
+        schema.pop("description", None)  # description added from the docstring
+        return schema

airbyte_cdk/sources/connector_state_manager.py

@@ -0,0 +1,161 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+import copy
+from dataclasses import dataclass
+from typing import Any, List, Mapping, MutableMapping, Optional, Tuple, Union, cast
+
+from airbyte_cdk.models import (
+    AirbyteMessage,
+    AirbyteStateBlob,
+    AirbyteStateMessage,
+    AirbyteStateType,
+    AirbyteStreamState,
+    StreamDescriptor,
+)
+from airbyte_cdk.models import Type as MessageType
+from airbyte_cdk.models.airbyte_protocol import AirbyteGlobalState, AirbyteStateBlob
+
+
+@dataclass(frozen=True)
+class HashableStreamDescriptor:
+    """
+    Helper class that overrides the existing StreamDescriptor class that is auto generated from the Airbyte Protocol and
+    freezes its fields so that it be used as a hash key. This is only marked public because we use it outside for unit tests.
+    """
+
+    name: str
+    namespace: Optional[str] = None
+
+
+class ConnectorStateManager:
+    """
+    ConnectorStateManager consolidates the various forms of a stream's incoming state message (STREAM / GLOBAL) under a common
+    interface. It also provides methods to extract and update state
+    """
+
+    def __init__(self, state: Optional[List[AirbyteStateMessage]] = None):
+        shared_state, per_stream_states = self._extract_from_state_message(state)
+
+        # We explicitly throw an error if we receive a GLOBAL state message that contains a shared_state because API sources are
+        # designed to checkpoint state independently of one another. API sources should never be emitting a state message where
+        # shared_state is populated. Rather than define how to handle shared_state without a clear use case, we're opting to throw an
+        # error instead and if/when we find one, we will then implement processing of the shared_state value.
+        if shared_state:
+            raise ValueError(
+                "Received a GLOBAL AirbyteStateMessage that contains a shared_state. This library only ever generates per-STREAM "
+                "STATE messages so this was not generated by this connector. This must be an orchestrator or platform error. GLOBAL "
+                "state messages with shared_state will not be processed correctly. "
+            )
+        self.per_stream_states = per_stream_states
+
+    def get_stream_state(
+        self, stream_name: str, namespace: Optional[str]
+    ) -> MutableMapping[str, Any]:
+        """
+        Retrieves the state of a given stream based on its descriptor (name + namespace).
+        :param stream_name: Name of the stream being fetched
+        :param namespace: Namespace of the stream being fetched
+        :return: The per-stream state for a stream
+        """
+        stream_state: AirbyteStateBlob | None = self.per_stream_states.get(
+            HashableStreamDescriptor(name=stream_name, namespace=namespace)
+        )
+        if stream_state:
+            return copy.deepcopy({k: v for k, v in stream_state.__dict__.items()})
+        return {}
+
+    def update_state_for_stream(
+        self, stream_name: str, namespace: Optional[str], value: Mapping[str, Any]
+    ) -> None:
+        """
+        Overwrites the state blob of a specific stream based on the provided stream name and optional namespace
+        :param stream_name: The name of the stream whose state is being updated
+        :param namespace: The namespace of the stream if it exists
+        :param value: A stream state mapping that is being updated for a stream
+        """
+        stream_descriptor = HashableStreamDescriptor(name=stream_name, namespace=namespace)
+        self.per_stream_states[stream_descriptor] = AirbyteStateBlob(value)
+
+    def create_state_message(self, stream_name: str, namespace: Optional[str]) -> AirbyteMessage:
+        """
+        Generates an AirbyteMessage using the current per-stream state of a specified stream
+        :param stream_name: The name of the stream for the message that is being created
+        :param namespace: The namespace of the stream for the message that is being created
+        :return: The Airbyte state message to be emitted by the connector during a sync
+        """
+        hashable_descriptor = HashableStreamDescriptor(name=stream_name, namespace=namespace)
+        stream_state = self.per_stream_states.get(hashable_descriptor) or AirbyteStateBlob()
+
+        return AirbyteMessage(
+            type=MessageType.STATE,
+            state=AirbyteStateMessage(
+                type=AirbyteStateType.STREAM,
+                stream=AirbyteStreamState(
+                    stream_descriptor=StreamDescriptor(name=stream_name, namespace=namespace),
+                    stream_state=stream_state,
+                ),
+            ),
+        )
+
+    @classmethod
+    def _extract_from_state_message(
+        cls,
+        state: Optional[List[AirbyteStateMessage]],
+    ) -> Tuple[
+        Optional[AirbyteStateBlob],
+        MutableMapping[HashableStreamDescriptor, Optional[AirbyteStateBlob]],
+    ]:
+        """
+        Takes an incoming list of state messages or a global state message and extracts state attributes according to
+        type which can then be assigned to the new state manager being instantiated
+        :param state: The incoming state input
+        :return: A tuple of shared state and per stream state assembled from the incoming state list
+        """
+        if state is None:
+            return None, {}
+
+        is_global = cls._is_global_state(state)
+
+        if is_global:
+            # We already validate that this is a global state message, not None:
+            global_state = cast(AirbyteGlobalState, state[0].global_)
+            # global_state has shared_state, also not None:
+            shared_state: AirbyteStateBlob = cast(
+                AirbyteStateBlob, copy.deepcopy(global_state.shared_state, {})
+            )
+            streams = {
+                HashableStreamDescriptor(
+                    name=per_stream_state.stream_descriptor.name,
+                    namespace=per_stream_state.stream_descriptor.namespace,
+                ): per_stream_state.stream_state
+                for per_stream_state in global_state.stream_states  # type: ignore[union-attr] # global_state has shared_state
+            }
+            return shared_state, streams
+        else:
+            streams = {
+                HashableStreamDescriptor(
+                    name=per_stream_state.stream.stream_descriptor.name,  # type: ignore[union-attr] # stream has stream_descriptor
+                    namespace=per_stream_state.stream.stream_descriptor.namespace,  # type: ignore[union-attr] # stream has stream_descriptor
+                ): per_stream_state.stream.stream_state  # type: ignore[union-attr] # stream has stream_state
+                for per_stream_state in state
+                if per_stream_state.type == AirbyteStateType.STREAM
+                and hasattr(per_stream_state, "stream")  # type: ignore # state is always a list of AirbyteStateMessage if is_per_stream is True
+            }
+            return None, streams
+
+    @staticmethod
+    def _is_global_state(state: Union[List[AirbyteStateMessage], MutableMapping[str, Any]]) -> bool:
+        return (
+            isinstance(state, List)
+            and len(state) == 1
+            and isinstance(state[0], AirbyteStateMessage)
+            and state[0].type == AirbyteStateType.GLOBAL
+        )
+
+    @staticmethod
+    def _is_per_stream_state(
+        state: Union[List[AirbyteStateMessage], MutableMapping[str, Any]],
+    ) -> bool:
+        return isinstance(state, List)

airbyte_cdk/sources/declarative/__init__.py

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

airbyte_cdk/sources/declarative/async_job/job.py

@@ -0,0 +1,52 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+
+
+from datetime import timedelta
+from typing import Optional
+
+from airbyte_cdk.sources.declarative.async_job.timer import Timer
+from airbyte_cdk.sources.types import StreamSlice
+
+from .status import AsyncJobStatus
+
+
+class AsyncJob:
+    """
+    Description of an API job.
+
+    Note that the timer will only stop once `update_status` is called so the job might be completed on the API side but until we query for
+    it and call `ApiJob.update_status`, `ApiJob.status` will not reflect the actual API side status.
+    """
+
+    def __init__(
+        self, api_job_id: str, job_parameters: StreamSlice, timeout: Optional[timedelta] = None
+    ) -> None:
+        self._api_job_id = api_job_id
+        self._job_parameters = job_parameters
+        self._status = AsyncJobStatus.RUNNING
+
+        timeout = timeout if timeout else timedelta(minutes=60)
+        self._timer = Timer(timeout)
+        self._timer.start()
+
+    def api_job_id(self) -> str:
+        return self._api_job_id
+
+    def status(self) -> AsyncJobStatus:
+        if self._timer.has_timed_out():
+            return AsyncJobStatus.TIMED_OUT
+        return self._status
+
+    def job_parameters(self) -> StreamSlice:
+        return self._job_parameters
+
+    def update_status(self, status: AsyncJobStatus) -> None:
+        if self._status != AsyncJobStatus.RUNNING and status == AsyncJobStatus.RUNNING:
+            self._timer.start()
+        elif status.is_terminal():
+            self._timer.stop()
+
+        self._status = status
+
+    def __repr__(self) -> str:
+        return f"AsyncJob(api_job_id={self.api_job_id()}, job_parameters={self.job_parameters()}, status={self.status()})"