diff --git a/airbyte_cdk/__init__.py b/airbyte_cdk/__init__.py index 9cb83008e6..d971690c86 100644 --- a/airbyte_cdk/__init__.py +++ b/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", diff --git a/airbyte_cdk/sources/concurrent_source/concurrent_read_processor.py b/airbyte_cdk/sources/concurrent_source/concurrent_read_processor.py index fa6a568bc8..206562e96a 100644 --- a/airbyte_cdk/sources/concurrent_source/concurrent_read_processor.py +++ b/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 diff --git a/airbyte_cdk/sources/declarative/concurrent_declarative_source.py b/airbyte_cdk/sources/declarative/concurrent_declarative_source.py index fb9c0ea0bb..6ac7754aab 100644 --- a/airbyte_cdk/sources/declarative/concurrent_declarative_source.py +++ b/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: diff --git a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml index 8a46a9e9be..778ba3bed6 100644 --- a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml +++ b/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,6 +4214,7 @@ definitions: - "$ref": "#/definitions/ListPartitionRouter" - "$ref": "#/definitions/SubstreamPartitionRouter" - "$ref": "#/definitions/GroupingPartitionRouter" + - "$ref": "#/definitions/UnionPartitionRouter" - "$ref": "#/definitions/CustomPartitionRouter" - type: array items: @@ -4219,6 +4222,7 @@ definitions: - "$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. diff --git a/airbyte_cdk/sources/declarative/migrations/legacy_to_per_partition_state_migration.py b/airbyte_cdk/sources/declarative/migrations/legacy_to_per_partition_state_migration.py index 022881bf5b..73fcf2b72f 100644 --- a/airbyte_cdk/sources/declarative/migrations/legacy_to_per_partition_state_migration.py +++ b/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 "" : { diff --git a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py index d7fead317d..3e7ce98e05 100644 --- a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py +++ b/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() diff --git a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py index 5a75912b36..f78c966930 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -458,6 +458,9 @@ from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( TypesMap as TypesMapModel, ) +from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( + UnionPartitionRouter as UnionPartitionRouterModel, +) from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( UnlimitedCallRatePolicy as UnlimitedCallRatePolicyModel, ) @@ -484,6 +487,7 @@ PartitionRouter, SinglePartitionRouter, SubstreamPartitionRouter, + UnionPartitionRouter, ) from airbyte_cdk.sources.declarative.partition_routers.async_job_partition_router import ( AsyncJobPartitionRouter, @@ -827,6 +831,7 @@ def _init_mappings(self) -> None: RateModel: self.create_rate, HttpRequestRegexMatcherModel: self.create_http_request_matcher, GroupingPartitionRouterModel: self.create_grouping_partition_router, + UnionPartitionRouterModel: self.create_union_partition_router, } # Needed for the case where we need to perform a second parse on the fields of a custom component @@ -1136,12 +1141,19 @@ def create_legacy_to_per_partition_state_migration( ) partition_router = retriever.partition_router if not isinstance( - partition_router, (SubstreamPartitionRouterModel, CustomPartitionRouterModel) + partition_router, + ( + SubstreamPartitionRouterModel, + CustomPartitionRouterModel, + UnionPartitionRouterModel, + ), ): raise ValueError( - f"LegacyToPerPartitionStateMigrations can only be applied on a SimpleRetriever with a Substream partition router. Got {type(partition_router)}" + f"LegacyToPerPartitionStateMigrations can only be applied on a SimpleRetriever with a SubstreamPartitionRouter, UnionPartitionRouter or CustomPartitionRouter. Got {type(partition_router)}" ) - if not hasattr(partition_router, "parent_stream_configs"): + if not isinstance(partition_router, UnionPartitionRouterModel) and not hasattr( + partition_router, "parent_stream_configs" + ): raise ValueError( "LegacyToPerPartitionStateMigrations can only be applied with a parent stream configuration." ) @@ -4545,6 +4557,84 @@ def create_grouping_partition_router( config=config, ) + def create_union_partition_router( + self, + model: UnionPartitionRouterModel, + config: Config, + *, + stream_name: str, + **kwargs: Any, + ) -> UnionPartitionRouter: + partition_routers = [ + self._create_component_from_model( + model=child, + config=config, + stream_name=stream_name, + **kwargs, + ) + for child in model.partition_routers + ] + + # partition_field depends only on config/parameters, so it is evaluated once at build + # time; the runtime component always receives a plain string. + partition_field = InterpolatedString.create( + model.partition_field, parameters=model.parameters or {} + ).eval(config) + + # Fail fast at build time when a built-in child router is statically known to emit a + # partition field different from the union's. CustomPartitionRouter children are opaque + # and can only be validated at runtime. + for child_model in model.partition_routers: + child_partition_fields: List[str] = [] + if isinstance(child_model, ListPartitionRouterModel): + child_partition_fields.append( + InterpolatedString.create( + child_model.cursor_field, parameters=child_model.parameters or {} + ).eval(config) + ) + elif isinstance(child_model, SubstreamPartitionRouterModel): + for parent_stream_config in child_model.parent_stream_configs: + child_partition_fields.append( + InterpolatedString.create( + parent_stream_config.partition_field, + parameters=child_model.parameters or {}, + ).eval(config) + ) + elif isinstance(child_model, UnionPartitionRouterModel): + child_partition_fields.append( + InterpolatedString.create( + child_model.partition_field, parameters=child_model.parameters or {} + ).eval(config) + ) + for child_partition_field in child_partition_fields: + if child_partition_field != partition_field: + raise ValueError( + f"UnionPartitionRouter expects all child partition routers to emit the " + f"partition field '{partition_field}', but a " + f"{child_model.type} child emits '{child_partition_field}'." + ) + + # A union slice comes from exactly one child partition router, so request options + # declared on children cannot be applied consistently to requests built from the + # normalized union slices. Partition values should be consumed via interpolation + # (e.g. stream_partition) instead. Note that this validation only covers built-in + # router types; CustomPartitionRouter children are opaque, so any request options + # they implement internally cannot be detected or rejected here. + for router in partition_routers: + if isinstance(router, SubstreamPartitionRouter): + if any( + parent_config.request_option for parent_config in router.parent_stream_configs + ): + raise ValueError("Request options are not supported for UnionPartitionRouter.") + if isinstance(router, ListPartitionRouter) and router.request_option: + raise ValueError("Request options are not supported for UnionPartitionRouter.") + + return UnionPartitionRouter( + partition_routers=partition_routers, + partition_field=partition_field, + parameters=model.parameters or {}, + ) + def _ensure_query_properties_to_model( self, requester: Union[HttpRequesterModel, CustomRequesterModel] ) -> None: diff --git a/airbyte_cdk/sources/declarative/partition_routers/__init__.py b/airbyte_cdk/sources/declarative/partition_routers/__init__.py index 2e99286d25..8d3f31c4a4 100644 --- a/airbyte_cdk/sources/declarative/partition_routers/__init__.py +++ b/airbyte_cdk/sources/declarative/partition_routers/__init__.py @@ -21,6 +21,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, +) __all__ = [ "AsyncJobPartitionRouter", @@ -29,5 +32,6 @@ "ListPartitionRouter", "SinglePartitionRouter", "SubstreamPartitionRouter", + "UnionPartitionRouter", "PartitionRouter", ] diff --git a/airbyte_cdk/sources/declarative/partition_routers/union_partition_router.py b/airbyte_cdk/sources/declarative/partition_routers/union_partition_router.py new file mode 100644 index 0000000000..66769fceb9 --- /dev/null +++ b/airbyte_cdk/sources/declarative/partition_routers/union_partition_router.py @@ -0,0 +1,136 @@ +# +# Copyright (c) 2026 Airbyte, Inc., all rights reserved. +# + +import logging +from dataclasses import InitVar, dataclass +from typing import Any, Iterable, List, Mapping, Optional + +from airbyte_cdk.sources.declarative.partition_routers.partition_router import PartitionRouter +from airbyte_cdk.sources.types import StreamSlice, StreamState + +logger = logging.getLogger("airbyte") + + +@dataclass +class UnionPartitionRouter(PartitionRouter): + """ + A partition router that yields the deduplicated union of its child partition routers' slices. + + Each emitted slice's partition is normalized to exactly `{partition_field: value}`. Any other + partition keys coming from a child router (e.g. a SubstreamPartitionRouter's `parent_slice`) + are moved into `extra_fields` so that per-partition state keys stay stable across children. + + Given two child routers producing: + A: [{"repository": "org/a"}, {"repository": "org/b"}] + B: [{"repository": "org/b", "parent_slice": {...}}, {"repository": "org/c", "parent_slice": {...}}] + the union yields: + [{"repository": "org/a"}, {"repository": "org/b"}, {"repository": "org/c"}] + + Attributes: + partition_routers (List[PartitionRouter]): The child partition routers to union. + partition_field (str): The single partition key all child slices are normalized to. + """ + + partition_routers: List[PartitionRouter] + partition_field: str + parameters: InitVar[Mapping[str, Any]] + + def __post_init__(self, parameters: Mapping[str, Any]) -> None: + self._parameters = parameters + + def stream_slices(self) -> Iterable[StreamSlice]: + """ + Iterate over the child partition routers in order, yielding each partition value once. + + The first occurrence of a partition value wins; later duplicates from any child are skipped. + """ + seen: set[Any] = set() + for router in self.partition_routers: + for stream_slice in router.stream_slices(): + if self.partition_field not in stream_slice.partition: + raise ValueError( + f"UnionPartitionRouter expects all child partition routers to emit the " + f"partition field '{self.partition_field}'. Got {stream_slice.partition}" + ) + value = stream_slice.partition[self.partition_field] + try: + if value in seen: + continue + seen.add(value) + except TypeError as exception: + raise ValueError( + f"UnionPartitionRouter can only deduplicate hashable partition values. " + f"Got unhashable value {value!r} for partition field '{self.partition_field}'." + ) from exception + carried_over_fields = { + key: field_value + for key, field_value in stream_slice.partition.items() + if key != self.partition_field + } + yield StreamSlice( + partition={self.partition_field: value}, + cursor_slice={}, + extra_fields={ + **(dict(stream_slice.extra_fields) if stream_slice.extra_fields else {}), + **carried_over_fields, + }, + ) + + def get_request_params( + self, + *, + stream_state: Optional[StreamState] = None, + stream_slice: Optional[StreamSlice] = None, + next_page_token: Optional[Mapping[str, Any]] = None, + ) -> Mapping[str, Any]: + return {} + + def get_request_headers( + self, + *, + stream_state: Optional[StreamState] = None, + stream_slice: Optional[StreamSlice] = None, + next_page_token: Optional[Mapping[str, Any]] = None, + ) -> Mapping[str, Any]: + return {} + + def get_request_body_data( + self, + *, + stream_state: Optional[StreamState] = None, + stream_slice: Optional[StreamSlice] = None, + next_page_token: Optional[Mapping[str, Any]] = None, + ) -> Mapping[str, Any]: + return {} + + def get_request_body_json( + self, + *, + stream_state: Optional[StreamState] = None, + stream_slice: Optional[StreamSlice] = None, + next_page_token: Optional[Mapping[str, Any]] = None, + ) -> Mapping[str, Any]: + return {} + + def get_stream_state(self) -> Optional[Mapping[str, StreamState]]: + """ + Merge the parent stream states of all child partition routers. + + States are merged with a shallow dict update, which assumes no two child routers + reference a parent stream with the same name. If they do, the last child's state + for that parent wins and a warning is logged. + """ + merged_state: dict[str, StreamState] = {} + for router in self.partition_routers: + child_state = router.get_stream_state() + if child_state: + colliding_keys = merged_state.keys() & child_state.keys() + if colliding_keys: + logger.warning( + f"UnionPartitionRouter child routers reference the same parent stream(s) " + f"{sorted(colliding_keys)}; the last child's state for each colliding parent " + f"stream overwrites the previous one." + ) + merged_state.update(child_state) + return merged_state diff --git a/unit_tests/sources/declarative/migrations/test_legacy_to_per_partition_migration.py b/unit_tests/sources/declarative/migrations/test_legacy_to_per_partition_migration.py index 11be8231a1..708a1e18cb 100644 --- a/unit_tests/sources/declarative/migrations/test_legacy_to_per_partition_migration.py +++ b/unit_tests/sources/declarative/migrations/test_legacy_to_per_partition_migration.py @@ -14,9 +14,11 @@ CustomRetriever, DatetimeBasedCursor, DeclarativeStream, + ListPartitionRouter, ParentStreamConfig, SimpleRetriever, SubstreamPartitionRouter, + UnionPartitionRouter, ) from airbyte_cdk.sources.declarative.models import ( LegacyToPerPartitionStateMigration as LegacyToPerPartitionStateMigrationModel, @@ -276,6 +278,71 @@ def _migrator_with_multiple_parent_streams(): return LegacyToPerPartitionStateMigration(partition_router, cursor, config, parameters) +def test_migrate_a_valid_legacy_state_to_per_partition_with_union_partition_router(): + input_state = { + "org/a": {"last_changed": "2022-12-27T08:34:39+00:00"}, + "org/b": {"last_changed": "2022-12-27T08:35:39+00:00"}, + } + + migrator = _union_migrator() + + assert migrator.should_migrate(input_state) + + expected_state = { + "states": [ + { + "partition": {"repository": "org/a"}, + "cursor": {"last_changed": "2022-12-27T08:34:39+00:00"}, + }, + { + "partition": {"repository": "org/b"}, + "cursor": {"last_changed": "2022-12-27T08:35:39+00:00"}, + }, + ] + } + + assert migrator.migrate(input_state) == expected_state + + +def _union_migrator(): + partition_router = UnionPartitionRouter( + type="UnionPartitionRouter", + partition_field="repository", + partition_routers=[ + SubstreamPartitionRouter( + type="SubstreamPartitionRouter", + parent_stream_configs=[ + ParentStreamConfig( + type="ParentStreamConfig", + parent_key="full_name", + partition_field="repository", + stream=DeclarativeStream( + type="DeclarativeStream", + retriever=CustomRetriever( + type="CustomRetriever", class_name="a_class_name" + ), + ), + ) + ], + ), + ListPartitionRouter( + type="ListPartitionRouter", + cursor_field="repository", + values=["org/a", "org/b"], + ), + ], + ) + cursor = DatetimeBasedCursor( + type="DatetimeBasedCursor", + cursor_field="{{ parameters['cursor_field'] }}", + datetime_format="%Y-%m-%dT%H:%M:%S.%fZ", + start_datetime="1970-01-01T00:00:00.0Z", + ) + config = {} + parameters = {"cursor_field": "last_changed"} + return LegacyToPerPartitionStateMigration(partition_router, cursor, config, parameters) + + @pytest.mark.parametrize( "retriever_type, partition_router_class, is_parent_stream_config, expected_exception, expected_error_message", [ @@ -292,7 +359,7 @@ def _migrator_with_multiple_parent_streams(): None, False, ValueError, - "LegacyToPerPartitionStateMigrations can only be applied on a SimpleRetriever with a Substream partition router. Got ", + "LegacyToPerPartitionStateMigrations can only be applied on a SimpleRetriever with a SubstreamPartitionRouter, UnionPartitionRouter or CustomPartitionRouter. Got ", ), ( SimpleRetriever, diff --git a/unit_tests/sources/declarative/parsers/test_model_to_component_factory.py b/unit_tests/sources/declarative/parsers/test_model_to_component_factory.py index 5149f46d2b..3bf388da24 100644 --- a/unit_tests/sources/declarative/parsers/test_model_to_component_factory.py +++ b/unit_tests/sources/declarative/parsers/test_model_to_component_factory.py @@ -89,6 +89,9 @@ from airbyte_cdk.sources.declarative.models import ( SubstreamPartitionRouter as SubstreamPartitionRouterModel, ) +from airbyte_cdk.sources.declarative.models import ( + UnionPartitionRouter as UnionPartitionRouterModel, +) from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( ConstantBackoffStrategy as ConstantBackoffStrategyModel, ) @@ -123,6 +126,7 @@ ListPartitionRouter, SinglePartitionRouter, SubstreamPartitionRouter, + UnionPartitionRouter, ) from airbyte_cdk.sources.declarative.requesters import HttpRequester from airbyte_cdk.sources.declarative.requesters.error_handlers import ( @@ -4678,6 +4682,240 @@ def test_create_grouping_partition_router_substream_with_request_option(): ) +def test_create_union_partition_router(): + content = """ + schema_loader: + file_path: "./source_example/schemas/{{ parameters['name'] }}.yaml" + name: "{{ parameters['stream_name'] }}" + retriever: + requester: + type: "HttpRequester" + path: "example" + record_selector: + extractor: + field_path: [] + stream_A: + type: DeclarativeStream + name: "A" + primary_key: "id" + $parameters: + retriever: "#/retriever" + url_base: "https://airbyte.io" + schema_loader: "#/schema_loader" + partition_router: + type: UnionPartitionRouter + partition_field: repository + partition_routers: + - type: SubstreamPartitionRouter + parent_stream_configs: + - stream: "#/stream_A" + parent_key: full_name + partition_field: repository + - type: ListPartitionRouter + cursor_field: repository + values: ["org/a", "org/b"] + """ + parsed_manifest = YamlDeclarativeSource._parse(content) + resolved_manifest = resolver.preprocess_manifest(parsed_manifest) + partition_router_manifest = transformer.propagate_types_and_parameters( + "", resolved_manifest["partition_router"], {} + ) + + partition_router = factory.create_component( + model_type=UnionPartitionRouterModel, + component_definition=partition_router_manifest, + config=input_config, + stream_name="child_stream", + ) + + assert isinstance(partition_router, UnionPartitionRouter) + assert partition_router.partition_field == "repository" + assert len(partition_router.partition_routers) == 2 + assert isinstance(partition_router.partition_routers[0], SubstreamPartitionRouter) + assert isinstance(partition_router.partition_routers[1], ListPartitionRouter) + + parent_stream_configs = partition_router.partition_routers[0].parent_stream_configs + assert len(parent_stream_configs) == 1 + assert parent_stream_configs[0].parent_key.eval({}) == "full_name" + assert parent_stream_configs[0].partition_field.eval({}) == "repository" + + +def test_create_union_partition_router_with_interpolated_partition_field(): + content = """ + partition_router: + type: UnionPartitionRouter + partition_field: "{{ config['union_partition_field'] }}" + partition_routers: + - type: ListPartitionRouter + cursor_field: "{{ config['union_partition_field'] }}" + values: ["org/a"] + - type: ListPartitionRouter + cursor_field: repository + values: ["org/b"] + """ + parsed_manifest = YamlDeclarativeSource._parse(content) + resolved_manifest = resolver.preprocess_manifest(parsed_manifest) + partition_router_manifest = transformer.propagate_types_and_parameters( + "", resolved_manifest["partition_router"], {} + ) + + partition_router = factory.create_component( + model_type=UnionPartitionRouterModel, + component_definition=partition_router_manifest, + config={**input_config, "union_partition_field": "repository"}, + stream_name="child_stream", + ) + + assert isinstance(partition_router, UnionPartitionRouter) + assert partition_router.partition_field == "repository" + + +@pytest.mark.parametrize( + "child_router_manifest", + [ + pytest.param( + """ + - type: SubstreamPartitionRouter + parent_stream_configs: + - stream: "#/stream_A" + parent_key: full_name + partition_field: repository + request_option: + inject_into: request_parameter + field_name: "repo" +""", + id="substream_child_with_request_option", + ), + pytest.param( + """ + - type: ListPartitionRouter + cursor_field: repository + values: ["org/a"] + request_option: + inject_into: request_parameter + field_name: "repo" +""", + id="list_child_with_request_option", + ), + ], +) +def test_create_union_partition_router_with_request_option(child_router_manifest): + content = f""" + schema_loader: + file_path: "./source_example/schemas/{{{{ parameters['name'] }}}}.yaml" + name: "{{{{ parameters['stream_name'] }}}}" + retriever: + requester: + type: "HttpRequester" + path: "example" + record_selector: + extractor: + field_path: [] + stream_A: + type: DeclarativeStream + name: "A" + primary_key: "id" + $parameters: + retriever: "#/retriever" + url_base: "https://airbyte.io" + schema_loader: "#/schema_loader" + partition_router: + type: UnionPartitionRouter + partition_field: repository + partition_routers: +{child_router_manifest} + - type: ListPartitionRouter + cursor_field: repository + values: ["org/b"] + """ + parsed_manifest = YamlDeclarativeSource._parse(content) + resolved_manifest = resolver.preprocess_manifest(parsed_manifest) + partition_router_manifest = transformer.propagate_types_and_parameters( + "", resolved_manifest["partition_router"], {} + ) + + with pytest.raises( + ValueError, match="Request options are not supported for UnionPartitionRouter." + ): + factory.create_component( + model_type=UnionPartitionRouterModel, + component_definition=partition_router_manifest, + config=input_config, + stream_name="child_stream", + ) + + +@pytest.mark.parametrize( + "child_router_manifest, mismatched_field", + [ + pytest.param( + """ + - type: SubstreamPartitionRouter + parent_stream_configs: + - stream: "#/stream_A" + parent_key: full_name + partition_field: repo +""", + "repo", + id="substream_child_with_mismatched_partition_field", + ), + pytest.param( + """ + - type: ListPartitionRouter + cursor_field: repo + values: ["org/a"] +""", + "repo", + id="list_child_with_mismatched_cursor_field", + ), + ], +) +def test_create_union_partition_router_with_mismatched_partition_field( + child_router_manifest, mismatched_field +): + content = f""" + schema_loader: + file_path: "./source_example/schemas/{{{{ parameters['name'] }}}}.yaml" + name: "{{{{ parameters['stream_name'] }}}}" + retriever: + requester: + type: "HttpRequester" + path: "example" + record_selector: + extractor: + field_path: [] + stream_A: + type: DeclarativeStream + name: "A" + primary_key: "id" + $parameters: + retriever: "#/retriever" + url_base: "https://airbyte.io" + schema_loader: "#/schema_loader" + partition_router: + type: UnionPartitionRouter + partition_field: repository + partition_routers: +{child_router_manifest} + - type: ListPartitionRouter + cursor_field: repository + values: ["org/b"] + """ + parsed_manifest = YamlDeclarativeSource._parse(content) + resolved_manifest = resolver.preprocess_manifest(parsed_manifest) + partition_router_manifest = transformer.propagate_types_and_parameters( + "", resolved_manifest["partition_router"], {} + ) + + with pytest.raises(ValueError, match=f"emits '{mismatched_field}'"): + factory.create_component( + model_type=UnionPartitionRouterModel, + component_definition=partition_router_manifest, + config=input_config, + stream_name="child_stream", + ) + + def test_simple_retriever_with_query_properties(): content = """ selector: diff --git a/unit_tests/sources/declarative/partition_routers/test_union_partition_router.py b/unit_tests/sources/declarative/partition_routers/test_union_partition_router.py new file mode 100644 index 0000000000..938f2609e6 --- /dev/null +++ b/unit_tests/sources/declarative/partition_routers/test_union_partition_router.py @@ -0,0 +1,231 @@ +# +# Copyright (c) 2026 Airbyte, Inc., all rights reserved. +# + +import pytest + +from airbyte_cdk.sources.declarative.partition_routers import ( + ListPartitionRouter, + UnionPartitionRouter, +) +from airbyte_cdk.sources.declarative.partition_routers.partition_router import PartitionRouter +from airbyte_cdk.sources.types import StreamSlice + + +class _StaticPartitionRouter(PartitionRouter): + """A test partition router that emits a fixed list of slices and exposes a fixed state.""" + + def __init__(self, slices, state=None): + self._slices = slices + self._state = state + + def stream_slices(self): + yield from self._slices + + def get_request_params(self, *, stream_state=None, stream_slice=None, next_page_token=None): + return {} + + def get_request_headers(self, *, stream_state=None, stream_slice=None, next_page_token=None): + return {} + + def get_request_body_data(self, *, stream_state=None, stream_slice=None, next_page_token=None): + return {} + + def get_request_body_json(self, *, stream_state=None, stream_slice=None, next_page_token=None): + return {} + + def get_stream_state(self): + return self._state + + +def _list_router(values, cursor_field="repository"): + return ListPartitionRouter(values=values, cursor_field=cursor_field, config={}, parameters={}) + + +def test_stream_slices_deduplicates_across_children(): + router = UnionPartitionRouter( + partition_routers=[ + _list_router(["org/a", "org/b"]), + _list_router(["org/b", "org/c"]), + ], + partition_field="repository", + parameters={}, + ) + + slices = list(router.stream_slices()) + + assert slices == [ + StreamSlice(partition={"repository": "org/a"}, cursor_slice={}), + StreamSlice(partition={"repository": "org/b"}, cursor_slice={}), + StreamSlice(partition={"repository": "org/c"}, cursor_slice={}), + ] + + +def test_stream_slices_first_occurrence_wins(): + child_a = _StaticPartitionRouter( + [ + StreamSlice( + partition={"repository": "org/b"}, + cursor_slice={}, + extra_fields={"origin": "a"}, + ) + ] + ) + child_b = _StaticPartitionRouter( + [ + StreamSlice( + partition={"repository": "org/b"}, + cursor_slice={}, + extra_fields={"origin": "b"}, + ) + ] + ) + + slices = list( + UnionPartitionRouter( + partition_routers=[child_a, child_b], + partition_field="repository", + parameters={}, + ).stream_slices() + ) + + assert len(slices) == 1 + assert slices[0].extra_fields["origin"] == "a" + + +def test_stream_slices_normalizes_partition_and_moves_extra_keys_to_extra_fields(): + substream_like_child = _StaticPartitionRouter( + [ + StreamSlice( + partition={"repository": "org/a", "parent_slice": {"organization": "org"}}, + cursor_slice={}, + extra_fields={"full_name": "org/a"}, + ) + ] + ) + + slices = list( + UnionPartitionRouter( + partition_routers=[substream_like_child], + partition_field="repository", + parameters={}, + ).stream_slices() + ) + + assert slices[0].partition == {"repository": "org/a"} + assert slices[0].extra_fields == { + "full_name": "org/a", + "parent_slice": {"organization": "org"}, + } + + +def test_stream_slices_raises_when_child_does_not_emit_partition_field(): + child = _StaticPartitionRouter( + [StreamSlice(partition={"other_field": "value"}, cursor_slice={})] + ) + router = UnionPartitionRouter( + partition_routers=[child], partition_field="repository", parameters={} + ) + + with pytest.raises(ValueError, match="partition field 'repository'"): + list(router.stream_slices()) + + +def test_nested_union_partition_router(): + inner = UnionPartitionRouter( + partition_routers=[ + _list_router(["org/a"]), + _list_router(["org/b"]), + ], + partition_field="repository", + parameters={}, + ) + outer = UnionPartitionRouter( + partition_routers=[inner, _list_router(["org/b", "org/c"])], + partition_field="repository", + parameters={}, + ) + + assert [s.partition for s in outer.stream_slices()] == [ + {"repository": "org/a"}, + {"repository": "org/b"}, + {"repository": "org/c"}, + ] + + +@pytest.mark.parametrize( + "child_states, expected_state", + [ + pytest.param([None, None], {}, id="all_children_without_state"), + pytest.param( + [{"parent_a": {"updated_at": "2024-01-01"}}, None], + {"parent_a": {"updated_at": "2024-01-01"}}, + id="one_child_with_state", + ), + pytest.param( + [ + {"parent_a": {"updated_at": "2024-01-01"}}, + {"parent_b": {"updated_at": "2024-02-01"}}, + ], + { + "parent_a": {"updated_at": "2024-01-01"}, + "parent_b": {"updated_at": "2024-02-01"}, + }, + id="merges_states_of_all_children", + ), + ], +) +def test_get_stream_state_merges_children_states(child_states, expected_state): + router = UnionPartitionRouter( + partition_routers=[_StaticPartitionRouter([], state=state) for state in child_states], + partition_field="repository", + parameters={}, + ) + + assert router.get_stream_state() == expected_state + + +def test_unhashable_partition_value_raises_value_error(): + router = UnionPartitionRouter( + partition_routers=[ + _StaticPartitionRouter( + [StreamSlice(partition={"repository": ["org/a"]}, cursor_slice={})] + ) + ], + partition_field="repository", + parameters={}, + ) + + with pytest.raises(ValueError, match="hashable"): + list(router.stream_slices()) + + +def test_get_stream_state_warns_on_parent_stream_name_collision(caplog): + router = UnionPartitionRouter( + partition_routers=[ + _StaticPartitionRouter([], state={"parent_a": {"updated_at": "2024-01-01"}}), + _StaticPartitionRouter([], state={"parent_a": {"updated_at": "2024-02-01"}}), + ], + partition_field="repository", + parameters={}, + ) + + with caplog.at_level("WARNING", logger="airbyte"): + state = router.get_stream_state() + + assert state == {"parent_a": {"updated_at": "2024-02-01"}} + assert any("parent_a" in record.message for record in caplog.records) + + +def test_request_options_are_empty(): + router = UnionPartitionRouter( + partition_routers=[_list_router(["org/a"])], + partition_field="repository", + parameters={}, + ) + stream_slice = StreamSlice(partition={"repository": "org/a"}, cursor_slice={}) + + assert router.get_request_params(stream_slice=stream_slice) == {} + assert router.get_request_headers(stream_slice=stream_slice) == {} + assert router.get_request_body_data(stream_slice=stream_slice) == {} + assert router.get_request_body_json(stream_slice=stream_slice) == {} diff --git a/unit_tests/sources/declarative/test_concurrent_declarative_source.py b/unit_tests/sources/declarative/test_concurrent_declarative_source.py index 5ffda84156..77adf8983e 100644 --- a/unit_tests/sources/declarative/test_concurrent_declarative_source.py +++ b/unit_tests/sources/declarative/test_concurrent_declarative_source.py @@ -1408,6 +1408,185 @@ def test_concurrent_declarative_source_runs_state_migrations_provided_in_manifes ], "State was migrated, but actual state don't match expected" +@freezegun.freeze_time(_NOW) +def test_read_resumes_from_legacy_state_with_union_partition_router(): + """ + Round-trip test: a legacy (pre-per-partition) state is migrated through + LegacyToPerPartitionStateMigration for a stream partitioned by a UnionPartitionRouter, + and the runtime per-partition state keys produced during the read match the migrated keys. + """ + + def _parent_stream(name: str) -> dict: + return { + "type": "DeclarativeStream", + "name": name, + "primary_key": "full_name", + "retriever": { + "type": "SimpleRetriever", + "requester": { + "type": "HttpRequester", + "url_base": "https://api.example.com", + "path": f"/{name}", + "http_method": "GET", + # Explicitly disabled to avoid SQLite-backed request caching in tests. + "use_cache": False, + }, + "record_selector": { + "type": "RecordSelector", + "extractor": {"type": "DpathExtractor", "field_path": []}, + }, + }, + "schema_loader": { + "type": "InlineSchemaLoader", + "schema": {"type": "object", "properties": {}}, + }, + } + + def _substream_router(parent_stream: dict) -> dict: + return { + "type": "SubstreamPartitionRouter", + "parent_stream_configs": [ + { + "type": "ParentStreamConfig", + "parent_key": "full_name", + "partition_field": "repository", + "stream": parent_stream, + } + ], + } + + manifest = { + "version": "5.0.0", + "definitions": {}, + "streams": [ + _parent_stream("repositories"), + _parent_stream("starred"), + { + "type": "DeclarativeStream", + "name": "issues", + "primary_key": "id", + "retriever": { + "type": "SimpleRetriever", + "requester": { + "type": "HttpRequester", + "url_base": "https://api.example.com", + "path": "/issues/{{ stream_partition['repository'] }}", + "http_method": "GET", + }, + "record_selector": { + "type": "RecordSelector", + "extractor": {"type": "DpathExtractor", "field_path": []}, + }, + "partition_router": { + "type": "UnionPartitionRouter", + "partition_field": "repository", + "partition_routers": [ + _substream_router(_parent_stream("repositories")), + _substream_router(_parent_stream("starred")), + ], + }, + }, + "incremental_sync": { + "type": "DatetimeBasedCursor", + "start_datetime": { + "datetime": "{{ format_datetime(config['start_date'], '%Y-%m-%d') }}" + }, + "end_datetime": {"datetime": "{{ now_utc().strftime('%Y-%m-%d') }}"}, + "datetime_format": "%Y-%m-%d", + "cursor_datetime_formats": ["%Y-%m-%d"], + "cursor_field": "updated_at", + }, + "state_migrations": [{"type": "LegacyToPerPartitionStateMigration"}], + "schema_loader": { + "type": "InlineSchemaLoader", + "schema": {"type": "object", "properties": {}}, + }, + }, + ], + "check": {"type": "CheckStream", "stream_names": ["repositories"]}, + } + + # Legacy (pre-per-partition) state format: {partition_value: {cursor_field: cursor_value}} + legacy_state = [ + AirbyteStateMessage( + type=AirbyteStateType.STREAM, + stream=AirbyteStreamState( + stream_descriptor=StreamDescriptor(name="issues", namespace=None), + stream_state=AirbyteStateBlob( + **{ + "org/repo-a": {"updated_at": "2024-08-21"}, + "org/repo-b": {"updated_at": "2024-08-22"}, + } + ), + ), + ), + ] + + catalog = ConfiguredAirbyteCatalog( + streams=[ + ConfiguredAirbyteStream( + stream=AirbyteStream( + name="issues", json_schema={}, supported_sync_modes=[SyncMode.incremental] + ), + sync_mode=SyncMode.incremental, + destination_sync_mode=DestinationSyncMode.append, + ), + ] + ) + + source = ConcurrentDeclarativeSource( + source_config=manifest, config=_CONFIG, catalog=catalog, state=legacy_state + ) + + # The migrated state keys are exactly `{partition_field: partition_value}`. + migrated_states = source.streams(_CONFIG)[2].cursor.state.get("states") + assert {json.dumps(state["partition"], sort_keys=True) for state in migrated_states} == { + '{"repository": "org/repo-a"}', + '{"repository": "org/repo-b"}', + } + + with HttpMocker() as http_mocker: + http_mocker.get( + HttpRequest("https://api.example.com/repositories"), + HttpResponse( + json.dumps([{"full_name": "org/repo-a"}, {"full_name": "org/repo-b"}]), 200 + ), + ) + http_mocker.get( + HttpRequest("https://api.example.com/starred"), + HttpResponse( + json.dumps([{"full_name": "org/repo-b"}, {"full_name": "org/repo-c"}]), 200 + ), + ) + for repository in ("org/repo-a", "org/repo-b", "org/repo-c"): + http_mocker.get( + HttpRequest(f"https://api.example.com/issues/{repository}"), + HttpResponse( + json.dumps([{"id": f"{repository}-1", "updated_at": "2024-09-01"}]), 200 + ), + ) + + messages = list( + source.read(logger=source.logger, config=_CONFIG, catalog=catalog, state=legacy_state) + ) + + # Deduplicated union: org/repo-b appears in both parents but is only read once. + issues_records = get_records_for_stream("issues", messages) + assert len(issues_records) == 3 + + # The runtime per-partition state keys match the migrated legacy keys exactly. + final_state = get_states_for_stream(stream_name="issues", messages=messages)[-1] + runtime_partitions = { + json.dumps(state["partition"], sort_keys=True) + for state in final_state.stream.stream_state.__dict__["states"] + } + assert runtime_partitions == { + '{"repository": "org/repo-a"}', + '{"repository": "org/repo-b"}', + '{"repository": "org/repo-c"}', + } + + @freezegun.freeze_time(_NOW) @patch( "airbyte_cdk.sources.streams.concurrent.state_converters.abstract_stream_state_converter.AbstractStreamStateConverter.__init__", @@ -5481,6 +5660,236 @@ def test_apply_stream_groups_raises_on_parent_child_in_same_group_with_grouping_ ConcurrentDeclarativeSource._apply_stream_groups(source, [parent, child]) +def _make_child_stream_with_union_router( + child_name: str, + parent_streams: list[DefaultStream], + wrap_in_grouping: bool = False, +) -> DefaultStream: + """Create a DefaultStream with a UnionPartitionRouter over SubstreamPartitionRouters.""" + from airbyte_cdk.sources.declarative.incremental.concurrent_partition_cursor import ( + ConcurrentCursorFactory, + ConcurrentPerPartitionCursor, + ) + from airbyte_cdk.sources.declarative.partition_routers.grouping_partition_router import ( + GroupingPartitionRouter, + ) + from airbyte_cdk.sources.declarative.partition_routers.substream_partition_router import ( + ParentStreamConfig, + SubstreamPartitionRouter, + ) + from airbyte_cdk.sources.declarative.partition_routers.union_partition_router import ( + UnionPartitionRouter, + ) + from airbyte_cdk.sources.declarative.stream_slicers.declarative_partition_generator import ( + DeclarativePartitionFactory, + StreamSlicerPartitionGenerator, + ) + from airbyte_cdk.sources.streams.concurrent.cursor import FinalStateCursor + from airbyte_cdk.sources.streams.concurrent.state_converters.datetime_stream_state_converter import ( + EpochValueConcurrentStreamStateConverter, + ) + + substream_routers = [ + SubstreamPartitionRouter( + parent_stream_configs=[ + ParentStreamConfig( + stream=parent_stream, + parent_key="id", + partition_field="parent_id", + config={}, + parameters={}, + ) + ], + config={}, + parameters={}, + ) + for parent_stream in parent_streams + ] + + union_router = UnionPartitionRouter( + partition_routers=substream_routers, + partition_field="parent_id", + parameters={}, + ) + + stream_slicer_router = ( + GroupingPartitionRouter( + group_size=10, + underlying_partition_router=union_router, + config={}, + ) + if wrap_in_grouping + else union_router + ) + + cursor_factory = ConcurrentCursorFactory(lambda *args, **kwargs: Mock()) + message_repository = InMemoryMessageRepository() + state_converter = EpochValueConcurrentStreamStateConverter() + + per_partition_cursor = ConcurrentPerPartitionCursor( + cursor_factory=cursor_factory, + partition_router=stream_slicer_router, + stream_name=child_name, + stream_namespace=None, + stream_state={}, + message_repository=message_repository, + connector_state_manager=Mock(), + connector_state_converter=state_converter, + cursor_field=Mock(cursor_field_key="updated_at"), + ) + + partition_factory = Mock(spec=DeclarativePartitionFactory) + partition_generator = StreamSlicerPartitionGenerator( + partition_factory=partition_factory, + stream_slicer=per_partition_cursor, + ) + + cursor = FinalStateCursor( + stream_name=child_name, stream_namespace=None, message_repository=message_repository + ) + return DefaultStream( + partition_generator=partition_generator, + name=child_name, + json_schema={}, + primary_key=[], + cursor_field=None, + logger=logging.getLogger(f"test.{child_name}"), + cursor=cursor, + ) + + +@pytest.mark.parametrize( + "grouped_parent,wrap_in_grouping", + [ + pytest.param("parent_a", False, id="first_union_child_parent"), + pytest.param("parent_b", False, id="second_union_child_parent"), + pytest.param("parent_a", True, id="union_nested_in_grouping"), + ], +) +def test_apply_stream_groups_raises_on_parent_child_in_same_group_with_union_router( + grouped_parent, wrap_in_grouping +): + """Test _apply_stream_groups detects deadlock through a UnionPartitionRouter's children.""" + parent_a = _make_default_stream("parent_a") + parent_b = _make_default_stream("parent_b") + child = _make_child_stream_with_union_router( + "child_stream", [parent_a, parent_b], wrap_in_grouping=wrap_in_grouping + ) + + source = Mock() + source._source_config = { + "stream_groups": { + "my_group": { + "streams": [ + {"name": grouped_parent, "type": "DeclarativeStream"}, + {"name": "child_stream", "type": "DeclarativeStream"}, + ], + "action": {"type": "BlockSimultaneousSyncsAction"}, + } + } + } + + with pytest.raises(ValueError, match="child stream must not share a group with its parent"): + ConcurrentDeclarativeSource._apply_stream_groups(source, [parent_a, parent_b, child]) + + +def test_union_partition_router_parent_streams_use_cache(): + """Parents referenced through a UnionPartitionRouter get use_cache force-enabled.""" + + def _stream_config(name: str) -> dict: + return { + "type": "DeclarativeStream", + "$parameters": { + "name": name, + "primary_key": "id", + "url_base": "https://api.example.com/v1/", + }, + "schema_loader": { + "type": "InlineSchemaLoader", + "schema": {"type": "object", "properties": {}}, + }, + "retriever": { + "type": "SimpleRetriever", + "requester": { + "type": "HttpRequester", + "path": name, + }, + "record_selector": {"extractor": {"type": "DpathExtractor", "field_path": []}}, + }, + } + + child_stream = _stream_config("repository_stats") + child_stream["retriever"]["partition_router"] = { + "type": "UnionPartitionRouter", + "partition_field": "repository", + "partition_routers": [ + { + "type": "SubstreamPartitionRouter", + "parent_stream_configs": [ + { + "type": "ParentStreamConfig", + "parent_key": "full_name", + "partition_field": "repository", + "stream": _stream_config("repositories"), + } + ], + }, + { + "type": "SubstreamPartitionRouter", + "parent_stream_configs": [ + { + "type": "ParentStreamConfig", + "parent_key": "full_name", + "partition_field": "repository", + "stream": _stream_config("starred_repositories"), + } + ], + }, + ], + } + + manifest = { + "version": "0.29.3", + "definitions": {}, + "streams": [ + _stream_config("repositories"), + _stream_config("starred_repositories"), + child_stream, + ], + "check": {"type": "CheckStream", "stream_names": ["repositories"]}, + } + + source = ConcurrentDeclarativeSource( + source_config=manifest, config={}, catalog=create_catalog("repositories"), state=None + ) + + streams = source.streams({}) + streams_by_name = {stream.name: stream for stream in streams} + assert set(streams_by_name) == {"repositories", "starred_repositories", "repository_stats"} + + def _use_cache(stream) -> bool: + return stream._stream_partition_generator._partition_factory._retriever.requester.use_cache + + # Both parents referenced through the union get caching enabled; the child does not. + assert _use_cache(streams_by_name["repositories"]) + assert _use_cache(streams_by_name["starred_repositories"]) + assert not _use_cache(streams_by_name["repository_stats"]) + + # The parent stream instances nested inside the union's substream routers are also cached. + union_router = streams_by_name["repository_stats"]._stream_partition_generator._stream_slicer + nested_parents = [ + parent_config.stream + for child_router in union_router.partition_routers + for parent_config in child_router.parent_stream_configs + ] + assert {parent.name for parent in nested_parents} == { + "repositories", + "starred_repositories", + } + for parent in nested_parents: + assert _use_cache(parent) + + @pytest.mark.parametrize( "stream_factory,expected_type", [ diff --git a/unit_tests/sources/streams/concurrent/test_concurrent_read_processor.py b/unit_tests/sources/streams/concurrent/test_concurrent_read_processor.py index 856bd11564..091dd3a33a 100644 --- a/unit_tests/sources/streams/concurrent/test_concurrent_read_processor.py +++ b/unit_tests/sources/streams/concurrent/test_concurrent_read_processor.py @@ -35,6 +35,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 LogMessage, MessageRepository from airbyte_cdk.sources.streams.concurrent.abstract_stream import AbstractStream from airbyte_cdk.sources.streams.concurrent.default_stream import DefaultStream @@ -1645,3 +1648,64 @@ def test_collect_parent_stream_names_unwraps_grouping_partition_router(): parent_names = handler._collect_all_parent_stream_names("child") assert parent_names == {"parent"} + + +@pytest.mark.parametrize( + "wrap_in_grouping", + [ + pytest.param(False, id="union_router"), + pytest.param(True, id="union_nested_in_grouping_router"), + ], +) +def test_collect_parent_stream_names_unwraps_union_partition_router(wrap_in_grouping): + """Test _collect_all_parent_stream_names collects parents from all UnionPartitionRouter children.""" + partition_enqueuer = Mock(spec=PartitionEnqueuer) + thread_pool_manager = Mock(spec=ThreadPoolManager) + logger = Mock(spec=logging.Logger) + slice_logger = Mock(spec=SliceLogger) + message_repository = Mock(spec=MessageRepository) + message_repository.consume_queue.return_value = [] + partition_reader = Mock(spec=PartitionReader) + + parent_a = Mock(spec=AbstractStream) + parent_a.name = "parent_a" + parent_a.block_simultaneous_read = "" + + parent_b = Mock(spec=AbstractStream) + parent_b.name = "parent_b" + parent_b.block_simultaneous_read = "" + + child_stream = Mock(spec=DefaultStream) + child_stream.name = "child" + child_stream.block_simultaneous_read = "" + + substream_routers = [] + for parent in (parent_a, parent_b): + substream_router = Mock(spec=SubstreamPartitionRouter) + parent_config = Mock() + parent_config.stream = parent + substream_router.parent_stream_configs = [parent_config] + substream_routers.append(substream_router) + + union_router = Mock(spec=UnionPartitionRouter) + union_router.partition_routers = substream_routers + + if wrap_in_grouping: + grouping_router = Mock(spec=GroupingPartitionRouter) + grouping_router.underlying_partition_router = union_router + child_stream.get_partition_router.return_value = grouping_router + else: + child_stream.get_partition_router.return_value = union_router + + handler = ConcurrentReadProcessor( + [parent_a, parent_b, child_stream], + partition_enqueuer, + thread_pool_manager, + logger, + slice_logger, + message_repository, + partition_reader, + ) + + parent_names = handler._collect_all_parent_stream_names("child") + assert parent_names == {"parent_a", "parent_b"}