airbyte-cdk 6.31.2.dev0__py3-none-any.whl → 6.33.0__py3-none-any.whl

airbyte_cdk/sources/declarative/concurrent_declarative_source.py

@@ -19,7 +19,11 @@ from airbyte_cdk.sources.declarative.extractors import RecordSelector
 from airbyte_cdk.sources.declarative.extractors.record_filter import (
     ClientSideIncrementalRecordFilterDecorator,
 )
+from airbyte_cdk.sources.declarative.incremental import ConcurrentPerPartitionCursor
 from airbyte_cdk.sources.declarative.incremental.datetime_based_cursor import DatetimeBasedCursor
+from airbyte_cdk.sources.declarative.incremental.per_partition_with_global import (
+    PerPartitionWithGlobalCursor,
+)
 from airbyte_cdk.sources.declarative.interpolation import InterpolatedString
 from airbyte_cdk.sources.declarative.manifest_declarative_source import ManifestDeclarativeSource
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
@@ -31,8 +35,9 @@ from airbyte_cdk.sources.declarative.models.declarative_component_schema import
 from airbyte_cdk.sources.declarative.parsers.model_to_component_factory import (
     ModelToComponentFactory,
 )
+from airbyte_cdk.sources.declarative.partition_routers import AsyncJobPartitionRouter
 from airbyte_cdk.sources.declarative.requesters import HttpRequester
-from airbyte_cdk.sources.declarative.retrievers import SimpleRetriever
+from airbyte_cdk.sources.declarative.retrievers import AsyncRetriever, Retriever, SimpleRetriever
 from airbyte_cdk.sources.declarative.stream_slicers.declarative_partition_generator import (
     DeclarativePartitionFactory,
     StreamSlicerPartitionGenerator,
@@ -45,7 +50,7 @@ from airbyte_cdk.sources.streams.concurrent.abstract_stream import AbstractStrea
 from airbyte_cdk.sources.streams.concurrent.availability_strategy import (
     AlwaysAvailableAvailabilityStrategy,
 )
-from airbyte_cdk.sources.streams.concurrent.cursor import FinalStateCursor
+from airbyte_cdk.sources.streams.concurrent.cursor import ConcurrentCursor, FinalStateCursor
 from airbyte_cdk.sources.streams.concurrent.default_stream import DefaultStream
 from airbyte_cdk.sources.streams.concurrent.helpers import get_primary_key_from_stream
 
@@ -66,6 +71,10 @@ class ConcurrentDeclarativeSource(ManifestDeclarativeSource, Generic[TState]):
         component_factory: Optional[ModelToComponentFactory] = None,
         **kwargs: Any,
     ) -> None:
+        # todo: We could remove state from initialization. Now that streams are grouped during the read(), a source
+        #  no longer needs to store the original incoming state. But maybe there's an edge case?
+        self._connector_state_manager = ConnectorStateManager(state=state)  # type: ignore  # state is always in the form of List[AirbyteStateMessage]. The ConnectorStateManager should use generics, but this can be done later
+
         # To reduce the complexity of the concurrent framework, we are not enabling RFR with synthetic
         # cursors. We do this by no longer automatically instantiating RFR cursors when converting
         # the declarative models into runtime components. Concurrent sources will continue to checkpoint
@@ -73,19 +82,17 @@ class ConcurrentDeclarativeSource(ManifestDeclarativeSource, Generic[TState]):
         component_factory = component_factory or ModelToComponentFactory(
             emit_connector_builder_messages=emit_connector_builder_messages,
             disable_resumable_full_refresh=True,
+            connector_state_manager=self._connector_state_manager,
         )
 
         super().__init__(
             source_config=source_config,
+            config=config,
             debug=debug,
             emit_connector_builder_messages=emit_connector_builder_messages,
             component_factory=component_factory,
         )
 
-        # todo: We could remove state from initialization. Now that streams are grouped during the read(), a source
-        #  no longer needs to store the original incoming state. But maybe there's an edge case?
-        self._state = state
-
         concurrency_level_from_manifest = self._source_config.get("concurrency_level")
         if concurrency_level_from_manifest:
             concurrency_level_component = self._constructor.create_component(
@@ -175,8 +182,6 @@ class ConcurrentDeclarativeSource(ManifestDeclarativeSource, Generic[TState]):
         concurrent_streams: List[AbstractStream] = []
         synchronous_streams: List[Stream] = []
 
-        state_manager = ConnectorStateManager(state=self._state)  # type: ignore  # state is always in the form of List[AirbyteStateMessage]. The ConnectorStateManager should use generics, but this can be done later
-
         # Combine streams and dynamic_streams. Note: both cannot be empty at the same time,
         # and this is validated during the initialization of the source.
         streams = self._stream_configs(self._source_config) + self._dynamic_stream_configs(
@@ -216,45 +221,52 @@ class ConcurrentDeclarativeSource(ManifestDeclarativeSource, Generic[TState]):
                 if self._is_datetime_incremental_without_partition_routing(
                     declarative_stream, incremental_sync_component_definition
                 ):
-                    stream_state = state_manager.get_stream_state(
+                    stream_state = self._connector_state_manager.get_stream_state(
                         stream_name=declarative_stream.name, namespace=declarative_stream.namespace
                     )
 
-                    cursor = self._constructor.create_concurrent_cursor_from_datetime_based_cursor(
-                        state_manager=state_manager,
-                        model_type=DatetimeBasedCursorModel,
-                        component_definition=incremental_sync_component_definition,  # type: ignore  # Not None because of the if condition above
-                        stream_name=declarative_stream.name,
-                        stream_namespace=declarative_stream.namespace,
-                        config=config or {},
-                        stream_state=stream_state,
-                    )
+                    retriever = self._get_retriever(declarative_stream, stream_state)
 
-                    retriever = declarative_stream.retriever
-
-                    # This is an optimization so that we don't invoke any cursor or state management flows within the
-                    # low-code framework because state management is handled through the ConcurrentCursor.
-                    if declarative_stream and isinstance(retriever, SimpleRetriever):
-                        # Also a temporary hack. In the legacy Stream implementation, as part of the read,
-                        # set_initial_state() is called to instantiate incoming state on the cursor. Although we no
-                        # longer rely on the legacy low-code cursor for concurrent checkpointing, low-code components
-                        # like StopConditionPaginationStrategyDecorator and ClientSideIncrementalRecordFilterDecorator
-                        # still rely on a DatetimeBasedCursor that is properly initialized with state.
-                        if retriever.cursor:
-                            retriever.cursor.set_initial_state(stream_state=stream_state)
-                        # We zero it out here, but since this is a cursor reference, the state is still properly
-                        # instantiated for the other components that reference it
-                        retriever.cursor = None
+                    if isinstance(declarative_stream.retriever, AsyncRetriever) and isinstance(
+                        declarative_stream.retriever.stream_slicer, AsyncJobPartitionRouter
+                    ):
+                        cursor = declarative_stream.retriever.stream_slicer.stream_slicer
 
-                    partition_generator = StreamSlicerPartitionGenerator(
-                        DeclarativePartitionFactory(
-                            declarative_stream.name,
-                            declarative_stream.get_json_schema(),
-                            retriever,
-                            self.message_repository,
-                        ),
-                        cursor,
-                    )
+                        if not isinstance(cursor, ConcurrentCursor | ConcurrentPerPartitionCursor):
+                            # This should never happen since we instantiate ConcurrentCursor in
+                            # model_to_component_factory.py
+                            raise ValueError(
+                                f"Expected AsyncJobPartitionRouter stream_slicer to be of type ConcurrentCursor, but received{cursor.__class__}"
+                            )
+
+                        partition_generator = StreamSlicerPartitionGenerator(
+                            partition_factory=DeclarativePartitionFactory(
+                                declarative_stream.name,
+                                declarative_stream.get_json_schema(),
+                                retriever,
+                                self.message_repository,
+                            ),
+                            stream_slicer=declarative_stream.retriever.stream_slicer,
+                        )
+                    else:
+                        cursor = (
+                            self._constructor.create_concurrent_cursor_from_datetime_based_cursor(
+                                model_type=DatetimeBasedCursorModel,
+                                component_definition=incremental_sync_component_definition,  # type: ignore  # Not None because of the if condition above
+                                stream_name=declarative_stream.name,
+                                stream_namespace=declarative_stream.namespace,
+                                config=config or {},
+                            )
+                        )
+                        partition_generator = StreamSlicerPartitionGenerator(
+                            partition_factory=DeclarativePartitionFactory(
+                                declarative_stream.name,
+                                declarative_stream.get_json_schema(),
+                                retriever,
+                                self.message_repository,
+                            ),
+                            stream_slicer=cursor,
+                        )
 
                     concurrent_streams.append(
                         DefaultStream(
@@ -304,6 +316,60 @@ class ConcurrentDeclarativeSource(ManifestDeclarativeSource, Generic[TState]):
                             cursor=final_state_cursor,
                         )
                     )
+                elif (
+                    incremental_sync_component_definition
+                    and incremental_sync_component_definition.get("type", "")
+                    == DatetimeBasedCursorModel.__name__
+                    and self._stream_supports_concurrent_partition_processing(
+                        declarative_stream=declarative_stream
+                    )
+                    and hasattr(declarative_stream.retriever, "stream_slicer")
+                    and isinstance(
+                        declarative_stream.retriever.stream_slicer, PerPartitionWithGlobalCursor
+                    )
+                ):
+                    stream_state = self._connector_state_manager.get_stream_state(
+                        stream_name=declarative_stream.name, namespace=declarative_stream.namespace
+                    )
+                    partition_router = declarative_stream.retriever.stream_slicer._partition_router
+
+                    perpartition_cursor = (
+                        self._constructor.create_concurrent_cursor_from_perpartition_cursor(
+                            state_manager=self._connector_state_manager,
+                            model_type=DatetimeBasedCursorModel,
+                            component_definition=incremental_sync_component_definition,
+                            stream_name=declarative_stream.name,
+                            stream_namespace=declarative_stream.namespace,
+                            config=config or {},
+                            stream_state=stream_state,
+                            partition_router=partition_router,
+                        )
+                    )
+
+                    retriever = self._get_retriever(declarative_stream, stream_state)
+
+                    partition_generator = StreamSlicerPartitionGenerator(
+                        DeclarativePartitionFactory(
+                            declarative_stream.name,
+                            declarative_stream.get_json_schema(),
+                            retriever,
+                            self.message_repository,
+                        ),
+                        perpartition_cursor,
+                    )
+
+                    concurrent_streams.append(
+                        DefaultStream(
+                            partition_generator=partition_generator,
+                            name=declarative_stream.name,
+                            json_schema=declarative_stream.get_json_schema(),
+                            availability_strategy=AlwaysAvailableAvailabilityStrategy(),
+                            primary_key=get_primary_key_from_stream(declarative_stream.primary_key),
+                            cursor_field=perpartition_cursor.cursor_field.cursor_field_key,
+                            logger=self.logger,
+                            cursor=perpartition_cursor,
+                        )
+                    )
                 else:
                     synchronous_streams.append(declarative_stream)
             else:
@@ -325,7 +391,10 @@ class ConcurrentDeclarativeSource(ManifestDeclarativeSource, Generic[TState]):
                 declarative_stream=declarative_stream
             )
             and hasattr(declarative_stream.retriever, "stream_slicer")
-            and isinstance(declarative_stream.retriever.stream_slicer, DatetimeBasedCursor)
+            and (
+                isinstance(declarative_stream.retriever.stream_slicer, DatetimeBasedCursor)
+                or isinstance(declarative_stream.retriever.stream_slicer, AsyncJobPartitionRouter)
+            )
         )
 
     def _stream_supports_concurrent_partition_processing(
@@ -394,6 +463,39 @@ class ConcurrentDeclarativeSource(ManifestDeclarativeSource, Generic[TState]):
                             return False
         return True
 
+    @staticmethod
+    def _get_retriever(
+        declarative_stream: DeclarativeStream, stream_state: Mapping[str, Any]
+    ) -> Retriever:
+        retriever = declarative_stream.retriever
+
+        # This is an optimization so that we don't invoke any cursor or state management flows within the
+        # low-code framework because state management is handled through the ConcurrentCursor.
+        if declarative_stream and isinstance(retriever, SimpleRetriever):
+            # Also a temporary hack. In the legacy Stream implementation, as part of the read,
+            # set_initial_state() is called to instantiate incoming state on the cursor. Although we no
+            # longer rely on the legacy low-code cursor for concurrent checkpointing, low-code components
+            # like StopConditionPaginationStrategyDecorator still rely on a DatetimeBasedCursor that is
+            # properly initialized with state.
+            if retriever.cursor:
+                retriever.cursor.set_initial_state(stream_state=stream_state)
+
+            # Similar to above, the ClientSideIncrementalRecordFilterDecorator cursor is a separate instance
+            # from the one initialized on the SimpleRetriever, so it also must also have state initialized
+            # for semi-incremental streams using is_client_side_incremental to filter properly
+            if isinstance(retriever.record_selector, RecordSelector) and isinstance(
+                retriever.record_selector.record_filter, ClientSideIncrementalRecordFilterDecorator
+            ):
+                retriever.record_selector.record_filter._cursor.set_initial_state(
+                    stream_state=stream_state
+                )  # type: ignore  # After non-concurrent cursors are deprecated we can remove these cursor workarounds
+
+            # We zero it out here, but since this is a cursor reference, the state is still properly
+            # instantiated for the other components that reference it
+            retriever.cursor = None
+
+        return retriever
+
     @staticmethod
     def _select_streams(
         streams: List[AbstractStream], configured_catalog: ConfiguredAirbyteCatalog

airbyte_cdk/sources/declarative/declarative_component_schema.yaml

@@ -320,6 +320,11 @@ definitions:
         title: Stream Count
         description: Numbers of the streams to try reading from when running a check operation.
         type: integer
+      use_check_availability:
+        title: Use Check Availability
+        description: Enables stream check availability. This field is automatically set by the CDK.
+        type: boolean
+        default: true
   CompositeErrorHandler:
     title: Composite Error Handler
     description: Error handler that sequentially iterates over a list of error handlers.
@@ -784,6 +789,29 @@ definitions:
       type:
         type: string
         enum: [DatetimeBasedCursor]
+      clamping:
+        title: Date Range Clamping
+        description: This option is used to adjust the upper and lower boundaries of each datetime window to beginning and end of the provided target period (day, week, month)
+        type: object
+        required:
+        - target
+        properties:
+          target:
+            title: Target
+            description: The period of time that datetime windows will be clamped by
+            # This should ideally be an enum. However, we don't use an enum because we want to allow for connectors
+            # to support interpolation on the connector config to get the target which is an arbitrary string
+            type: string
+            interpolation_context:
+              - config
+            examples:
+              - "DAY"
+              - "WEEK"
+              - "MONTH"
+              - "{{ config['target'] }}"
+          target_details:
+            type: object
+            additionalProperties: true
       cursor_field:
         title: Cursor Field
         description: The location of the value on a record that will be used as a bookmark during sync. To ensure no data loss, the API must return records in ascending order based on the cursor field. Nested fields are not supported, so the field must be at the top level of the record. You can use a combination of Add Field and Remove Field transformations to move the nested field to the top.
@@ -1058,8 +1086,6 @@ definitions:
     type: object
     required:
       - type
-      - client_id
-      - client_secret
     properties:
       type:
         type: string
@@ -1254,6 +1280,15 @@ definitions:
             default: []
             examples:
               - ["invalid_grant", "invalid_permissions"]
+      profile_assertion:
+        title: Profile Assertion
+        description: The authenticator being used to authenticate the client authenticator.
+        "$ref": "#/definitions/JwtAuthenticator"
+      use_profile_assertion:
+        title: Use Profile Assertion
+        description: Enable using profile assertion as a flow for OAuth authorization.
+        type: boolean
+        default: false
       $parameters:
         type: object
         additionalProperties: true
@@ -1770,6 +1805,19 @@ definitions:
       $parameters:
         type: object
         additionalProperties: true
+  ComplexFieldType:
+    title: Schema Field Type
+    description: (This component is experimental. Use at your own risk.) Represents a complex field type.
+    type: object
+    required:
+      - field_type
+    properties:
+      field_type:
+        type: string
+      items:
+        anyOf:
+          - type: string
+          - "$ref": "#/definitions/ComplexFieldType"
   TypesMap:
     title: Types Map
     description: (This component is experimental. Use at your own risk.) Represents a mapping between a current type and its corresponding target type.
@@ -1784,6 +1832,7 @@ definitions:
           - type: array
             items:
               type: string
+          - "$ref": "#/definitions/ComplexFieldType"
       current_type:
         anyOf:
           - type: string

airbyte_cdk/sources/declarative/declarative_stream.py

@@ -138,7 +138,9 @@ class DeclarativeStream(Stream):
         """
         :param: stream_state We knowingly avoid using stream_state as we want cursors to manage their own state.
         """
-        if stream_slice is None or stream_slice == {}:
+        if stream_slice is None or (
+            not isinstance(stream_slice, StreamSlice) and stream_slice == {}
+        ):
             # As the parameter is Optional, many would just call `read_records(sync_mode)` during testing without specifying the field
             # As part of the declarative model without custom components, this should never happen as the CDK would wire up a
             # SinglePartitionRouter that would create this StreamSlice properly

airbyte_cdk/sources/declarative/extractors/record_filter.py

@@ -59,13 +59,11 @@ class ClientSideIncrementalRecordFilterDecorator(RecordFilter):
 
     def __init__(
         self,
-        date_time_based_cursor: DatetimeBasedCursor,
-        substream_cursor: Optional[Union[PerPartitionWithGlobalCursor, GlobalSubstreamCursor]],
+        cursor: Union[DatetimeBasedCursor, PerPartitionWithGlobalCursor, GlobalSubstreamCursor],
         **kwargs: Any,
     ):
         super().__init__(**kwargs)
-        self._date_time_based_cursor = date_time_based_cursor
-        self._substream_cursor = substream_cursor
+        self._cursor = cursor
 
     def filter_records(
         self,
@@ -77,7 +75,7 @@ class ClientSideIncrementalRecordFilterDecorator(RecordFilter):
         records = (
             record
             for record in records
-            if (self._substream_cursor or self._date_time_based_cursor).should_be_synced(
+            if self._cursor.should_be_synced(
                 # Record is created on the fly to align with cursors interface; stream name is ignored as we don't need it here
                 # Record stream name is empty cause it is not used durig the filtering
                 Record(data=record, associated_slice=stream_slice, stream_name="")

airbyte_cdk/sources/declarative/incremental/__init__.py

@@ -2,6 +2,10 @@
 # Copyright (c) 2022 Airbyte, Inc., all rights reserved.
 #
 
+from airbyte_cdk.sources.declarative.incremental.concurrent_partition_cursor import (
+    ConcurrentCursorFactory,
+    ConcurrentPerPartitionCursor,
+)
 from airbyte_cdk.sources.declarative.incremental.datetime_based_cursor import DatetimeBasedCursor
 from airbyte_cdk.sources.declarative.incremental.declarative_cursor import DeclarativeCursor
 from airbyte_cdk.sources.declarative.incremental.global_substream_cursor import (
@@ -21,6 +25,8 @@ from airbyte_cdk.sources.declarative.incremental.resumable_full_refresh_cursor i
 
 __all__ = [
     "CursorFactory",
+    "ConcurrentCursorFactory",
+    "ConcurrentPerPartitionCursor",
     "DatetimeBasedCursor",
     "DeclarativeCursor",
     "GlobalSubstreamCursor",