WIP push

Author: maxime.c

airbyte_cdk/sources/declarative/concurrent_declarative_source.py

@@ -313,7 +313,7 @@ def _group_streams(
                         partition_generator = StreamSlicerPartitionGenerator(
                             partition_factory=DeclarativePartitionFactory(
                                 declarative_stream.name,
-                                declarative_stream.get_json_schema(),
+                                declarative_stream.schema_loader,
                                 retriever,
                                 self.message_repository,
                             ),
@@ -344,7 +344,7 @@ def _group_streams(
                     partition_generator = StreamSlicerPartitionGenerator(
                         DeclarativePartitionFactory(
                             declarative_stream.name,
-                            declarative_stream.get_json_schema(),
+                            declarative_stream.schema_loader,
                             declarative_stream.retriever,
                             self.message_repository,
                         ),
@@ -361,7 +361,7 @@ def _group_streams(
                         DefaultStream(
                             partition_generator=partition_generator,
                             name=declarative_stream.name,
-                            json_schema=declarative_stream.get_json_schema(),
+                            json_schema=lambda: declarative_stream.schema_loader.get_json_schema(),
                             availability_strategy=AlwaysAvailableAvailabilityStrategy(),
                             primary_key=get_primary_key_from_stream(declarative_stream.primary_key),
                             cursor_field=None,
@@ -455,49 +455,14 @@ def _is_concurrent_cursor_incremental_without_partition_routing(
                 in (DatetimeBasedCursorModel.__name__, IncrementingCountCursorModel.__name__)
             )
             and hasattr(declarative_stream.retriever, "stream_slicer")
-            and (
-                isinstance(declarative_stream.retriever.stream_slicer, DatetimeBasedCursor)
-                # IncrementingCountCursorModel is hardcoded to be of type DatetimeBasedCursor
-                # add isintance check here if we want to create a Declarative IncrementingCountCursor
-                # or isinstance(
-                #     declarative_stream.retriever.stream_slicer, IncrementingCountCursor
-                # )
-                or isinstance(declarative_stream.retriever.stream_slicer, AsyncJobPartitionRouter)
-            )
+            and isinstance(declarative_stream.retriever.stream_slicer, (ConcurrentCursor, AsyncJobPartitionRouter))
         )
 
     @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
+        return declarative_stream.retriever
 
     @staticmethod
     def _select_streams(

airbyte_cdk/sources/declarative/extractors/record_filter.py

@@ -7,7 +7,7 @@
 from airbyte_cdk.sources.declarative.incremental import (
     DatetimeBasedCursor,
     GlobalSubstreamCursor,
-    PerPartitionWithGlobalCursor,
+    PerPartitionWithGlobalCursor, ConcurrentPerPartitionCursor,
 )
 from airbyte_cdk.sources.declarative.interpolation.interpolated_boolean import InterpolatedBoolean
 from airbyte_cdk.sources.types import Config, Record, StreamSlice, StreamState

airbyte_cdk/sources/declarative/manifest_declarative_source.py

@@ -36,9 +36,6 @@
 from airbyte_cdk.sources.declarative.checks.connection_checker import ConnectionChecker
 from airbyte_cdk.sources.declarative.declarative_source import DeclarativeSource
 from airbyte_cdk.sources.declarative.interpolation import InterpolatedBoolean
-from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
-    ConditionalStreams as ConditionalStreamsModel,
-)
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
     DeclarativeStream as DeclarativeStreamModel,
 )
@@ -297,7 +294,7 @@ def connection_checker(self) -> ConnectionChecker:
                 f"Expected to generate a ConnectionChecker component, but received {check_stream.__class__}"
             )
 
-    def streams(self, config: Mapping[str, Any]) -> List[Stream]:
+    def streams(self, config: Mapping[str, Any]) -> List[Stream]:  # FIXME with the recent change, this method can return Stream and AbstractStream which means it does not align with AbstractSource interface anymore. How big of a deal is this?
         if self._spec_component:
             self._spec_component.validate_config(config)
 

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -1620,7 +1620,6 @@ def create_concurrent_cursor_from_perpartition_cursor(
                 stream_namespace=stream_namespace,
                 config=config,
                 message_repository=NoopMessageRepository(),
-                stream_state_migrations=stream_state_migrations,
             )
         )
 
@@ -1955,7 +1954,7 @@ def create_declarative_stream(
             )
             if combined_slicers and not isinstance(combined_slicers, Cursor):
                 raise ValueError(
-                    "Unsupported Slicer is used. PerPartitionWithGlobalCursor should be used here instead"
+                    "Unsupported Slicer is used. PerPartitionWithGlobalCursor should be used here instead"  # FIXME update error message
                 )
             cursor = (
                 combined_slicers
@@ -2090,15 +2089,15 @@ def create_declarative_stream(
                 options["name"] = model.name
             schema_loader = DefaultSchemaLoader(config=config, parameters=options)
 
-        if concurrent or (hasattr(model.retriever, "partition_router") and model.retriever.partition_router):
+        if concurrent or (hasattr(model.retriever, "partition_router") and model.retriever.partition_router and model.incremental_sync):
             stream_name = model.name or ""
             cursor = combined_slicers \
                 if combined_slicers and isinstance(combined_slicers, Cursor) \
                 else FinalStateCursor(stream_name, None, self._message_repository)
             partition_generator = StreamSlicerPartitionGenerator(
                 DeclarativePartitionFactory(
                     stream_name,
-                    schema_loader.get_json_schema(),
+                    schema_loader,
                     retriever,
                     self._message_repository,
                 ),
@@ -2107,12 +2106,13 @@ def create_declarative_stream(
             return DefaultStream(
                 partition_generator=partition_generator,
                 name=stream_name,
-                json_schema=schema_loader.get_json_schema(),
+                json_schema=lambda: schema_loader.get_json_schema(),
                 availability_strategy=AlwaysAvailableAvailabilityStrategy(),  # FIXME it seems this is what we do in the ConcurrentDeclarativeSource but it feels wrong
                 primary_key=get_primary_key_from_stream(primary_key),
                 cursor_field=cursor.cursor_field.cursor_field_key if hasattr(cursor, "cursor_field") else "",  # FIXME we should have the cursor field has part of the interface of cursor
                 logger=logging.getLogger(f"airbyte.{stream_name}"),  # FIXME this is a breaking change compared to the old implementation,
                 cursor=cursor,
+                supports_file_transfer=hasattr(model, "file_uploader") and bool(model.file_uploader)
             )
         return DeclarativeStream(
             name=model.name or "",
@@ -2188,14 +2188,26 @@ def _build_incremental_cursor(
                 partition_router=stream_slicer,
             )
         elif model.incremental_sync:
-            return self.create_concurrent_cursor_from_datetime_based_cursor(  # type: ignore # This is a known issue that we are creating and returning a ConcurrentCursor which does not technically implement the (low-code) StreamSlicer. However, (low-code) StreamSlicer and ConcurrentCursor both implement StreamSlicer.stream_slices() which is the primary method needed for checkpointing
-                model_type=DatetimeBasedCursorModel,
-                component_definition=model.incremental_sync.__dict__,
-                stream_name=model.name or "",
-                stream_namespace=None,
-                config=config or {},
-                stream_state_migrations=model.state_migrations,
-            )
+            if type(model.incremental_sync) == IncrementingCountCursorModel:
+                return self.create_concurrent_cursor_from_incrementing_count_cursor(
+                    model_type=IncrementingCountCursorModel,
+                    component_definition=model.incremental_sync.__dict__,
+                    stream_name=model.name or "",
+                    stream_namespace=None,
+                    config=config or {},
+                    stream_state_migrations=model.state_migrations,
+                )
+            elif type(model.incremental_sync) == DatetimeBasedCursorModel:
+                return self.create_concurrent_cursor_from_datetime_based_cursor(  # type: ignore # This is a known issue that we are creating and returning a ConcurrentCursor which does not technically implement the (low-code) StreamSlicer. However, (low-code) StreamSlicer and ConcurrentCursor both implement StreamSlicer.stream_slices() which is the primary method needed for checkpointing
+                    model_type=type(model.incremental_sync),
+                    component_definition=model.incremental_sync.__dict__,
+                    stream_name=model.name or "",
+                    stream_namespace=None,
+                    config=config or {},
+                    stream_state_migrations=model.state_migrations,
+                )
+            else:
+                raise ValueError(f"Incremental sync of type {type(model.incremental_sync)} is not supported")
         return None
 
     def _build_resumable_cursor(
@@ -3235,9 +3247,6 @@ def _get_url() -> str:
             config=config,
         )
 
-        # Define cursor only if per partition or common incremental support is needed
-        cursor = stream_slicer if isinstance(stream_slicer, DeclarativeCursor) else None
-
         if (
             not isinstance(stream_slicer, (ConcurrentCursor, ConcurrentPerPartitionCursor))
             or type(stream_slicer) not in [ConcurrentCursor, ConcurrentPerPartitionCursor]
@@ -3260,6 +3269,7 @@ def _get_url() -> str:
                 ),
             )
 
+        cursor = stream_slicer if isinstance(stream_slicer, (ConcurrentCursor, DeclarativeCursor)) else None
         cursor_used_for_stop_condition = cursor if stop_condition_on_cursor else None
         paginator = (
             self._create_component_from_model(
@@ -3311,7 +3321,7 @@ def _get_url() -> str:
                 record_selector=record_selector,
                 stream_slicer=stream_slicer,
                 request_option_provider=request_options_provider,
-                cursor=cursor,
+                cursor=None,
                 config=config,
                 ignore_stream_slicer_parameters_on_paginated_requests=ignore_stream_slicer_parameters_on_paginated_requests,
                 parameters=model.parameters or {},
@@ -3325,7 +3335,7 @@ def _get_url() -> str:
             record_selector=record_selector,
             stream_slicer=stream_slicer,
             request_option_provider=request_options_provider,
-            cursor=cursor,
+            cursor=None,
             config=config,
             ignore_stream_slicer_parameters_on_paginated_requests=ignore_stream_slicer_parameters_on_paginated_requests,
             additional_query_properties=query_properties,
@@ -3713,7 +3723,7 @@ def _create_message_repository_substream_wrapper(
         self, model: ParentStreamConfigModel, config: Config, **kwargs: Any
     ) -> Any:
         # getting the parent state
-        child_state = self._connector_state_manager.get_stream_state(kwargs["stream_name"], None)
+        child_state = self._connector_state_manager.get_stream_state(kwargs["stream_name"], None)  # FIXME adding `stream_name` as a parameter means it will be a breaking change. I assume this is mostly called internally so I don't think we need to bother that much about this but still raising the flag
         if model.incremental_dependency and child_state:
             parent_stream_name = model.stream.name or ""
             parent_state = ConcurrentPerPartitionCursor.get_parent_state(child_state, parent_stream_name)
@@ -4070,7 +4080,7 @@ def create_grouping_partition_router(
         self, model: GroupingPartitionRouterModel, config: Config, **kwargs: Any
     ) -> GroupingPartitionRouter:
         underlying_router = self._create_component_from_model(
-            model=model.underlying_partition_router, config=config
+            model=model.underlying_partition_router, config=config, **kwargs,
         )
         if model.group_size < 1:
             raise ValueError(f"Group size must be greater than 0, got {model.group_size}")

airbyte_cdk/sources/declarative/partition_routers/substream_partition_router.py

@@ -23,7 +23,7 @@
 from airbyte_cdk.sources.types import Config, StreamSlice, StreamState
 
 if TYPE_CHECKING:
-    from airbyte_cdk.sources.streams.concurrent.default_stream import DefaultStream
+    from airbyte_cdk.sources.streams.concurrent.abstract_stream import AbstractStream
 
 
 def iterate_with_last_flag(generator: Iterable[Partition]) -> Iterable[tuple[Partition, bool]]:
@@ -55,7 +55,7 @@ class ParentStreamConfig:
     incremental_dependency (bool): Indicates if the parent stream should be read incrementally.
     """
 
-    stream: "DefaultStream"
+    stream: "AbstractStream"
     parent_key: Union[InterpolatedString, str]
     partition_field: Union[InterpolatedString, str]
     config: Config

airbyte_cdk/sources/declarative/schema/dynamic_schema_loader.py

@@ -5,6 +5,7 @@
 
 from copy import deepcopy
 from dataclasses import InitVar, dataclass, field
+from functools import lru_cache
 from typing import Any, List, Mapping, MutableMapping, Optional, Union
 
 import dpath
@@ -129,6 +130,7 @@ class DynamicSchemaLoader(SchemaLoader):
     schema_transformations: List[RecordTransformation] = field(default_factory=lambda: [])
     schema_filter: Optional[RecordFilter] = None
 
+    @lru_cache(maxsize=None)
     def get_json_schema(self) -> Mapping[str, Any]:
         """
         Constructs a JSON Schema based on retrieved data.
@@ -298,3 +300,36 @@ def _extract_data(
         ]
 
         return dpath.get(body, path, default=default)  # type: ignore # extracted will be a MutableMapping, given input data structure
+
+    def __hash__(self) -> int:
+        """
+        Autogenerated by cursor
+
+        Returns a hash of the DynamicSchemaLoader instance based on its configuration.
+        
+        Note: This implementation handles potentially unhashable components by converting
+        them to hashable types or using fallback strategies.
+        """
+        try:
+            # Convert config to hashable type if it's a dict
+            config_hash = frozenset(self.config.items()) if isinstance(self.config, dict) else self.config
+            
+            # Convert list of transformations to tuple for hashing
+            transformations_hash = tuple(self.schema_transformations)
+            
+            return hash((
+                self.retriever,
+                config_hash,
+                self.schema_type_identifier,
+                transformations_hash,
+                self.schema_filter,
+            ))
+        except TypeError:
+            # Fallback to id-based hashing if any component is not hashable
+            return hash((
+                id(self.retriever),
+                id(self.config),
+                id(self.schema_type_identifier),
+                tuple(id(t) for t in self.schema_transformations),
+                id(self.schema_filter) if self.schema_filter else None,
+            ))

airbyte_cdk/sources/declarative/schema/json_file_schema_loader.py

@@ -6,6 +6,7 @@
 import pkgutil
 import sys
 from dataclasses import InitVar, dataclass, field
+from functools import lru_cache
 from typing import Any, Mapping, Tuple, Union
 
 from airbyte_cdk.sources.declarative.interpolation.interpolated_string import InterpolatedString