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

@@ -13,12 +13,18 @@
 )
 from airbyte_cdk.sources.concurrent_source.stream_thread_exception import StreamThreadException
 from airbyte_cdk.sources.concurrent_source.thread_pool_manager import ThreadPoolManager
+from airbyte_cdk.sources.declarative.partition_routers.cartesian_product_stream_slicer import (
+    CartesianProductStreamSlicer,
+)
 from airbyte_cdk.sources.declarative.partition_routers.grouping_partition_router import (
     GroupingPartitionRouter,
 )
 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 +444,20 @@ 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 is not None 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, CartesianProductStreamSlicer):
+                routers.extend(router.stream_slicers)
+            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

@@ -76,12 +76,18 @@
 from airbyte_cdk.sources.declarative.parsers.model_to_component_factory import (
     ModelToComponentFactory,
 )
+from airbyte_cdk.sources.declarative.partition_routers.cartesian_product_stream_slicer import (
+    CartesianProductStreamSlicer,
+)
 from airbyte_cdk.sources.declarative.partition_routers.grouping_partition_router import (
     GroupingPartitionRouter,
 )
 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
@@ -464,15 +470,21 @@ 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))
+            partition_router = inst.get_partition_router()
+            routers = [partition_router] if partition_router is not None 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, CartesianProductStreamSlicer):
+                    routers.extend(router.stream_slicers)
+                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:
@@ -535,14 +547,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

@@ -4183,6 +4183,7 @@ definitions:
           - "$ref": "#/definitions/SubstreamPartitionRouter"
           - "$ref": "#/definitions/ListPartitionRouter"
           - "$ref": "#/definitions/GroupingPartitionRouter"
+          - "$ref": "#/definitions/UnionPartitionRouter"
           - "$ref": "#/definitions/CustomPartitionRouter"
           - type: array
             title: Multiple Partition Routers
@@ -4191,6 +4192,7 @@ definitions:
                 - "$ref": "#/definitions/SubstreamPartitionRouter"
                 - "$ref": "#/definitions/ListPartitionRouter"
                 - "$ref": "#/definitions/GroupingPartitionRouter"
+                - "$ref": "#/definitions/UnionPartitionRouter"
                 - "$ref": "#/definitions/CustomPartitionRouter"
       $parameters:
         type: object
@@ -4393,13 +4395,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
@@ -4647,6 +4651,7 @@ definitions:
         anyOf:
           - "$ref": "#/definitions/ListPartitionRouter"
           - "$ref": "#/definitions/SubstreamPartitionRouter"
+          - "$ref": "#/definitions/UnionPartitionRouter"
           - "$ref": "#/definitions/CustomPartitionRouter"
       deduplicate:
         title: Deduplicate Partitions
@@ -4656,6 +4661,51 @@ 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. Partition values must be hashable; deduplication holds every distinct value
+      in memory for the duration of the sync, and the resulting partition count is the deduplicated sum
+      of the child routers' partition counts, which affects per-partition state size. Because child routers
+      must emit scalar partition values, GroupingPartitionRouter (which emits list-valued partitions)
+      cannot be used as a child.
+    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

@@ -3182,12 +3182,14 @@ class SimpleRetriever(BaseModel):
             SubstreamPartitionRouter,
             ListPartitionRouter,
             GroupingPartitionRouter,
+            UnionPartitionRouter,
             CustomPartitionRouter,
             List[
                 Union[
                     SubstreamPartitionRouter,
                     ListPartitionRouter,
                     GroupingPartitionRouter,
+                    UnionPartitionRouter,
                     CustomPartitionRouter,
                 ]
             ],
@@ -3261,12 +3263,14 @@ class AsyncRetriever(BaseModel):
             ListPartitionRouter,
             SubstreamPartitionRouter,
             GroupingPartitionRouter,
+            UnionPartitionRouter,
             CustomPartitionRouter,
             List[
                 Union[
                     ListPartitionRouter,
                     SubstreamPartitionRouter,
                     GroupingPartitionRouter,
+                    UnionPartitionRouter,
                     CustomPartitionRouter,
                 ]
             ],
@@ -3349,7 +3353,10 @@ class GroupingPartitionRouter(BaseModel):
         title="Group Size",
     )
     underlying_partition_router: Union[
-        ListPartitionRouter, SubstreamPartitionRouter, CustomPartitionRouter
+        ListPartitionRouter,
+        SubstreamPartitionRouter,
+        "UnionPartitionRouter",
+        CustomPartitionRouter,
     ] = Field(
         ...,
         description="The partition router whose output will be grouped. This can be any valid partition router component.",
@@ -3363,6 +3370,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(
@@ -3412,3 +3442,5 @@ class DynamicDeclarativeStream(BaseModel):
 PropertiesFromEndpoint.update_forward_refs()
 SimpleRetriever.update_forward_refs()
 AsyncRetriever.update_forward_refs()
+GroupingPartitionRouter.update_forward_refs()
+UnionPartitionRouter.update_forward_refs()