feat(low-code cdk): add StateDelegatingStream (#318)

Co-authored-by: octavia-squidington-iii <[email protected]>

Author: lazebnyi

airbyte_cdk/sources/declarative/concurrent_declarative_source.py

@@ -162,6 +162,10 @@ def read(
         else:
             filtered_catalog = catalog
 
+        # It is no need run read for synchronous streams if they are not exists.
+        if not filtered_catalog.streams:
+            return
+
         yield from super().read(logger, config, filtered_catalog, state)
 
     def discover(self, logger: logging.Logger, config: Mapping[str, Any]) -> AirbyteCatalog:
@@ -201,6 +205,18 @@ def _group_streams(
             # Some low-code sources use a combination of DeclarativeStream and regular Python streams. We can't inspect
             # these legacy Python streams the way we do low-code streams to determine if they are concurrent compatible,
             # so we need to treat them as synchronous
+
+            if name_to_stream_mapping[declarative_stream.name]["type"] == "StateDelegatingStream":
+                stream_state = self._connector_state_manager.get_stream_state(
+                    stream_name=declarative_stream.name, namespace=declarative_stream.namespace
+                )
+
+                name_to_stream_mapping[declarative_stream.name] = (
+                    name_to_stream_mapping[declarative_stream.name]["incremental_stream"]
+                    if stream_state
+                    else name_to_stream_mapping[declarative_stream.name]["full_refresh_stream"]
+                )
+
             if isinstance(declarative_stream, DeclarativeStream) and (
                 name_to_stream_mapping[declarative_stream.name]["retriever"]["type"]
                 == "SimpleRetriever"

airbyte_cdk/sources/declarative/declarative_component_schema.yaml

@@ -24,7 +24,9 @@ properties:
   streams:
     type: array
     items:
-      "$ref": "#/definitions/DeclarativeStream"
+      anyOf:
+        - "$ref": "#/definitions/DeclarativeStream"
+        - "$ref": "#/definitions/StateDelegatingStream"
   dynamic_streams:
     type: array
     items:
@@ -2881,7 +2883,9 @@ definitions:
       stream:
         title: Parent Stream
         description: Reference to the parent stream.
-        "$ref": "#/definitions/DeclarativeStream"
+        anyOf:
+          - "$ref": "#/definitions/DeclarativeStream"
+          - "$ref": "#/definitions/StateDelegatingStream"
       partition_field:
         title: Current Parent Key Value Identifier
         description: While iterating over parent records during a sync, the parent_key value can be referenced by using this field.
@@ -3154,6 +3158,36 @@ definitions:
       $parameters:
         type: object
         additionalProperties: true
+  StateDelegatingStream:
+    description: (This component is experimental. Use at your own risk.) Orchestrate the retriever's usage based on the state value.
+    type: object
+    required:
+      - type
+      - name
+      - full_refresh_stream
+      - incremental_stream
+    properties:
+      type:
+        type: string
+        enum: [ StateDelegatingStream ]
+      name:
+        title: Name
+        description: The stream name.
+        type: string
+        default: ""
+        example:
+          - "Users"
+      full_refresh_stream:
+        title: Retriever
+        description: Component used to coordinate how records are extracted across stream slices and request pages when the state is empty or not provided.
+        "$ref": "#/definitions/DeclarativeStream"
+      incremental_stream:
+        title: Retriever
+        description: Component used to coordinate how records are extracted across stream slices and request pages when the state provided.
+        "$ref": "#/definitions/DeclarativeStream"
+      $parameters:
+        type: object
+        additionalProperties: true
   SimpleRetriever:
     description: Retrieves records by synchronously sending requests to fetch records. The retriever acts as an orchestrator between the requester, the record selector, the paginator, and the partition router.
     type: object

airbyte_cdk/sources/declarative/manifest_declarative_source.py

@@ -30,6 +30,9 @@
     DeclarativeStream as DeclarativeStreamModel,
 )
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import Spec as SpecModel
+from airbyte_cdk.sources.declarative.models.declarative_component_schema import (
+    StateDelegatingStream as StateDelegatingStreamModel,
+)
 from airbyte_cdk.sources.declarative.parsers.custom_code_compiler import (
     get_registered_components_module,
 )
@@ -146,7 +149,9 @@ def streams(self, config: Mapping[str, Any]) -> List[Stream]:
 
         source_streams = [
             self._constructor.create_component(
-                DeclarativeStreamModel,
+                StateDelegatingStreamModel
+                if stream_config.get("type") == StateDelegatingStreamModel.__name__
+                else DeclarativeStreamModel,
                 stream_config,
                 config,
                 emit_connector_builder_messages=self._emit_connector_builder_messages,
@@ -165,7 +170,15 @@ def _initialize_cache_for_parent_streams(
         def update_with_cache_parent_configs(parent_configs: list[dict[str, Any]]) -> None:
             for parent_config in parent_configs:
                 parent_streams.add(parent_config["stream"]["name"])
-                parent_config["stream"]["retriever"]["requester"]["use_cache"] = True
+                if parent_config["stream"]["type"] == "StateDelegatingStream":
+                    parent_config["stream"]["full_refresh_stream"]["retriever"]["requester"][
+                        "use_cache"
+                    ] = True
+                    parent_config["stream"]["incremental_stream"]["retriever"]["requester"][
+                        "use_cache"
+                    ] = True
+                else:
+                    parent_config["stream"]["retriever"]["requester"]["use_cache"] = True
 
         for stream_config in stream_configs:
             if stream_config.get("incremental_sync", {}).get("parent_stream"):
@@ -188,7 +201,15 @@ def update_with_cache_parent_configs(parent_configs: list[dict[str, Any]]) -> No
 
         for stream_config in stream_configs:
             if stream_config["name"] in parent_streams:
-                stream_config["retriever"]["requester"]["use_cache"] = True
+                if stream_config["type"] == "StateDelegatingStream":
+                    stream_config["full_refresh_stream"]["retriever"]["requester"]["use_cache"] = (
+                        True
+                    )
+                    stream_config["incremental_stream"]["retriever"]["requester"]["use_cache"] = (
+                        True
+                    )
+                else:
+                    stream_config["retriever"]["requester"]["use_cache"] = True
 
         return stream_configs
 

airbyte_cdk/sources/declarative/models/declarative_component_schema.py

@@ -1860,7 +1860,7 @@ class Config:
 
     type: Literal["DeclarativeSource"]
     check: Union[CheckStream, CheckDynamicStream]
-    streams: List[DeclarativeStream]
+    streams: List[Union[DeclarativeStream, StateDelegatingStream]]
     dynamic_streams: Optional[List[DynamicDeclarativeStream]] = None
     version: str = Field(
         ...,
@@ -1892,7 +1892,7 @@ class Config:
 
     type: Literal["DeclarativeSource"]
     check: Union[CheckStream, CheckDynamicStream]
-    streams: Optional[List[DeclarativeStream]] = None
+    streams: Optional[List[Union[DeclarativeStream, StateDelegatingStream]]] = None
     dynamic_streams: List[DynamicDeclarativeStream]
     version: str = Field(
         ...,
@@ -2211,7 +2211,7 @@ class ParentStreamConfig(BaseModel):
         examples=["id", "{{ config['parent_record_id'] }}"],
         title="Parent Key",
     )
-    stream: DeclarativeStream = Field(
+    stream: Union[DeclarativeStream, StateDelegatingStream] = Field(
         ..., description="Reference to the parent stream.", title="Parent Stream"
     )
     partition_field: str = Field(
@@ -2238,6 +2238,22 @@ class ParentStreamConfig(BaseModel):
     parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters")
 
 
+class StateDelegatingStream(BaseModel):
+    type: Literal["StateDelegatingStream"]
+    name: str = Field(..., description="The stream name.", example=["Users"], title="Name")
+    full_refresh_stream: DeclarativeStream = Field(
+        ...,
+        description="Component used to coordinate how records are extracted across stream slices and request pages when the state is empty or not provided.",
+        title="Retriever",
+    )
+    incremental_stream: DeclarativeStream = Field(
+        ...,
+        description="Component used to coordinate how records are extracted across stream slices and request pages when the state provided.",
+        title="Retriever",
+    )
+    parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters")
+
+
 class SimpleRetriever(BaseModel):
     type: Literal["SimpleRetriever"]
     record_selector: RecordSelector = Field(
@@ -2423,5 +2439,6 @@ class DynamicDeclarativeStream(BaseModel):
 DeclarativeStream.update_forward_refs()
 SessionTokenAuthenticator.update_forward_refs()
 DynamicSchemaLoader.update_forward_refs()
+ParentStreamConfig.update_forward_refs()
 SimpleRetriever.update_forward_refs()
 AsyncRetriever.update_forward_refs()