refactor regarding async stuff

Author: maxime.c

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -94,16 +94,13 @@
     ClientSideIncrementalRecordFilterDecorator,
 )
 from airbyte_cdk.sources.declarative.incremental import (
-    ChildPartitionResumableFullRefreshCursor,
     ConcurrentCursorFactory,
     ConcurrentPerPartitionCursor,
     CursorFactory,
     DatetimeBasedCursor,
     DeclarativeCursor,
     GlobalSubstreamCursor,
-    PerPartitionCursor,
     PerPartitionWithGlobalCursor,
-    ResumableFullRefreshCursor,
 )
 from airbyte_cdk.sources.declarative.interpolation import InterpolatedString
 from airbyte_cdk.sources.declarative.interpolation.interpolated_mapping import InterpolatedMapping
@@ -446,10 +443,6 @@
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
     ZipfileDecoder as ZipfileDecoderModel,
 )
-from airbyte_cdk.sources.declarative.parsers.custom_code_compiler import (
-    COMPONENTS_MODULE_NAME,
-    SDM_COMPONENTS_MODULE_NAME,
-)
 from airbyte_cdk.sources.declarative.partition_routers import (
     CartesianProductStreamSlicer,
     GroupingPartitionRouter,
@@ -508,7 +501,7 @@
     RequestOptionsProvider,
 )
 from airbyte_cdk.sources.declarative.requesters.request_path import RequestPath
-from airbyte_cdk.sources.declarative.requesters.requester import HttpMethod, Requester
+from airbyte_cdk.sources.declarative.requesters.requester import HttpMethod
 from airbyte_cdk.sources.declarative.resolvers import (
     ComponentMappingDefinition,
     ConfigComponentsResolver,
@@ -1941,10 +1934,10 @@ def create_declarative_stream(
 
         primary_key = model.primary_key.__root__ if model.primary_key else None
 
-        stream_slicer = self._build_stream_slicer_from_partition_router(
+        partition_router = self._build_stream_slicer_from_partition_router(
             model.retriever, config, stream_name=model.name
         )
-        concurrent_cursor = self._build_concurrent_cursor(model, stream_slicer, config)
+        concurrent_cursor = self._build_concurrent_cursor(model, partition_router, config)
 
         if model.incremental_sync and isinstance(model.incremental_sync, DatetimeBasedCursorModel):
             cursor_model = model.incremental_sync
@@ -2019,7 +2012,9 @@ def create_declarative_stream(
             primary_key=primary_key,
             stream_slicer=combined_slicers,
             request_options_provider=request_options_provider,
-            stop_condition_cursor=concurrent_cursor if self._is_stop_condition_on_cursor(model) else None,
+            stop_condition_cursor=concurrent_cursor
+            if self._is_stop_condition_on_cursor(model)
+            else None,
             client_side_incremental_sync={"cursor": concurrent_cursor}
             if self._is_client_side_filtering_enabled(model)
             else None,
@@ -2055,7 +2050,10 @@ def create_declarative_stream(
             schema_loader = DefaultSchemaLoader(config=config, parameters=options)
 
         if (
-            (isinstance(combined_slicers, PartitionRouter) or isinstance(concurrent_cursor, ConcurrentCursor))
+            (
+                isinstance(combined_slicers, PartitionRouter)
+                or isinstance(concurrent_cursor, ConcurrentCursor)
+            )
             and not is_parent
             and not self._emit_connector_builder_messages
         ):
@@ -2067,7 +2065,9 @@ def create_declarative_stream(
             # We specifically exclude parent streams here because SubstreamPartitionRouter has not been updated yet
             # We specifically exclude Connector Builder stuff for now as Brian is working on this anyway
 
+            stream_name = model.name or ""
             stream_slicer = concurrent_cursor
+            cursor = FinalStateCursor(stream_name, None, self._message_repository)
             if isinstance(retriever, AsyncRetriever):
                 # The AsyncRetriever only ever worked with a cursor from the concurrent package. Hence, the method
                 # `_build_incremental_cursor` which we would usually think would return only declarative stuff has a
@@ -2076,10 +2076,13 @@ def create_declarative_stream(
                 # We can't re-use `concurrent_cursor` because it is a different instance than the one passed in
                 # AsyncJobPartitionRouter.
                 stream_slicer = retriever.stream_slicer
+                if isinstance(combined_slicers, Cursor):
+                    cursor = combined_slicers
             elif isinstance(combined_slicers, PartitionRouter):
                 stream_slicer = combined_slicers
+            else:
+                cursor = concurrent_cursor
 
-            stream_name = model.name or ""
             partition_generator = StreamSlicerPartitionGenerator(
                 DeclarativePartitionFactory(
                     stream_name,
@@ -2089,16 +2092,17 @@ def create_declarative_stream(
                 ),
                 stream_slicer,
             )
-            cursor = concurrent_cursor if concurrent_cursor else FinalStateCursor(stream_name, None, self._message_repository)
+
             return DefaultStream(
                 partition_generator=partition_generator,
                 name=stream_name,
                 json_schema=schema_loader.get_json_schema,
                 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,
-                # FIXME we should have the cursor field has part of the interface of cursor
+                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,
+                # FIXME this is a breaking change compared to the old implementation which used the source name instead
                 cursor=cursor,
             )
 
@@ -2121,18 +2125,18 @@ def create_declarative_stream(
             parameters=model.parameters or {},
         )
 
-    def _is_stop_condition_on_cursor(self, model):
+    def _is_stop_condition_on_cursor(self, model: DeclarativeStreamModel) -> bool:
         return (
-                model.incremental_sync
-                and hasattr(model.incremental_sync, "is_data_feed")
-                and model.incremental_sync.is_data_feed
+            model.incremental_sync
+            and hasattr(model.incremental_sync, "is_data_feed")
+            and model.incremental_sync.is_data_feed
         )
 
-    def _is_client_side_filtering_enabled(self, model):
+    def _is_client_side_filtering_enabled(self, model: DeclarativeStreamModel) -> bool:
         client_side_filtering_enabled = (
-                model.incremental_sync
-                and hasattr(model.incremental_sync, "is_client_side_incremental")
-                and model.incremental_sync.is_client_side_incremental
+            model.incremental_sync
+            and hasattr(model.incremental_sync, "is_client_side_incremental")
+            and model.incremental_sync.is_client_side_incremental
         )
         return client_side_filtering_enabled
 

unit_tests/sources/declarative/parsers/test_model_to_component_factory.py

@@ -373,38 +373,33 @@ def test_full_config_stream():
 
     assert isinstance(retriever.record_selector.extractor, DpathExtractor)
     assert isinstance(retriever.record_selector.extractor.decoder, JsonDecoder)
-    assert [
-        fp.eval(input_config) for fp in retriever.record_selector.extractor._field_path
-    ] == ["lists"]
+    assert [fp.eval(input_config) for fp in retriever.record_selector.extractor._field_path] == [
+        "lists"
+    ]
 
     assert isinstance(retriever.record_selector.record_filter, RecordFilter)
     assert (
-            retriever.record_selector.record_filter._filter_interpolator.condition
-            == "{{ record['id'] > stream_state['id'] }}"
+        retriever.record_selector.record_filter._filter_interpolator.condition
+        == "{{ record['id'] > stream_state['id'] }}"
     )
 
     assert isinstance(retriever.paginator, DefaultPaginator)
     assert isinstance(retriever.paginator.decoder, PaginationDecoderDecorator)
     assert retriever.paginator.page_size_option.field_name.eval(input_config) == "page_size"
-    assert (
-            retriever.paginator.page_size_option.inject_into
-            == RequestOptionType.request_parameter
-    )
+    assert retriever.paginator.page_size_option.inject_into == RequestOptionType.request_parameter
     assert isinstance(retriever.paginator.page_token_option, RequestPath)
     assert retriever.paginator.url_base.string == "https://api.sendgrid.com/v3/"
     assert retriever.paginator.url_base.default == "https://api.sendgrid.com/v3/"
 
     assert isinstance(retriever.paginator.pagination_strategy, CursorPaginationStrategy)
-    assert isinstance(
-        retriever.paginator.pagination_strategy.decoder, PaginationDecoderDecorator
-    )
+    assert isinstance(retriever.paginator.pagination_strategy.decoder, PaginationDecoderDecorator)
     assert (
-            retriever.paginator.pagination_strategy._cursor_value.string
-            == "{{ response._metadata.next }}"
+        retriever.paginator.pagination_strategy._cursor_value.string
+        == "{{ response._metadata.next }}"
     )
     assert (
-            retriever.paginator.pagination_strategy._cursor_value.default
-            == "{{ response._metadata.next }}"
+        retriever.paginator.pagination_strategy._cursor_value.default
+        == "{{ response._metadata.next }}"
     )
     assert retriever.paginator.pagination_strategy.page_size == 10
 
@@ -416,24 +411,20 @@ def test_full_config_stream():
 
     assert isinstance(retriever.request_option_provider, DatetimeBasedRequestOptionsProvider)
     assert (
-            retriever.request_option_provider.start_time_option.inject_into
-            == RequestOptionType.request_parameter
+        retriever.request_option_provider.start_time_option.inject_into
+        == RequestOptionType.request_parameter
     )
     assert (
-            retriever.request_option_provider.start_time_option.field_name.eval(
-            config=input_config
-        )
-            == "after"
+        retriever.request_option_provider.start_time_option.field_name.eval(config=input_config)
+        == "after"
     )
     assert (
-            retriever.request_option_provider.end_time_option.inject_into
-            == RequestOptionType.request_parameter
+        retriever.request_option_provider.end_time_option.inject_into
+        == RequestOptionType.request_parameter
     )
     assert (
-            retriever.request_option_provider.end_time_option.field_name.eval(
-            config=input_config
-        )
-            == "before"
+        retriever.request_option_provider.end_time_option.field_name.eval(config=input_config)
+        == "before"
     )
     assert retriever.request_option_provider._partition_field_start.string == "start_time"
     assert retriever.request_option_provider._partition_field_end.string == "end_time"
@@ -444,9 +435,7 @@ def test_full_config_stream():
     assert isinstance(
         retriever.requester.request_options_provider, InterpolatedRequestOptionsProvider
     )
-    assert (
-            retriever.requester.request_options_provider.request_parameters.get("unit") == "day"
-    )
+    assert retriever.requester.request_options_provider.request_parameters.get("unit") == "day"
 
     checker = factory.create_component(
         model_type=CheckStreamModel, component_definition=manifest["check"], config=input_config
@@ -1118,7 +1107,8 @@ def test_incremental_data_feed():
     )
 
     assert isinstance(
-        get_retriever(stream).paginator.pagination_strategy, StopConditionPaginationStrategyDecorator
+        get_retriever(stream).paginator.pagination_strategy,
+        StopConditionPaginationStrategyDecorator,
     )
 
 
@@ -4149,9 +4139,7 @@ def test_simple_retriever_with_query_properties():
     assert property_chunking.property_limit_type == PropertyLimitType.property_count
     assert property_chunking.property_limit == 3
 
-    merge_strategy = (
-        retriever.additional_query_properties.property_chunking.record_merge_strategy
-    )
+    merge_strategy = retriever.additional_query_properties.property_chunking.record_merge_strategy
     assert isinstance(merge_strategy, GroupByKey)
     assert merge_strategy.key == ["id"]