airbyte-cdk 0.53.9__py3-none-any.whl → 0.55.0__py3-none-any.whl

airbyte_cdk/sources/declarative/requesters/paginators/default_paginator.py

@@ -3,7 +3,7 @@
 #
 
 from dataclasses import InitVar, dataclass
-from typing import Any, List, Mapping, Optional, Union
+from typing import Any, List, Mapping, MutableMapping, Optional, Union
 
 import requests
 from airbyte_cdk.sources.declarative.decoders.decoder import Decoder
@@ -94,7 +94,7 @@ class DefaultPaginator(Paginator):
     page_size_option: Optional[RequestOption] = None
     page_token_option: Optional[Union[RequestPath, RequestOption]] = None
 
-    def __post_init__(self, parameters: Mapping[str, Any]):
+    def __post_init__(self, parameters: Mapping[str, Any]) -> None:
         if self.page_size_option and not self.pagination_strategy.get_page_size():
             raise ValueError("page_size_option cannot be set if the pagination strategy does not have a page_size")
         if isinstance(self.url_base, str):
@@ -108,10 +108,10 @@ class DefaultPaginator(Paginator):
         else:
             return None
 
-    def path(self):
+    def path(self) -> Optional[str]:
         if self._token and self.page_token_option and isinstance(self.page_token_option, RequestPath):
             # Replace url base to only return the path
-            return str(self._token).replace(self.url_base.eval(self.config), "")
+            return str(self._token).replace(self.url_base.eval(self.config), "")  # type: ignore # url_base is casted to a InterpolatedString in __post_init__
         else:
             return None
 
@@ -121,7 +121,7 @@ class DefaultPaginator(Paginator):
         stream_state: Optional[StreamState] = None,
         stream_slice: Optional[StreamSlice] = None,
         next_page_token: Optional[Mapping[str, Any]] = None,
-    ) -> Mapping[str, Any]:
+    ) -> MutableMapping[str, Any]:
         return self._get_request_options(RequestOptionType.request_parameter)
 
     def get_request_headers(
@@ -151,11 +151,11 @@ class DefaultPaginator(Paginator):
     ) -> Mapping[str, Any]:
         return self._get_request_options(RequestOptionType.body_json)
 
-    def reset(self):
+    def reset(self) -> None:
         self.pagination_strategy.reset()
-        self._token = None
+        self._token = self.pagination_strategy.initial_token
 
-    def _get_request_options(self, option_type: RequestOptionType) -> Mapping[str, Any]:
+    def _get_request_options(self, option_type: RequestOptionType) -> MutableMapping[str, Any]:
         options = {}
 
         if (
@@ -178,7 +178,7 @@ class PaginatorTestReadDecorator(Paginator):
 
     _PAGE_COUNT_BEFORE_FIRST_NEXT_CALL = 1
 
-    def __init__(self, decorated, maximum_number_of_pages: int = 5):
+    def __init__(self, decorated: Paginator, maximum_number_of_pages: int = 5) -> None:
         if maximum_number_of_pages and maximum_number_of_pages < 1:
             raise ValueError(f"The maximum number of pages on a test read needs to be strictly positive. Got {maximum_number_of_pages}")
         self._maximum_number_of_pages = maximum_number_of_pages
@@ -192,7 +192,7 @@ class PaginatorTestReadDecorator(Paginator):
         self._page_count += 1
         return self._decorated.next_page_token(response, last_records)
 
-    def path(self):
+    def path(self) -> Optional[str]:
         return self._decorated.path()
 
     def get_request_params(
@@ -201,7 +201,7 @@ class PaginatorTestReadDecorator(Paginator):
         stream_state: Optional[StreamState] = None,
         stream_slice: Optional[StreamSlice] = None,
         next_page_token: Optional[Mapping[str, Any]] = None,
-    ) -> Mapping[str, Any]:
+    ) -> MutableMapping[str, Any]:
         return self._decorated.get_request_params(stream_state=stream_state, stream_slice=stream_slice, next_page_token=next_page_token)
 
     def get_request_headers(
@@ -219,7 +219,7 @@ class PaginatorTestReadDecorator(Paginator):
         stream_state: Optional[StreamState] = None,
         stream_slice: Optional[StreamSlice] = None,
         next_page_token: Optional[Mapping[str, Any]] = None,
-    ) -> Mapping[str, Any]:
+    ) -> Optional[Union[Mapping[str, Any], str]]:
         return self._decorated.get_request_body_data(stream_state=stream_state, stream_slice=stream_slice, next_page_token=next_page_token)
 
     def get_request_body_json(
@@ -228,9 +228,9 @@ class PaginatorTestReadDecorator(Paginator):
         stream_state: Optional[StreamState] = None,
         stream_slice: Optional[StreamSlice] = None,
         next_page_token: Optional[Mapping[str, Any]] = None,
-    ) -> Mapping[str, Any]:
+    ) -> Optional[Mapping[str, Any]]:
         return self._decorated.get_request_body_json(stream_state=stream_state, stream_slice=stream_slice, next_page_token=next_page_token)
 
-    def reset(self):
+    def reset(self) -> None:
         self._decorated.reset()
         self._page_count = self._PAGE_COUNT_BEFORE_FIRST_NEXT_CALL

airbyte_cdk/sources/declarative/requesters/request_options/request_options_provider.py

@@ -53,7 +53,7 @@ class RequestOptionsProvider:
         stream_state: Optional[StreamState] = None,
         stream_slice: Optional[StreamSlice] = None,
         next_page_token: Optional[Mapping[str, Any]] = None,
-    ) -> Optional[Union[Mapping, str]]:
+    ) -> Optional[Union[Mapping[str, Any], str]]:
         """
         Specifies how to populate the body of the request with a non-JSON payload.
 
@@ -71,7 +71,7 @@ class RequestOptionsProvider:
         stream_state: Optional[StreamState] = None,
         stream_slice: Optional[StreamSlice] = None,
         next_page_token: Optional[Mapping[str, Any]] = None,
-    ) -> Optional[Mapping]:
+    ) -> Optional[Mapping[str, Any]]:
         """
         Specifies how to populate the body of the request with a JSON payload.
 

airbyte_cdk/sources/streams/concurrent/abstract_stream.py

@@ -7,7 +7,7 @@ from typing import Any, Iterable, Mapping, Optional
 
 from airbyte_cdk.models import AirbyteStream
 from airbyte_cdk.sources.streams.concurrent.availability_strategy import StreamAvailability
-from airbyte_cdk.sources.streams.concurrent.partitions.record import Record
+from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
 from deprecated.classic import deprecated
 
 
@@ -37,10 +37,10 @@ class AbstractStream(ABC):
     """
 
     @abstractmethod
-    def read(self) -> Iterable[Record]:
+    def generate_partitions(self) -> Iterable[Partition]:
         """
-        Read a stream in full refresh mode
-        :return: The stream's records
+        Generates the partitions that will be read by this stream.
+        :return: An iterable of partitions.
         """
 
     @property

airbyte_cdk/sources/streams/concurrent/adapters.py

@@ -22,11 +22,11 @@ from airbyte_cdk.sources.streams.concurrent.availability_strategy import (
     StreamUnavailable,
 )
 from airbyte_cdk.sources.streams.concurrent.cursor import Cursor, NoopCursor
+from airbyte_cdk.sources.streams.concurrent.default_stream import DefaultStream
 from airbyte_cdk.sources.streams.concurrent.exceptions import ExceptionWithDisplayMessage
 from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
 from airbyte_cdk.sources.streams.concurrent.partitions.partition_generator import PartitionGenerator
 from airbyte_cdk.sources.streams.concurrent.partitions.record import Record
-from airbyte_cdk.sources.streams.concurrent.thread_based_concurrent_stream import ThreadBasedConcurrentStream
 from airbyte_cdk.sources.streams.core import StreamData
 from airbyte_cdk.sources.utils.schema_helpers import InternalConfig
 from airbyte_cdk.sources.utils.slice_logger import SliceLogger
@@ -52,7 +52,6 @@ class StreamFacade(Stream):
         stream: Stream,
         source: AbstractSource,
         logger: logging.Logger,
-        max_workers: int,
         state: Optional[MutableMapping[str, Any]],
         cursor: Cursor,
     ) -> Stream:
@@ -73,28 +72,27 @@ class StreamFacade(Stream):
 
         message_repository = source.message_repository
         return StreamFacade(
-            ThreadBasedConcurrentStream(
+            DefaultStream(
                 partition_generator=StreamPartitionGenerator(
                     stream,
                     message_repository,
                     SyncMode.full_refresh if isinstance(cursor, NoopCursor) else SyncMode.incremental,
                     [cursor_field] if cursor_field is not None else None,
                     state,
+                    cursor,
                 ),
-                max_workers=max_workers,
                 name=stream.name,
                 namespace=stream.namespace,
                 json_schema=stream.get_json_schema(),
                 availability_strategy=StreamAvailabilityStrategy(stream, source),
                 primary_key=pk,
                 cursor_field=cursor_field,
-                slice_logger=source._slice_logger,
-                message_repository=message_repository,
                 logger=logger,
-                cursor=cursor,
             ),
             stream,
             cursor,
+            slice_logger=source._slice_logger,
+            logger=logger,
         )
 
     @property
@@ -132,13 +130,15 @@ class StreamFacade(Stream):
         else:
             return stream.cursor_field
 
-    def __init__(self, stream: AbstractStream, legacy_stream: Stream, cursor: Cursor):
+    def __init__(self, stream: AbstractStream, legacy_stream: Stream, cursor: Cursor, slice_logger: SliceLogger, logger: logging.Logger):
         """
         :param stream: The underlying AbstractStream
         """
         self._abstract_stream = stream
         self._legacy_stream = legacy_stream
         self._cursor = cursor
+        self._slice_logger = slice_logger
+        self._logger = logger
 
     def read_full_refresh(
         self,
@@ -177,8 +177,11 @@ class StreamFacade(Stream):
         yield from self._read_records()
 
     def _read_records(self) -> Iterable[StreamData]:
-        for record in self._abstract_stream.read():
-            yield record.data
+        for partition in self._abstract_stream.generate_partitions():
+            if self._slice_logger.should_log_slice_message(self._logger):
+                yield self._slice_logger.create_slice_log_message(partition.to_slice())
+            for record in partition.read():
+                yield record.data
 
     @property
     def name(self) -> str:
@@ -259,6 +262,7 @@ class StreamPartition(Partition):
         sync_mode: SyncMode,
         cursor_field: Optional[List[str]],
         state: Optional[MutableMapping[str, Any]],
+        cursor: Cursor,
     ):
         """
         :param stream: The stream to delegate to
@@ -271,6 +275,8 @@ class StreamPartition(Partition):
         self._sync_mode = sync_mode
         self._cursor_field = cursor_field
         self._state = state
+        self._cursor = cursor
+        self._is_closed = False
 
     def read(self) -> Iterable[Record]:
         """
@@ -294,7 +300,9 @@ class StreamPartition(Partition):
                 if isinstance(record_data, Mapping):
                     data_to_return = dict(record_data)
                     self._stream.transformer.transform(data_to_return, self._stream.get_json_schema())
-                    yield Record(data_to_return)
+                    record = Record(data_to_return, self._stream.name)
+                    self._cursor.observe(record)
+                    yield Record(data_to_return, self._stream.name)
                 else:
                     self._message_repository.emit_message(record_data)
         except Exception as e:
@@ -315,6 +323,16 @@ class StreamPartition(Partition):
         else:
             return hash(self._stream.name)
 
+    def stream_name(self) -> str:
+        return self._stream.name
+
+    def close(self) -> None:
+        self._cursor.close_partition(self)
+        self._is_closed = True
+
+    def is_closed(self) -> bool:
+        return self._is_closed
+
     def __repr__(self) -> str:
         return f"StreamPartition({self._stream.name}, {self._slice})"
 
@@ -334,6 +352,7 @@ class StreamPartitionGenerator(PartitionGenerator):
         sync_mode: SyncMode,
         cursor_field: Optional[List[str]],
         state: Optional[MutableMapping[str, Any]],
+        cursor: Cursor,
     ):
         """
         :param stream: The stream to delegate to
@@ -344,10 +363,13 @@ class StreamPartitionGenerator(PartitionGenerator):
         self._sync_mode = sync_mode
         self._cursor_field = cursor_field
         self._state = state
+        self._cursor = cursor
 
     def generate(self) -> Iterable[Partition]:
         for s in self._stream.stream_slices(sync_mode=self._sync_mode, cursor_field=self._cursor_field, stream_state=self._state):
-            yield StreamPartition(self._stream, copy.deepcopy(s), self.message_repository, self._sync_mode, self._cursor_field, self._state)
+            yield StreamPartition(
+                self._stream, copy.deepcopy(s), self.message_repository, self._sync_mode, self._cursor_field, self._state, self._cursor
+            )
 
 
 @deprecated("This class is experimental. Use at your own risk.")

airbyte_cdk/sources/streams/concurrent/default_stream.py

@@ -0,0 +1,79 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+from functools import lru_cache
+from logging import Logger
+from typing import Any, Iterable, List, Mapping, Optional
+
+from airbyte_cdk.models import AirbyteStream, SyncMode
+from airbyte_cdk.sources.streams.concurrent.abstract_stream import AbstractStream
+from airbyte_cdk.sources.streams.concurrent.availability_strategy import AbstractAvailabilityStrategy, StreamAvailability
+from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
+from airbyte_cdk.sources.streams.concurrent.partitions.partition_generator import PartitionGenerator
+
+
+class DefaultStream(AbstractStream):
+    def __init__(
+        self,
+        partition_generator: PartitionGenerator,
+        name: str,
+        json_schema: Mapping[str, Any],
+        availability_strategy: AbstractAvailabilityStrategy,
+        primary_key: List[str],
+        cursor_field: Optional[str],
+        logger: Logger,
+        namespace: Optional[str] = None,
+    ) -> None:
+        self._stream_partition_generator = partition_generator
+        self._name = name
+        self._json_schema = json_schema
+        self._availability_strategy = availability_strategy
+        self._primary_key = primary_key
+        self._cursor_field = cursor_field
+        self._logger = logger
+        self._namespace = namespace
+
+    def generate_partitions(self) -> Iterable[Partition]:
+        yield from self._stream_partition_generator.generate()
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    def check_availability(self) -> StreamAvailability:
+        return self._availability_strategy.check_availability(self._logger)
+
+    @property
+    def cursor_field(self) -> Optional[str]:
+        return self._cursor_field
+
+    @lru_cache(maxsize=None)
+    def get_json_schema(self) -> Mapping[str, Any]:
+        return self._json_schema
+
+    def as_airbyte_stream(self) -> AirbyteStream:
+        stream = AirbyteStream(name=self.name, json_schema=dict(self._json_schema), supported_sync_modes=[SyncMode.full_refresh])
+
+        if self._namespace:
+            stream.namespace = self._namespace
+
+        if self._cursor_field:
+            stream.source_defined_cursor = True
+            stream.supported_sync_modes.append(SyncMode.incremental)
+            stream.default_cursor_field = [self._cursor_field]
+
+        keys = self._primary_key
+        if keys and len(keys) > 0:
+            stream.source_defined_primary_key = [keys]
+
+        return stream
+
+    def log_stream_sync_configuration(self) -> None:
+        self._logger.debug(
+            f"Syncing stream instance: {self.name}",
+            extra={
+                "primary_key": self._primary_key,
+                "cursor_field": self.cursor_field,
+            },
+        )

airbyte_cdk/sources/streams/concurrent/partition_enqueuer.py

@@ -4,8 +4,9 @@
 
 from queue import Queue
 
-from airbyte_cdk.sources.streams.concurrent.partitions.partition_generator import PartitionGenerator
-from airbyte_cdk.sources.streams.concurrent.partitions.types import PARTITIONS_GENERATED_SENTINEL, QueueItem
+from airbyte_cdk.sources.concurrent_source.partition_generation_completed_sentinel import PartitionGenerationCompletedSentinel
+from airbyte_cdk.sources.streams.concurrent.abstract_stream import AbstractStream
+from airbyte_cdk.sources.streams.concurrent.partitions.types import QueueItem
 
 
 class PartitionEnqueuer:
@@ -13,15 +14,14 @@ class PartitionEnqueuer:
     Generates partitions from a partition generator and puts them in a queue.
     """
 
-    def __init__(self, queue: Queue[QueueItem], sentinel: PARTITIONS_GENERATED_SENTINEL) -> None:
+    def __init__(self, queue: Queue[QueueItem]) -> None:
         """
         :param queue:  The queue to put the partitions in.
         :param sentinel: The sentinel to put in the queue when all the partitions have been generated.
         """
         self._queue = queue
-        self._sentinel = sentinel
 
-    def generate_partitions(self, partition_generator: PartitionGenerator) -> None:
+    def generate_partitions(self, stream: AbstractStream) -> None:
         """
         Generate partitions from a partition generator and put them in a queue.
         When all the partitions are added to the queue, a sentinel is added to the queue to indicate that all the partitions have been generated.
@@ -33,8 +33,8 @@ class PartitionEnqueuer:
         :return:
         """
         try:
-            for partition in partition_generator.generate():
+            for partition in stream.generate_partitions():
                 self._queue.put(partition)
-            self._queue.put(self._sentinel)
+            self._queue.put(PartitionGenerationCompletedSentinel(stream))
         except Exception as e:
             self._queue.put(e)

airbyte_cdk/sources/streams/concurrent/partitions/partition.py

@@ -32,6 +32,29 @@ class Partition(ABC):
         """
         pass
 
+    @abstractmethod
+    def stream_name(self) -> str:
+        """
+        Returns the name of the stream that this partition is reading from.
+        :return: The name of the stream.
+        """
+        pass
+
+    @abstractmethod
+    def close(self) -> None:
+        """
+        Closes the partition.
+        """
+        pass
+
+    @abstractmethod
+    def is_closed(self) -> bool:
+        """
+        Returns whether the partition is closed.
+        :return:
+        """
+        pass
+
     @abstractmethod
     def __hash__(self) -> int:
         """

airbyte_cdk/sources/streams/concurrent/partitions/record.py

@@ -10,13 +10,14 @@ class Record:
     Represents a record read from a stream.
     """
 
-    def __init__(self, data: Mapping[str, Any]):
+    def __init__(self, data: Mapping[str, Any], stream_name: str):
         self.data = data
+        self.stream_name = stream_name
 
     def __eq__(self, other: Any) -> bool:
         if not isinstance(other, Record):
             return False
-        return self.data == other.data
+        return self.data == other.data and self.stream_name == other.stream_name
 
     def __repr__(self) -> str:
-        return f"Record(data={self.data})"
+        return f"Record(data={self.data}, stream_name={self.stream_name})"

airbyte_cdk/sources/streams/concurrent/partitions/types.py

@@ -4,11 +4,10 @@
 
 from typing import Union
 
+from airbyte_cdk.sources.concurrent_source.partition_generation_completed_sentinel import PartitionGenerationCompletedSentinel
 from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
 from airbyte_cdk.sources.streams.concurrent.partitions.record import Record
 
-PARTITIONS_GENERATED_SENTINEL = object
-
 
 class PartitionCompleteSentinel:
     """
@@ -26,4 +25,4 @@ class PartitionCompleteSentinel:
 """
 Typedef representing the items that can be added to the ThreadBasedConcurrentStream
 """
-QueueItem = Union[Record, Partition, PartitionCompleteSentinel, PARTITIONS_GENERATED_SENTINEL, Partition, Exception]
+QueueItem = Union[Record, Partition, PartitionCompleteSentinel, PartitionGenerationCompletedSentinel, Exception]

airbyte_cdk/sources/utils/slice_logger.py

@@ -12,6 +12,11 @@ from airbyte_cdk.models import Type as MessageType
 
 
 class SliceLogger(ABC):
+    """
+    SliceLogger is an interface that allows us to log slices of data in a uniform way.
+    It is responsible for determining whether or not a slice should be logged and for creating the log message.
+    """
+
     SLICE_LOG_PREFIX = "slice:"
 
     def create_slice_log_message(self, _slice: Optional[Mapping[str, Any]]) -> AirbyteMessage:

{airbyte_cdk-0.53.9.dist-info → airbyte_cdk-0.55.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: airbyte-cdk
-Version: 0.53.9
+Version: 0.55.0
 Summary: A framework for writing Airbyte Connectors.
 Home-page: https://github.com/airbytehq/airbyte
 Author: Airbyte