feat(low-code): add UnionPartitionRouter component

airbyte_cdk/__init__.py

@@ -114,6 +114,7 @@
     CartesianProductStreamSlicer,
     SinglePartitionRouter,
     SubstreamPartitionRouter,
+    UnionPartitionRouter,
 )
 from .sources.declarative.partition_routers.substream_partition_router import ParentStreamConfig
 from .sources.declarative.requesters import HttpRequester, Requester
@@ -272,6 +273,7 @@
     "StopConditionPaginationStrategyDecorator",
     "StreamSlice",
     "SubstreamPartitionRouter",
+    "UnionPartitionRouter",
     "YamlDeclarativeSource",
     # Entrypoint
     "launch",

airbyte_cdk/sources/concurrent_source/concurrent_read_processor.py

@@ -19,6 +19,9 @@
 from airbyte_cdk.sources.declarative.partition_routers.substream_partition_router import (
     SubstreamPartitionRouter,
 )
+from airbyte_cdk.sources.declarative.partition_routers.union_partition_router import (
+    UnionPartitionRouter,
+)
 from airbyte_cdk.sources.message import MessageRepository
 from airbyte_cdk.sources.streams.concurrent.abstract_stream import AbstractStream
 from airbyte_cdk.sources.streams.concurrent.default_stream import DefaultStream
@@ -438,14 +441,18 @@ def _collect_all_parent_stream_names(self, stream_name: str) -> Set[str]:
         partition_router = (
             stream.get_partition_router() if isinstance(stream, DefaultStream) else None
         )
-        if isinstance(partition_router, GroupingPartitionRouter):
-            partition_router = partition_router.underlying_partition_router
-
-        if isinstance(partition_router, SubstreamPartitionRouter):
-            for parent_config in partition_router.parent_stream_configs:
-                parent_name = parent_config.stream.name
-                parent_names.add(parent_name)
-                parent_names.update(self._collect_all_parent_stream_names(parent_name))
+        routers = [partition_router] if partition_router else []
+        while routers:
+            router = routers.pop()
+            if isinstance(router, GroupingPartitionRouter):
+                routers.append(router.underlying_partition_router)
+            elif isinstance(router, UnionPartitionRouter):
+                routers.extend(router.partition_routers)
+            elif isinstance(router, SubstreamPartitionRouter):
+                for parent_config in router.parent_stream_configs:
+                    parent_name = parent_config.stream.name
+                    parent_names.add(parent_name)
+                    parent_names.update(self._collect_all_parent_stream_names(parent_name))
 
         return parent_names
 

airbyte_cdk/sources/declarative/concurrent_declarative_source.py

@@ -82,6 +82,9 @@
 from airbyte_cdk.sources.declarative.partition_routers.substream_partition_router import (
     SubstreamPartitionRouter,
 )
+from airbyte_cdk.sources.declarative.partition_routers.union_partition_router import (
+    UnionPartitionRouter,
+)
 from airbyte_cdk.sources.declarative.resolvers import COMPONENTS_RESOLVER_TYPE_MAPPING
 from airbyte_cdk.sources.declarative.spec.spec import Spec
 from airbyte_cdk.sources.declarative.types import Config, ConnectionDefinition
@@ -461,15 +464,18 @@ def _collect_all_ancestor_names(stream_name: str) -> Set[str]:
             inst = stream_name_to_instance.get(stream_name)
             if not isinstance(inst, DefaultStream):
                 return ancestors
-            router = inst.get_partition_router()
-            if isinstance(router, GroupingPartitionRouter):
-                router = router.underlying_partition_router
-            if not isinstance(router, SubstreamPartitionRouter):
-                return ancestors
-            for parent_config in router.parent_stream_configs:
-                parent_name = parent_config.stream.name
-                ancestors.add(parent_name)
-                ancestors.update(_collect_all_ancestor_names(parent_name))
+            routers = [inst.get_partition_router()]
+            while routers:
+                router = routers.pop()
+                if isinstance(router, GroupingPartitionRouter):
+                    routers.append(router.underlying_partition_router)
+                elif isinstance(router, UnionPartitionRouter):
+                    routers.extend(router.partition_routers)
+                elif isinstance(router, SubstreamPartitionRouter):
+                    for parent_config in router.parent_stream_configs:
+                        parent_name = parent_config.stream.name
+                        ancestors.add(parent_name)
+                        ancestors.update(_collect_all_ancestor_names(parent_name))
             return ancestors
 
         for stream in streams:
@@ -532,14 +538,23 @@ def update_with_cache_parent_configs(
             elif stream_config.get("retriever", {}).get("partition_router", {}):
                 partition_router = stream_config["retriever"]["partition_router"]
 
-                if isinstance(partition_router, dict) and partition_router.get(
-                    "parent_stream_configs"
-                ):
-                    update_with_cache_parent_configs(partition_router["parent_stream_configs"])
-                elif isinstance(partition_router, list):
-                    for router in partition_router:
-                        if router.get("parent_stream_configs"):
-                            update_with_cache_parent_configs(router["parent_stream_configs"])
+                routers = (
+                    list(partition_router)
+                    if isinstance(partition_router, list)
+                    else [partition_router]
+                )
+                while routers:
+                    router = routers.pop()
+                    if not isinstance(router, dict):
+                        continue
+                    if router.get("parent_stream_configs"):
+                        update_with_cache_parent_configs(router["parent_stream_configs"])
+                    # Descend into composed partition routers (e.g. UnionPartitionRouter's
+                    # partition_routers or GroupingPartitionRouter's underlying_partition_router)
+                    # so nested parent streams also get caching enabled.
+                    routers.extend(router.get("partition_routers") or [])
+                    if router.get("underlying_partition_router"):
+                        routers.append(router["underlying_partition_router"])
 
         for stream_config in stream_configs:
             if stream_config["name"] in parent_streams:

airbyte_cdk/sources/declarative/declarative_component_schema.yaml

@@ -4002,6 +4002,7 @@ definitions:
           - "$ref": "#/definitions/SubstreamPartitionRouter"
           - "$ref": "#/definitions/ListPartitionRouter"
           - "$ref": "#/definitions/GroupingPartitionRouter"
+          - "$ref": "#/definitions/UnionPartitionRouter"
           - "$ref": "#/definitions/CustomPartitionRouter"
           - type: array
             title: Multiple Partition Routers
@@ -4010,6 +4011,7 @@ definitions:
                 - "$ref": "#/definitions/SubstreamPartitionRouter"
                 - "$ref": "#/definitions/ListPartitionRouter"
                 - "$ref": "#/definitions/GroupingPartitionRouter"
+                - "$ref": "#/definitions/UnionPartitionRouter"
                 - "$ref": "#/definitions/CustomPartitionRouter"
       $parameters:
         type: object
@@ -4212,13 +4214,15 @@ definitions:
           - "$ref": "#/definitions/ListPartitionRouter"
           - "$ref": "#/definitions/SubstreamPartitionRouter"
           - "$ref": "#/definitions/GroupingPartitionRouter"
+          - "$ref": "#/definitions/UnionPartitionRouter"
           - "$ref": "#/definitions/CustomPartitionRouter"
           - type: array
             items:
               anyOf:
                 - "$ref": "#/definitions/ListPartitionRouter"
                 - "$ref": "#/definitions/SubstreamPartitionRouter"
                 - "$ref": "#/definitions/GroupingPartitionRouter"
+                - "$ref": "#/definitions/UnionPartitionRouter"
                 - "$ref": "#/definitions/CustomPartitionRouter"
       decoder:
         title: HTTP Response Format
@@ -4475,6 +4479,47 @@ definitions:
       $parameters:
         type: object
         additionalProperties: true
+  UnionPartitionRouter:
+    title: Union Partition Router
+    description: >
+      A partition router that yields the deduplicated union of the partitions produced by its child
+      partition routers. Every emitted partition is normalized to a single key defined by `partition_field`;
+      any other partition keys coming from a child router (such as a SubstreamPartitionRouter's `parent_slice`)
+      are moved into the slice's extra fields. The first occurrence of a partition value wins and later
+      duplicates are skipped.
+    type: object
+    required:
+      - type
+      - partition_field
+      - partition_routers
+    properties:
+      type:
+        type: string
+        enum: [UnionPartitionRouter]
+      partition_field:
+        title: Partition Field
+        description: The single partition key that all child partition routers' slices are normalized to. Each child router must emit this key in its partitions. Interpolation is evaluated once when the connector is built, using the connector config and $parameters.
+        type: string
+        interpolation_context:
+          - config
+          - parameters
+        examples:
+          - "repository"
+          - "{{ config['partition_field'] }}"
+      partition_routers:
+        title: Partition Routers
+        description: The child partition routers whose partitions are unioned. Request options are not supported on child partition routers; partition values should be consumed via interpolation (e.g. `stream_partition`).
+        type: array
+        minItems: 2
+        items:
+          anyOf:
+            - "$ref": "#/definitions/ListPartitionRouter"
+            - "$ref": "#/definitions/SubstreamPartitionRouter"
+            - "$ref": "#/definitions/UnionPartitionRouter"
+            - "$ref": "#/definitions/CustomPartitionRouter"
+      $parameters:
+        type: object
+        additionalProperties: true
   WaitUntilTimeFromHeader:
     title: Wait Until Time Defined In Response Header
     description: Extract time at which we can retry the request from response header and wait for the difference between now and that time.

airbyte_cdk/sources/declarative/migrations/legacy_to_per_partition_state_migration.py

@@ -1,12 +1,14 @@
 # Copyright (c) 2024 Airbyte, Inc., all rights reserved.
 
-from typing import Any, Mapping
+from typing import Any, Mapping, Union
 
 from airbyte_cdk.sources.declarative.interpolation.interpolated_string import InterpolatedString
 from airbyte_cdk.sources.declarative.migrations.state_migration import StateMigration
 from airbyte_cdk.sources.declarative.models import (
+    CustomPartitionRouter,
     DatetimeBasedCursor,
     SubstreamPartitionRouter,
+    UnionPartitionRouter,
 )
 from airbyte_cdk.sources.declarative.models.declarative_component_schema import ParentStreamConfig
 
@@ -34,7 +36,9 @@ class LegacyToPerPartitionStateMigration(StateMigration):
 
     def __init__(
         self,
-        partition_router: SubstreamPartitionRouter,
+        partition_router: Union[
+            SubstreamPartitionRouter, UnionPartitionRouter, CustomPartitionRouter
+        ],
         cursor: DatetimeBasedCursor,
         config: Mapping[str, Any],
         parameters: Mapping[str, Any],
@@ -44,14 +48,22 @@ def __init__(
         self._config = config
         self._parameters = parameters
         self._partition_key_field = InterpolatedString.create(
-            self._get_partition_field(self._partition_router), parameters=self._parameters
+            self._get_partition_field(partition_router), parameters=self._parameters
         ).eval(self._config)
         self._cursor_field = InterpolatedString.create(
             self._cursor.cursor_field, parameters=self._parameters
         ).eval(self._config)
 
-    def _get_partition_field(self, partition_router: SubstreamPartitionRouter) -> str:
-        parent_stream_config = partition_router.parent_stream_configs[0]
+    def _get_partition_field(
+        self,
+        partition_router: Union[
+            SubstreamPartitionRouter, UnionPartitionRouter, CustomPartitionRouter
+        ],
+    ) -> str:
+        if isinstance(partition_router, UnionPartitionRouter):
+            return partition_router.partition_field
+
+        parent_stream_config = partition_router.parent_stream_configs[0]  # type: ignore # custom partition routers are expected to expose parent_stream_configs
 
         # Retrieve the partition field with a condition, as properties are returned as a dictionary for custom components.
         partition_field = (
@@ -66,11 +78,14 @@ def should_migrate(self, stream_state: Mapping[str, Any]) -> bool:
         if _is_already_migrated(stream_state):
             return False
 
-        # There is exactly one parent stream
-        number_of_parent_streams = len(self._partition_router.parent_stream_configs)  # type: ignore # custom partition will introduce this attribute if needed
-        if number_of_parent_streams != 1:
-            # There should be exactly one parent stream
-            return False
+        # UnionPartitionRouter has no parent_stream_configs; its partitions are already
+        # normalized to a single partition field so the parent stream check does not apply.
+        if not isinstance(self._partition_router, UnionPartitionRouter):
+            # There is exactly one parent stream
+            number_of_parent_streams = len(self._partition_router.parent_stream_configs)  # type: ignore # custom partition will introduce this attribute if needed
+            if number_of_parent_streams != 1:
+                # There should be exactly one parent stream
+                return False
         """
         The expected state format is
         "<parent_key_id>" : {

airbyte_cdk/sources/declarative/models/declarative_component_schema.py

@@ -3065,12 +3065,14 @@ class SimpleRetriever(BaseModel):
             SubstreamPartitionRouter,
             ListPartitionRouter,
             GroupingPartitionRouter,
+            UnionPartitionRouter,
             CustomPartitionRouter,
             List[
                 Union[
                     SubstreamPartitionRouter,
                     ListPartitionRouter,
                     GroupingPartitionRouter,
+                    UnionPartitionRouter,
                     CustomPartitionRouter,
                 ]
             ],
@@ -3144,12 +3146,14 @@ class AsyncRetriever(BaseModel):
             ListPartitionRouter,
             SubstreamPartitionRouter,
             GroupingPartitionRouter,
+            UnionPartitionRouter,
             CustomPartitionRouter,
             List[
                 Union[
                     ListPartitionRouter,
                     SubstreamPartitionRouter,
                     GroupingPartitionRouter,
+                    UnionPartitionRouter,
                     CustomPartitionRouter,
                 ]
             ],
@@ -3246,6 +3250,29 @@ class GroupingPartitionRouter(BaseModel):
     parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters")
 
 
+class UnionPartitionRouter(BaseModel):
+    type: Literal["UnionPartitionRouter"]
+    partition_field: str = Field(
+        ...,
+        description="The single partition key that all child partition routers' slices are normalized to. Each child router must emit this key in its partitions. Interpolation is evaluated once when the connector is built, using the connector config and $parameters.",
+        examples=["repository", "{{ config['partition_field'] }}"],
+        title="Partition Field",
+    )
+    partition_routers: List[
+        Union[
+            ListPartitionRouter,
+            SubstreamPartitionRouter,
+            UnionPartitionRouter,
+            CustomPartitionRouter,
+        ]
+    ] = Field(
+        ...,
+        description="The child partition routers whose partitions are unioned. Request options are not supported on child partition routers; partition values should be consumed via interpolation (e.g. `stream_partition`).",
+        title="Partition Routers",
+    )
+    parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters")
+
+
 class HttpComponentsResolver(BaseModel):
     type: Literal["HttpComponentsResolver"]
     retriever: Union[SimpleRetriever, AsyncRetriever, CustomRetriever] = Field(
@@ -3295,3 +3322,4 @@ class DynamicDeclarativeStream(BaseModel):
 PropertiesFromEndpoint.update_forward_refs()
 SimpleRetriever.update_forward_refs()
 AsyncRetriever.update_forward_refs()
+UnionPartitionRouter.update_forward_refs()