airbyte-cdk 0.72.1__py3-none-any.whl → 6.17.1.dev0__py3-none-any.whl

airbyte_cdk/sources/concurrent_source/concurrent_source.py

@@ -8,15 +8,21 @@ from typing import Iterable, Iterator, List
 
 from airbyte_cdk.models import AirbyteMessage
 from airbyte_cdk.sources.concurrent_source.concurrent_read_processor import ConcurrentReadProcessor
-from airbyte_cdk.sources.concurrent_source.partition_generation_completed_sentinel import PartitionGenerationCompletedSentinel
+from airbyte_cdk.sources.concurrent_source.partition_generation_completed_sentinel import (
+    PartitionGenerationCompletedSentinel,
+)
+from airbyte_cdk.sources.concurrent_source.stream_thread_exception import StreamThreadException
 from airbyte_cdk.sources.concurrent_source.thread_pool_manager import ThreadPoolManager
 from airbyte_cdk.sources.message import InMemoryMessageRepository, MessageRepository
 from airbyte_cdk.sources.streams.concurrent.abstract_stream import AbstractStream
 from airbyte_cdk.sources.streams.concurrent.partition_enqueuer import PartitionEnqueuer
 from airbyte_cdk.sources.streams.concurrent.partition_reader import PartitionReader
 from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
-from airbyte_cdk.sources.streams.concurrent.partitions.record import Record
-from airbyte_cdk.sources.streams.concurrent.partitions.types import PartitionCompleteSentinel, QueueItem
+from airbyte_cdk.sources.streams.concurrent.partitions.types import (
+    PartitionCompleteSentinel,
+    QueueItem,
+)
+from airbyte_cdk.sources.types import Record
 from airbyte_cdk.sources.utils.slice_logger import DebugSliceLogger, SliceLogger
 
 
@@ -40,14 +46,25 @@ class ConcurrentSource:
         timeout_seconds: int = DEFAULT_TIMEOUT_SECONDS,
     ) -> "ConcurrentSource":
         is_single_threaded = initial_number_of_partitions_to_generate == 1 and num_workers == 1
-        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"
+        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"
         threadpool = ThreadPoolManager(
-            concurrent.futures.ThreadPoolExecutor(max_workers=num_workers, thread_name_prefix="workerpool"),
+            concurrent.futures.ThreadPoolExecutor(
+                max_workers=num_workers, thread_name_prefix="workerpool"
+            ),
             logger,
         )
         return ConcurrentSource(
-            threadpool, logger, slice_logger, message_repository, initial_number_of_partitions_to_generate, timeout_seconds
+            threadpool,
+            logger,
+            slice_logger,
+            message_repository,
+            initial_number_of_partitions_to_generate,
+            timeout_seconds,
         )
 
     def __init__(
@@ -79,11 +96,6 @@ class ConcurrentSource:
         streams: List[AbstractStream],
     ) -> Iterator[AirbyteMessage]:
         self._logger.info("Starting syncing")
-        stream_instances_to_read_from = self._get_streams_to_read_from(streams)
-
-        # Return early if there are no streams to read from
-        if not stream_instances_to_read_from:
-            return
 
         # We set a maxsize to for the main thread to process record items when the queue size grows. This assumes that there are less
         # threads generating partitions that than are max number of workers. If it weren't the case, we could have threads only generating
@@ -91,7 +103,7 @@ class ConcurrentSource:
         # information and might even need to be configurable depending on the source
         queue: Queue[QueueItem] = Queue(maxsize=10_000)
         concurrent_stream_processor = ConcurrentReadProcessor(
-            stream_instances_to_read_from,
+            streams,
             PartitionEnqueuer(queue, self._threadpool),
             self._threadpool,
             self._logger,
@@ -111,7 +123,9 @@ class ConcurrentSource:
         self._threadpool.check_for_errors_and_shutdown()
         self._logger.info("Finished syncing")
 
-    def _submit_initial_partition_generators(self, concurrent_stream_processor: ConcurrentReadProcessor) -> Iterable[AirbyteMessage]:
+    def _submit_initial_partition_generators(
+        self, concurrent_stream_processor: ConcurrentReadProcessor
+    ) -> Iterable[AirbyteMessage]:
         for _ in range(self._initial_number_partitions_to_generate):
             status_message = concurrent_stream_processor.start_next_partition_generator()
             if status_message:
@@ -123,11 +137,6 @@ class ConcurrentSource:
         concurrent_stream_processor: ConcurrentReadProcessor,
     ) -> Iterable[AirbyteMessage]:
         while airbyte_message_or_record_or_exception := queue.get():
-            try:
-                self._threadpool.shutdown_if_exception()
-            except Exception as exception:
-                concurrent_stream_processor.on_exception(exception)
-
             yield from self._handle_item(
                 airbyte_message_or_record_or_exception,
                 concurrent_stream_processor,
@@ -142,7 +151,7 @@ class ConcurrentSource:
         concurrent_stream_processor: ConcurrentReadProcessor,
     ) -> Iterable[AirbyteMessage]:
         # handle queue item and call the appropriate handler depending on the type of the queue item
-        if isinstance(queue_item, Exception):
+        if isinstance(queue_item, StreamThreadException):
             yield from concurrent_stream_processor.on_exception(queue_item)
         elif isinstance(queue_item, PartitionGenerationCompletedSentinel):
             yield from concurrent_stream_processor.on_partition_generation_completed(queue_item)
@@ -154,19 +163,3 @@ class ConcurrentSource:
             yield from concurrent_stream_processor.on_record(queue_item)
         else:
             raise ValueError(f"Unknown queue item type: {type(queue_item)}")
-
-    def _get_streams_to_read_from(self, streams: List[AbstractStream]) -> List[AbstractStream]:
-        """
-        Iterate over the configured streams and return a list of streams to read from.
-        If a stream is not configured, it will be skipped.
-        If a stream is configured but does not exist in the source and self.raise_exception_on_missing_stream is True, an exception will be raised
-        If a stream is not available, it will be skipped
-        """
-        stream_instances_to_read_from = []
-        for stream in streams:
-            stream_availability = stream.check_availability()
-            if not stream_availability.is_available():
-                self._logger.warning(f"Skipped syncing stream '{stream.name}' because it was unavailable. {stream_availability.message()}")
-                continue
-            stream_instances_to_read_from.append(stream)
-        return stream_instances_to_read_from

airbyte_cdk/sources/concurrent_source/concurrent_source_adapter.py

@@ -1,16 +1,33 @@
 #
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
+
 import logging
 from abc import ABC
-from typing import Any, Iterator, List, Mapping, MutableMapping, Optional, Union
+from datetime import timedelta
+from typing import Any, Callable, Iterator, List, Mapping, MutableMapping, Optional, Tuple
 
 from airbyte_cdk.models import AirbyteMessage, AirbyteStateMessage, ConfiguredAirbyteCatalog
 from airbyte_cdk.sources import AbstractSource
 from airbyte_cdk.sources.concurrent_source.concurrent_source import ConcurrentSource
+from airbyte_cdk.sources.connector_state_manager import ConnectorStateManager
 from airbyte_cdk.sources.streams import Stream
 from airbyte_cdk.sources.streams.concurrent.abstract_stream import AbstractStream
 from airbyte_cdk.sources.streams.concurrent.abstract_stream_facade import AbstractStreamFacade
+from airbyte_cdk.sources.streams.concurrent.adapters import StreamFacade
+from airbyte_cdk.sources.streams.concurrent.cursor import (
+    ConcurrentCursor,
+    Cursor,
+    CursorField,
+    CursorValueType,
+    FinalStateCursor,
+    GapType,
+)
+from airbyte_cdk.sources.streams.concurrent.state_converters.abstract_stream_state_converter import (
+    AbstractStreamStateConverter,
+)
+
+DEFAULT_LOOKBACK_SECONDS = 0
 
 
 class ConcurrentSourceAdapter(AbstractSource, ABC):
@@ -30,19 +47,25 @@ class ConcurrentSourceAdapter(AbstractSource, ABC):
         logger: logging.Logger,
         config: Mapping[str, Any],
         catalog: ConfiguredAirbyteCatalog,
-        state: Optional[Union[List[AirbyteStateMessage], MutableMapping[str, Any]]] = None,
+        state: Optional[List[AirbyteStateMessage]] = None,
     ) -> Iterator[AirbyteMessage]:
         abstract_streams = self._select_abstract_streams(config, catalog)
         concurrent_stream_names = {stream.name for stream in abstract_streams}
         configured_catalog_for_regular_streams = ConfiguredAirbyteCatalog(
-            streams=[stream for stream in catalog.streams if stream.stream.name not in concurrent_stream_names]
+            streams=[
+                stream
+                for stream in catalog.streams
+                if stream.stream.name not in concurrent_stream_names
+            ]
         )
         if abstract_streams:
             yield from self._concurrent_source.read(abstract_streams)
         if configured_catalog_for_regular_streams.streams:
             yield from super().read(logger, config, configured_catalog_for_regular_streams, state)
 
-    def _select_abstract_streams(self, config: Mapping[str, Any], configured_catalog: ConfiguredAirbyteCatalog) -> List[AbstractStream]:
+    def _select_abstract_streams(
+        self, config: Mapping[str, Any], configured_catalog: ConfiguredAirbyteCatalog
+    ) -> List[AbstractStream]:
         """
         Selects streams that can be processed concurrently and returns their abstract representations.
         """
@@ -52,12 +75,73 @@ class ConcurrentSourceAdapter(AbstractSource, ABC):
         for configured_stream in configured_catalog.streams:
             stream_instance = stream_name_to_instance.get(configured_stream.stream.name)
             if not stream_instance:
-                if not self.raise_exception_on_missing_stream:
-                    continue
-                raise KeyError(
-                    f"The stream {configured_stream.stream.name} no longer exists in the configuration. "
-                    f"Refresh the schema in replication settings and remove this stream from future sync attempts."
-                )
+                continue
+
             if isinstance(stream_instance, AbstractStreamFacade):
                 abstract_streams.append(stream_instance.get_underlying_stream())
         return abstract_streams
+
+    def convert_to_concurrent_stream(
+        self,
+        logger: logging.Logger,
+        stream: Stream,
+        state_manager: ConnectorStateManager,
+        cursor: Optional[Cursor] = None,
+    ) -> Stream:
+        """
+        Prepares a stream for concurrent processing by initializing or assigning a cursor,
+        managing the stream's state, and returning an updated Stream instance.
+        """
+        state: MutableMapping[str, Any] = {}
+
+        if cursor:
+            state = state_manager.get_stream_state(stream.name, stream.namespace)
+
+            stream.cursor = cursor  # type: ignore[assignment]  # cursor is of type ConcurrentCursor, which inherits from Cursor
+            if hasattr(stream, "parent"):
+                stream.parent.cursor = cursor
+        else:
+            cursor = FinalStateCursor(
+                stream_name=stream.name,
+                stream_namespace=stream.namespace,
+                message_repository=self.message_repository,  # type: ignore[arg-type]  # _default_message_repository will be returned in the worst case
+            )
+        return StreamFacade.create_from_stream(stream, self, logger, state, cursor)
+
+    def initialize_cursor(
+        self,
+        stream: Stream,
+        state_manager: ConnectorStateManager,
+        converter: AbstractStreamStateConverter,
+        slice_boundary_fields: Optional[Tuple[str, str]],
+        start: Optional[CursorValueType],
+        end_provider: Callable[[], CursorValueType],
+        lookback_window: Optional[GapType] = None,
+        slice_range: Optional[GapType] = None,
+    ) -> Optional[ConcurrentCursor]:
+        lookback_window = lookback_window or timedelta(seconds=DEFAULT_LOOKBACK_SECONDS)
+
+        cursor_field_name = stream.cursor_field
+
+        if cursor_field_name:
+            if not isinstance(cursor_field_name, str):
+                raise ValueError(
+                    f"Cursor field type must be a string, but received {type(cursor_field_name).__name__}."
+                )
+
+            return ConcurrentCursor(
+                stream.name,
+                stream.namespace,
+                state_manager.get_stream_state(stream.name, stream.namespace),
+                self.message_repository,  # type: ignore[arg-type]  # _default_message_repository will be returned in the worst case
+                state_manager,
+                converter,
+                CursorField(cursor_field_name),
+                slice_boundary_fields,
+                start,
+                end_provider,
+                lookback_window,
+                slice_range,
+            )
+
+        return None

airbyte_cdk/sources/concurrent_source/stream_thread_exception.py

@@ -0,0 +1,25 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+
+from typing import Any
+
+
+class StreamThreadException(Exception):
+    def __init__(self, exception: Exception, stream_name: str):
+        self._exception = exception
+        self._stream_name = stream_name
+
+    @property
+    def stream_name(self) -> str:
+        return self._stream_name
+
+    @property
+    def exception(self) -> Exception:
+        return self._exception
+
+    def __str__(self) -> str:
+        return f"Exception while syncing stream {self._stream_name}: {self._exception}"
+
+    def __eq__(self, other: Any) -> bool:
+        if isinstance(other, StreamThreadException):
+            return self._exception == other._exception and self._stream_name == other._stream_name
+        return False

airbyte_cdk/sources/concurrent_source/thread_pool_manager.py

@@ -37,7 +37,9 @@ class ThreadPoolManager:
     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)})")
+            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:
@@ -71,39 +73,43 @@ class ThreadPoolManager:
                         )
                     futures.pop(index)
 
-    def shutdown(self) -> None:
+    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 shutdown_if_exception(self) -> None:
-        """
-        This method will raise if there is an exception so that the caller can use it.
-        """
-        if self._most_recently_seen_exception:
-            self._stop_and_raise_exception(self._most_recently_seen_exception)
-
     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:
         """
-        self.shutdown_if_exception()
+        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]
+        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}")
+                exception = RuntimeError(
+                    f"Failed reading with futures not done: {futures_not_done}"
+                )
                 self._stop_and_raise_exception(exception)
             else:
-                self.shutdown()
+                self._shutdown()
 
     def _stop_and_raise_exception(self, exception: BaseException) -> None:
-        self.shutdown()
+        self._shutdown()
         raise exception

airbyte_cdk/sources/config.py

@@ -4,8 +4,9 @@
 
 from typing import Any, Dict
 
+from pydantic.v1 import BaseModel
+
 from airbyte_cdk.sources.utils.schema_helpers import expand_refs, rename_key
-from pydantic import BaseModel
 
 
 class BaseConfig(BaseModel):
@@ -17,7 +18,7 @@ class BaseConfig(BaseModel):
     """
 
     @classmethod
-    def schema(cls, *args, **kwargs) -> Dict[str, Any]:
+    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

airbyte_cdk/sources/connector_state_manager.py

@@ -3,45 +3,40 @@
 #
 
 import copy
-from typing import Any, List, Mapping, MutableMapping, Optional, Tuple, Union
+from dataclasses import dataclass
+from typing import Any, List, Mapping, MutableMapping, Optional, Tuple, Union, cast
 
 from airbyte_cdk.models import (
     AirbyteMessage,
     AirbyteStateBlob,
     AirbyteStateMessage,
     AirbyteStateType,
-    AirbyteStream,
     AirbyteStreamState,
     StreamDescriptor,
 )
 from airbyte_cdk.models import Type as MessageType
-from airbyte_cdk.sources.streams import Stream
-from pydantic import Extra
+from airbyte_cdk.models.airbyte_protocol import AirbyteGlobalState, AirbyteStateBlob
 
 
-class HashableStreamDescriptor(StreamDescriptor):
+@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.
     """
 
-    class Config:
-        extra = Extra.allow
-        frozen = True
+    name: str
+    namespace: Optional[str] = None
 
 
 class ConnectorStateManager:
     """
-    ConnectorStateManager consolidates the various forms of a stream's incoming state message (STREAM / GLOBAL / LEGACY) under a common
+    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,
-        stream_instance_map: Mapping[str, Union[Stream, AirbyteStream]],
-        state: Optional[Union[List[AirbyteStateMessage], MutableMapping[str, Any]]] = None,
-    ):
-        shared_state, per_stream_states = self._extract_from_state_message(state, stream_instance_map)
+    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
@@ -55,19 +50,25 @@ class ConnectorStateManager:
             )
         self.per_stream_states = per_stream_states
 
-    def get_stream_state(self, stream_name: str, namespace: Optional[str]) -> MutableMapping[str, Any]:
+    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 = self.per_stream_states.get(HashableStreamDescriptor(name=stream_name, namespace=namespace))
+        stream_state: AirbyteStateBlob | None = self.per_stream_states.get(
+            HashableStreamDescriptor(name=stream_name, namespace=namespace)
+        )
         if stream_state:
-            return stream_state.dict()  # type: ignore # mypy thinks dict() returns any, but it returns a dict
+            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:
+    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
@@ -75,11 +76,11 @@ class ConnectorStateManager:
         :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.parse_obj(value)
+        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 in either the per-stream or legacy format
+        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
@@ -92,7 +93,8 @@ class ConnectorStateManager:
             state=AirbyteStateMessage(
                 type=AirbyteStateType.STREAM,
                 stream=AirbyteStreamState(
-                    stream_descriptor=StreamDescriptor(name=stream_name, namespace=namespace), stream_state=stream_state
+                    stream_descriptor=StreamDescriptor(name=stream_name, namespace=namespace),
+                    stream_state=stream_state,
                 ),
             ),
         )
@@ -100,86 +102,48 @@ class ConnectorStateManager:
     @classmethod
     def _extract_from_state_message(
         cls,
-        state: Optional[Union[List[AirbyteStateMessage], MutableMapping[str, Any]]],
-        stream_instance_map: Mapping[str, Union[Stream, AirbyteStream]],
-    ) -> Tuple[Optional[AirbyteStateBlob], MutableMapping[HashableStreamDescriptor, Optional[AirbyteStateBlob]]]:
+        state: Optional[List[AirbyteStateMessage]],
+    ) -> Tuple[
+        Optional[AirbyteStateBlob],
+        MutableMapping[HashableStreamDescriptor, Optional[AirbyteStateBlob]],
+    ]:
         """
-        Takes an incoming list of state messages or the legacy state format and extracts state attributes according to type
-        which can then be assigned to the new state manager being instantiated
+        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_legacy = cls._is_legacy_dict_state(state)
-        is_migrated_legacy = cls._is_migrated_legacy_state(state)
         is_global = cls._is_global_state(state)
-        is_per_stream = cls._is_per_stream_state(state)
-
-        # Incoming pure legacy object format
-        if is_legacy:
-            streams = cls._create_descriptor_to_stream_state_mapping(state, stream_instance_map)  # type: ignore # We verified state is a dict in _is_legacy_dict_state
-            return None, streams
-
-        # When processing incoming state in source.read_state(), legacy state gets deserialized into List[AirbyteStateMessage]
-        # which can be translated into independent per-stream state values
-        if is_migrated_legacy:
-            streams = cls._create_descriptor_to_stream_state_mapping(state[0].data, stream_instance_map)  # type: ignore # We verified that state is a list in _is_migrated_legacy_state
-            return None, streams
 
         if is_global:
-            global_state = state[0].global_  # type: ignore # We verified state is a list in _is_global_state
-            shared_state = copy.deepcopy(global_state.shared_state, {})
+            # 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
+                    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
+                for per_stream_state in global_state.stream_states  # type: ignore[union-attr] # global_state has shared_state
             }
             return shared_state, streams
-
-        if is_per_stream:
+        else:
             streams = {
                 HashableStreamDescriptor(
-                    name=per_stream_state.stream.stream_descriptor.name, namespace=per_stream_state.stream.stream_descriptor.namespace
-                ): per_stream_state.stream.stream_state
+                    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
+                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
-        else:
-            raise ValueError("Input state should come in the form of list of Airbyte state messages or a mapping of states")
-
-    @staticmethod
-    def _create_descriptor_to_stream_state_mapping(
-        state: MutableMapping[str, Any], stream_to_instance_map: Mapping[str, Union[Stream, AirbyteStream]]
-    ) -> MutableMapping[HashableStreamDescriptor, Optional[AirbyteStateBlob]]:
-        """
-        Takes incoming state received in the legacy format and transforms it into a mapping of StreamDescriptor to AirbyteStreamState
-        :param state: A mapping object representing the complete state of all streams in the legacy format
-        :param stream_to_instance_map: A mapping of stream name to stream instance used to retrieve a stream's namespace
-        :return: The mapping of all of a sync's streams to the corresponding stream state
-        """
-        streams = {}
-        for stream_name, state_value in state.items():
-            namespace = stream_to_instance_map[stream_name].namespace if stream_name in stream_to_instance_map else None
-            stream_descriptor = HashableStreamDescriptor(name=stream_name, namespace=namespace)
-            streams[stream_descriptor] = AirbyteStateBlob.parse_obj(state_value or {})
-        return streams
-
-    @staticmethod
-    def _is_legacy_dict_state(state: Union[List[AirbyteStateMessage], MutableMapping[str, Any]]) -> bool:
-        return isinstance(state, dict)
-
-    @staticmethod
-    def _is_migrated_legacy_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.LEGACY
-        )
 
     @staticmethod
     def _is_global_state(state: Union[List[AirbyteStateMessage], MutableMapping[str, Any]]) -> bool:
@@ -191,5 +155,7 @@ class ConnectorStateManager:
         )
 
     @staticmethod
-    def _is_per_stream_state(state: Union[List[AirbyteStateMessage], MutableMapping[str, Any]]) -> bool:
+    def _is_per_stream_state(
+        state: Union[List[AirbyteStateMessage], MutableMapping[str, Any]],
+    ) -> bool:
         return isinstance(state, List)

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()})"