From eca483644a228388f4678ea64865e0a528a8b64d Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 10 Jul 2026 13:27:23 +0000 Subject: [PATCH 1/6] feat(low-code): add UnionPartitionRouter component Co-Authored-By: Daryna Ishchenko --- airbyte_cdk/__init__.py | 2 + .../concurrent_read_processor.py | 23 +- .../concurrent_declarative_source.py | 49 ++-- .../declarative_component_schema.yaml | 41 +++ ...legacy_to_per_partition_state_migration.py | 19 +- .../models/declarative_component_schema.py | 237 ++++++++++-------- .../parsers/model_to_component_factory.py | 53 +++- .../declarative/partition_routers/__init__.py | 4 + .../union_partition_router.py | 114 +++++++++ .../test_legacy_to_per_partition_migration.py | 67 +++++ .../test_model_to_component_factory.py | 137 ++++++++++ .../test_union_partition_router.py | 199 +++++++++++++++ 12 files changed, 807 insertions(+), 138 deletions(-) create mode 100644 airbyte_cdk/sources/declarative/partition_routers/union_partition_router.py create mode 100644 unit_tests/sources/declarative/partition_routers/test_union_partition_router.py 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..9390ede19a 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,43 @@ 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. + type: string + examples: + - "repository" + 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..43f58cfeac 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 @@ -7,6 +7,7 @@ from airbyte_cdk.sources.declarative.models import ( DatetimeBasedCursor, SubstreamPartitionRouter, + UnionPartitionRouter, ) from airbyte_cdk.sources.declarative.models.declarative_component_schema import ParentStreamConfig @@ -51,6 +52,11 @@ def __init__( ).eval(self._config) def _get_partition_field(self, partition_router: SubstreamPartitionRouter) -> str: + # UnionPartitionRouter normalizes all its children's partitions to a single + # explicit partition field, so it can be read directly off the component. + if isinstance(partition_router, UnionPartitionRouter): + return partition_router.partition_field + parent_stream_config = partition_router.parent_stream_configs[0] # Retrieve the partition field with a condition, as properties are returned as a dictionary for custom components. @@ -66,11 +72,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..96f5a7a39c 100644 --- a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py +++ b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py @@ -6,7 +6,7 @@ from enum import Enum from typing import Any, Dict, List, Literal, Optional, Union -from pydantic.v1 import BaseModel, Extra, Field +from pydantic.v1 import BaseModel, Extra, Field, confloat, conint from airbyte_cdk.sources.declarative.models.base_model_with_deprecations import ( BaseModelWithDeprecations, @@ -18,12 +18,6 @@ class AuthFlowType(Enum): oauth1_0 = "oauth1.0" -class ScopesJoinStrategy(Enum): - space = "space" - comma = "comma" - plus = "plus" - - class BasicHttpAuthenticator(BaseModel): type: Literal["BasicHttpAuthenticator"] username: str = Field( @@ -57,10 +51,9 @@ class DynamicStreamCheckConfig(BaseModel): dynamic_stream_name: str = Field( ..., description="The dynamic stream name.", title="Dynamic Stream Name" ) - stream_count: Optional[int] = Field( + stream_count: Optional[conint(ge=1)] = Field( None, description="The number of streams to attempt reading from during a check operation. If unset, all generated streams are checked. Must be a positive integer; if it exceeds the total number of available streams, all streams are checked.", - ge=1, title="Stream Count", ) @@ -98,17 +91,16 @@ class ConcurrencyLevel(BaseModel): class ConstantBackoffStrategy(BaseModel): type: Literal["ConstantBackoffStrategy"] - backoff_time_in_seconds: Union[float, str] = Field( + backoff_time_in_seconds: Union[confloat(ge=0.0), str] = Field( ..., description="Backoff time in seconds.", examples=[30, 30.5, "{{ config['backoff_time'] }}"], title="Backoff Time", ) - jitter_range_in_seconds: Optional[float] = Field( + jitter_range_in_seconds: Optional[confloat(ge=0.0)] = Field( None, description="Optional additive jitter range in seconds. When set, the backoff time is uniformly distributed between backoff_time_in_seconds and backoff_time_in_seconds + (jitter_range_in_seconds * 2), so jitter only increases the base backoff.", examples=[15], - ge=0, title="Jitter Range", ) parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") @@ -496,7 +488,7 @@ class Config: ) weight: Optional[Union[int, str]] = Field( None, - description="The weight of a request matching this matcher when acquiring a call from the rate limiter. Different endpoints can consume different amounts from a shared budget by specifying different weights. If not set, each request counts as 1.", + description="The weight of a request matching this matcher when acquiring a call from the rate limiter. Different endpoints can consume different amounts from a shared budget by specifying different weights. If not set, each request counts as 1.\n", title="Weight", ) @@ -516,6 +508,32 @@ class OnNoRecords(Enum): emit_parent = "emit_parent" +class RecordExpander(BaseModel): + type: Literal["RecordExpander"] + expand_records_from_field: List[str] = Field( + ..., + description="Path to a nested array field within each record. Items from this array will be extracted and emitted as separate records. Supports wildcards (*) for matching multiple arrays.", + examples=[ + ["lines", "data"], + ["items"], + ["nested", "array"], + ["sections", "*", "items"], + ], + title="Expand Records From Field", + ) + remain_original_record: Optional[bool] = Field( + False, + description='If true, each expanded record will include the original parent record in an "original_record" field. Defaults to false.', + title="Remain Original Record", + ) + on_no_records: Optional[OnNoRecords] = Field( + OnNoRecords.skip, + description='Behavior when the expansion path is missing, not a list, or an empty list. "skip" (default) emits nothing. "emit_parent" emits the original parent record unchanged.', + title="On No Records", + ) + parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") + + class ExponentialBackoffStrategy(BaseModel): type: Literal["ExponentialBackoffStrategy"] factor: Optional[Union[float, str]] = Field( @@ -524,11 +542,10 @@ class ExponentialBackoffStrategy(BaseModel): examples=[5, 5.5, "10"], title="Factor", ) - jitter_range_in_seconds: Optional[float] = Field( + jitter_range_in_seconds: Optional[confloat(ge=0.0)] = Field( None, description="Optional additive jitter range in seconds. When set, the backoff time is uniformly distributed between computed_backoff and computed_backoff + (jitter_range_in_seconds * 2), so jitter only increases the computed backoff.", examples=[2], - ge=0, title="Jitter Range", ) parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") @@ -683,12 +700,13 @@ class JsonItemsDecoder(BaseModel): type: Literal["JsonItemsDecoder"] items_path: str = Field( ..., - description="Dot-separated path to the JSON array whose elements should be yielded as records. Uses `ijson` path syntax (e.g. `data.users`), not JSONPath syntax \u2014 do not include leading `$.` or trailing `[*]`.", + description="Dot-separated path to the JSON array whose elements should be yielded as records. Uses `ijson` path syntax (e.g. `data.users`), not JSONPath syntax — do not include leading `$.` or trailing `[*]`.", + examples=["dataByDepartmentAndSearchTerm", "dataByAsin", "data.users"], title="Items Path", ) encoding: Optional[str] = Field( "utf-8", - description="The character encoding of the JSON data. Defaults to UTF-8.", + description="Text encoding used to decode the streamed bytes before JSON parsing.", title="Encoding", ) @@ -851,22 +869,32 @@ class NoPagination(BaseModel): type: Literal["NoPagination"] -class State(BaseModel): +class Scope(BaseModel): class Config: extra = Extra.allow - min: int - max: int + scope: str = Field(..., description="The OAuth scope string to request from the provider.") -class OAuthScope(BaseModel): +class OptionalScope(BaseModel): class Config: extra = Extra.allow - scope: str = Field( - ..., - description="The OAuth scope string to request from the provider.", - ) + scope: str = Field(..., description="The OAuth scope string to request from the provider.") + + +class ScopesJoinStrategy(Enum): + space = "space" + comma = "comma" + plus = "plus" + + +class State(BaseModel): + class Config: + extra = Extra.allow + + min: int + max: int class OauthConnectorInputSpecification(BaseModel): @@ -888,17 +916,13 @@ class Config: examples=["user:read user:read_orders workspaces:read"], title="Scopes", ) - # NOTE: scopes, optional_scopes, and scopes_join_strategy are processed by the - # platform OAuth handler (DeclarativeOAuthSpecHandler.kt), not by the CDK runtime. - # The CDK schema defines the manifest contract; the platform reads these fields - # during the OAuth consent flow to build the authorization URL. - scopes: Optional[List[OAuthScope]] = Field( + scopes: Optional[List[Scope]] = Field( None, description="List of OAuth scope objects. When present, takes precedence over the `scope` string property.\nThe scope values are joined using the `scopes_join_strategy` (default: space) before being\nsent to the OAuth provider.", examples=[[{"scope": "user:read"}, {"scope": "user:write"}]], title="Scopes", ) - optional_scopes: Optional[List[OAuthScope]] = Field( + optional_scopes: Optional[List[OptionalScope]] = Field( None, description="Optional OAuth scope objects that may or may not be granted.", examples=[[{"scope": "admin:read"}]], @@ -1271,7 +1295,14 @@ class AsyncJobStatusMap(BaseModel): completed: List[str] failed: List[str] timeout: List[str] - skipped: Optional[List[str]] = None + skipped: Optional[List[str]] = Field( + None, + description="Statuses that indicate the job was skipped because there is no data to return. Jobs with these statuses will not be retried and no records will be fetched.", + ) + + +class BlockSimultaneousSyncsAction(BaseModel): + type: Literal["BlockSimultaneousSyncsAction"] class ValueType(Enum): @@ -2118,28 +2149,23 @@ class DefaultPaginator(BaseModel): parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") -class RecordExpander(BaseModel): - type: Literal["RecordExpander"] - expand_records_from_field: List[str] = Field( +class DpathExtractor(BaseModel): + type: Literal["DpathExtractor"] + field_path: List[str] = Field( ..., - description="Path to a nested array field within each record. Items from this array will be extracted and emitted as separate records. Supports wildcards (*) for matching multiple arrays.", + description='List of potentially nested fields describing the full path of the field to extract. Use "*" to extract all values from an array. See more info in the [docs](https://docs.airbyte.com/connector-development/config-based/understanding-the-yaml-file/record-selector).', examples=[ - ["lines", "data"], - ["items"], - ["nested", "array"], - ["sections", "*", "items"], + ["data"], + ["data", "records"], + ["data", "{{ parameters.name }}"], + ["data", "*", "record"], ], - title="Expand Records From Field", - ) - remain_original_record: Optional[bool] = Field( - False, - description='If true, each expanded record will include the original parent record in an "original_record" field. Defaults to false.', - title="Remain Original Record", + title="Field Path", ) - on_no_records: Optional[OnNoRecords] = Field( - OnNoRecords.skip, - description='Behavior when the expansion path is missing, not a list, or an empty list. "skip" (default) emits nothing. "emit_parent" emits the original parent record unchanged.', - title="On No Records", + record_expander: Optional[RecordExpander] = Field( + None, + description="Optional component to expand records by extracting items from nested array fields.", + title="Record Expander", ) parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") @@ -2212,6 +2238,27 @@ class ListPartitionRouter(BaseModel): parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") +class RecordSelector(BaseModel): + type: Literal["RecordSelector"] + extractor: Union[DpathExtractor, CustomRecordExtractor] + record_filter: Optional[Union[RecordFilter, CustomRecordFilter]] = Field( + None, + description="Responsible for filtering records to be emitted by the Source.", + title="Record Filter", + ) + schema_normalization: Optional[Union[SchemaNormalization, CustomSchemaNormalization]] = Field( + None, + description="Responsible for normalization according to the schema.", + title="Schema Normalization", + ) + transform_before_filtering: Optional[bool] = Field( + None, + description="If true, transformation will be applied before record filtering.", + title="Transform Before Filtering", + ) + parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") + + class PaginationReset(BaseModel): type: Literal["PaginationReset"] action: Action1 @@ -2332,27 +2379,6 @@ class Config: ) -class DpathExtractor(BaseModel): - type: Literal["DpathExtractor"] - field_path: List[str] = Field( - ..., - description='List of potentially nested fields describing the full path of the field to extract. Use "*" to extract all values from an array. See more info in the [docs](https://docs.airbyte.com/connector-development/config-based/understanding-the-yaml-file/record-selector).', - examples=[ - ["data"], - ["data", "records"], - ["data", "{{ parameters.name }}"], - ["data", "*", "record"], - ], - title="Field Path", - ) - record_expander: Optional[RecordExpander] = Field( - None, - description="Optional component to expand records by extracting items from nested array fields.", - title="Record Expander", - ) - parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") - - class ZipfileDecoder(BaseModel): class Config: extra = Extra.allow @@ -2365,27 +2391,6 @@ class Config: ) -class RecordSelector(BaseModel): - type: Literal["RecordSelector"] - extractor: Union[DpathExtractor, CustomRecordExtractor] - record_filter: Optional[Union[RecordFilter, CustomRecordFilter]] = Field( - None, - description="Responsible for filtering records to be emitted by the Source.", - title="Record Filter", - ) - schema_normalization: Optional[Union[SchemaNormalization, CustomSchemaNormalization]] = Field( - None, - description="Responsible for normalization according to the schema.", - title="Schema Normalization", - ) - transform_before_filtering: Optional[bool] = Field( - None, - description="If true, transformation will be applied before record filtering.", - title="Transform Before Filtering", - ) - parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") - - class ConfigMigration(BaseModel): type: Literal["ConfigMigration"] description: Optional[str] = Field( @@ -2478,7 +2483,7 @@ class Config: api_budget: Optional[HTTPAPIBudget] = None stream_groups: Optional[Dict[str, StreamGroup]] = Field( None, - description="Groups of streams that share a common resource and should not be read simultaneously. Each group defines a set of stream references and an action that controls how concurrent reads are managed. Only applies to ConcurrentDeclarativeSource.", + description="Groups of streams that share a common resource and should not be read simultaneously. Each group defines a set of stream references and an action that controls how concurrent reads are managed. Only applies to ConcurrentDeclarativeSource.\n", title="Stream Groups", ) max_concurrent_async_job_count: Optional[Union[int, str]] = Field( @@ -2518,7 +2523,7 @@ class Config: api_budget: Optional[HTTPAPIBudget] = None stream_groups: Optional[Dict[str, StreamGroup]] = Field( None, - description="Groups of streams that share a common resource and should not be read simultaneously. Each group defines a set of stream references and an action that controls how concurrent reads are managed. Only applies to ConcurrentDeclarativeSource.", + description="Groups of streams that share a common resource and should not be read simultaneously. Each group defines a set of stream references and an action that controls how concurrent reads are managed. Only applies to ConcurrentDeclarativeSource.\n", title="Stream Groups", ) max_concurrent_async_job_count: Optional[Union[int, str]] = Field( @@ -3014,7 +3019,7 @@ class StateDelegatingStream(BaseModel): ) api_retention_period: Optional[str] = Field( None, - description="The data retention period of the incremental API (ISO8601 duration). If the cursor value is older than this retention period, the connector will automatically fall back to a full refresh to avoid data loss.\nThis is useful for APIs like Stripe Events API which only retain data for 30 days.\n * **PT1H**: 1 hour\n * **P1D**: 1 day\n * **P1W**: 1 week\n * **P1M**: 1 month\n * **P1Y**: 1 year\n * **P30D**: 30 days\n", + description="The data retention period of the incremental API (ISO8601 duration). If the cursor value is older than this retention period, the connector will automatically fall back to a full refresh to avoid data loss.\nThis is useful for APIs like Stripe Events API which only retain data for 30 days.\n* **PT1H**: 1 hour\n* **P1D**: 1 day\n* **P1W**: 1 week\n* **P1M**: 1 month\n* **P1Y**: 1 year\n* **P30D**: 30 days\n", examples=["P30D", "P90D", "P1Y"], title="API Retention Period", ) @@ -3065,12 +3070,14 @@ class SimpleRetriever(BaseModel): SubstreamPartitionRouter, ListPartitionRouter, GroupingPartitionRouter, + UnionPartitionRouter, CustomPartitionRouter, List[ Union[ SubstreamPartitionRouter, ListPartitionRouter, GroupingPartitionRouter, + UnionPartitionRouter, CustomPartitionRouter, ] ], @@ -3114,10 +3121,9 @@ class AsyncRetriever(BaseModel): None, description="The time in minutes after which the single Async Job should be considered as Timed Out.", ) - failed_retry_wait_time_in_seconds: Optional[Union[int, str]] = Field( + failed_retry_wait_time_in_seconds: Optional[Union[conint(ge=1), str]] = Field( None, description="Time in seconds to wait before retrying a failed async job. Only applies to jobs that ran on the API side and reported a FAILED status (e.g. report generation failed due to a cooldown). Creation failures (HTTP errors when starting a job, such as 429s) and TIMED_OUT jobs are retried immediately and are not affected by this setting. When set, the orchestrator defers retry of real failed jobs until the wait time has elapsed, without blocking other jobs.", - ge=1, ) download_target_requester: Optional[Union[HttpRequester, CustomRequester]] = Field( None, @@ -3144,12 +3150,14 @@ class AsyncRetriever(BaseModel): ListPartitionRouter, SubstreamPartitionRouter, GroupingPartitionRouter, + UnionPartitionRouter, CustomPartitionRouter, List[ Union[ ListPartitionRouter, SubstreamPartitionRouter, GroupingPartitionRouter, + UnionPartitionRouter, CustomPartitionRouter, ] ], @@ -3196,20 +3204,14 @@ class AsyncRetriever(BaseModel): parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") -class BlockSimultaneousSyncsAction(BaseModel): - type: Literal["BlockSimultaneousSyncsAction"] - - class StreamGroup(BaseModel): - streams: List[str] = Field( + streams: List[DeclarativeStream] = Field( ..., - description='List of references to streams that belong to this group. Use JSON references to stream definitions (e.g., "#/definitions/my_stream").', + description="List of references to streams that belong to this group.\n", title="Streams", ) action: BlockSimultaneousSyncsAction = Field( - ..., - description="The action to apply to streams in this group.", - title="Action", + ..., description="The action to apply to streams in this group.", title="Action" ) @@ -3246,6 +3248,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.", + examples=["repository"], + 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( 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..9b464ae50c 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)}" ) - 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,43 @@ 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 + ] + + # 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. + 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=model.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..2c29818b16 --- /dev/null +++ b/airbyte_cdk/sources/declarative/partition_routers/union_partition_router.py @@ -0,0 +1,114 @@ +# +# Copyright (c) 2026 Airbyte, Inc., all rights reserved. +# + +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 + + +@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] + if value in seen: + continue + seen.add(value) + 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.""" + merged_state: dict[str, StreamState] = {} + for router in self.partition_routers: + child_state = router.get_stream_state() + if child_state: + merged_state.update(child_state) + return merged_state or None 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..62afa393b6 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", [ 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..6753fc67fa 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,139 @@ 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" + + +@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", + ) + + 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..eb74d7bc6b --- /dev/null +++ b/unit_tests/sources/declarative/partition_routers/test_union_partition_router.py @@ -0,0 +1,199 @@ +# +# 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], 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_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) == {} From 30e56f8c9583baf77aa4265c5f182eea23c081d2 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 10 Jul 2026 13:33:37 +0000 Subject: [PATCH 2/6] fix: restore generated models style, apply minimal UnionPartitionRouter additions Co-Authored-By: Daryna Ishchenko --- .../models/declarative_component_schema.py | 211 +++++++++--------- 1 file changed, 107 insertions(+), 104 deletions(-) diff --git a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py index 96f5a7a39c..646602d042 100644 --- a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py +++ b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py @@ -6,7 +6,7 @@ from enum import Enum from typing import Any, Dict, List, Literal, Optional, Union -from pydantic.v1 import BaseModel, Extra, Field, confloat, conint +from pydantic.v1 import BaseModel, Extra, Field from airbyte_cdk.sources.declarative.models.base_model_with_deprecations import ( BaseModelWithDeprecations, @@ -18,6 +18,12 @@ class AuthFlowType(Enum): oauth1_0 = "oauth1.0" +class ScopesJoinStrategy(Enum): + space = "space" + comma = "comma" + plus = "plus" + + class BasicHttpAuthenticator(BaseModel): type: Literal["BasicHttpAuthenticator"] username: str = Field( @@ -51,9 +57,10 @@ class DynamicStreamCheckConfig(BaseModel): dynamic_stream_name: str = Field( ..., description="The dynamic stream name.", title="Dynamic Stream Name" ) - stream_count: Optional[conint(ge=1)] = Field( + stream_count: Optional[int] = Field( None, description="The number of streams to attempt reading from during a check operation. If unset, all generated streams are checked. Must be a positive integer; if it exceeds the total number of available streams, all streams are checked.", + ge=1, title="Stream Count", ) @@ -91,16 +98,17 @@ class ConcurrencyLevel(BaseModel): class ConstantBackoffStrategy(BaseModel): type: Literal["ConstantBackoffStrategy"] - backoff_time_in_seconds: Union[confloat(ge=0.0), str] = Field( + backoff_time_in_seconds: Union[float, str] = Field( ..., description="Backoff time in seconds.", examples=[30, 30.5, "{{ config['backoff_time'] }}"], title="Backoff Time", ) - jitter_range_in_seconds: Optional[confloat(ge=0.0)] = Field( + jitter_range_in_seconds: Optional[float] = Field( None, description="Optional additive jitter range in seconds. When set, the backoff time is uniformly distributed between backoff_time_in_seconds and backoff_time_in_seconds + (jitter_range_in_seconds * 2), so jitter only increases the base backoff.", examples=[15], + ge=0, title="Jitter Range", ) parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") @@ -488,7 +496,7 @@ class Config: ) weight: Optional[Union[int, str]] = Field( None, - description="The weight of a request matching this matcher when acquiring a call from the rate limiter. Different endpoints can consume different amounts from a shared budget by specifying different weights. If not set, each request counts as 1.\n", + description="The weight of a request matching this matcher when acquiring a call from the rate limiter. Different endpoints can consume different amounts from a shared budget by specifying different weights. If not set, each request counts as 1.", title="Weight", ) @@ -508,32 +516,6 @@ class OnNoRecords(Enum): emit_parent = "emit_parent" -class RecordExpander(BaseModel): - type: Literal["RecordExpander"] - expand_records_from_field: List[str] = Field( - ..., - description="Path to a nested array field within each record. Items from this array will be extracted and emitted as separate records. Supports wildcards (*) for matching multiple arrays.", - examples=[ - ["lines", "data"], - ["items"], - ["nested", "array"], - ["sections", "*", "items"], - ], - title="Expand Records From Field", - ) - remain_original_record: Optional[bool] = Field( - False, - description='If true, each expanded record will include the original parent record in an "original_record" field. Defaults to false.', - title="Remain Original Record", - ) - on_no_records: Optional[OnNoRecords] = Field( - OnNoRecords.skip, - description='Behavior when the expansion path is missing, not a list, or an empty list. "skip" (default) emits nothing. "emit_parent" emits the original parent record unchanged.', - title="On No Records", - ) - parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") - - class ExponentialBackoffStrategy(BaseModel): type: Literal["ExponentialBackoffStrategy"] factor: Optional[Union[float, str]] = Field( @@ -542,10 +524,11 @@ class ExponentialBackoffStrategy(BaseModel): examples=[5, 5.5, "10"], title="Factor", ) - jitter_range_in_seconds: Optional[confloat(ge=0.0)] = Field( + jitter_range_in_seconds: Optional[float] = Field( None, description="Optional additive jitter range in seconds. When set, the backoff time is uniformly distributed between computed_backoff and computed_backoff + (jitter_range_in_seconds * 2), so jitter only increases the computed backoff.", examples=[2], + ge=0, title="Jitter Range", ) parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") @@ -700,13 +683,12 @@ class JsonItemsDecoder(BaseModel): type: Literal["JsonItemsDecoder"] items_path: str = Field( ..., - description="Dot-separated path to the JSON array whose elements should be yielded as records. Uses `ijson` path syntax (e.g. `data.users`), not JSONPath syntax — do not include leading `$.` or trailing `[*]`.", - examples=["dataByDepartmentAndSearchTerm", "dataByAsin", "data.users"], + description="Dot-separated path to the JSON array whose elements should be yielded as records. Uses `ijson` path syntax (e.g. `data.users`), not JSONPath syntax \u2014 do not include leading `$.` or trailing `[*]`.", title="Items Path", ) encoding: Optional[str] = Field( "utf-8", - description="Text encoding used to decode the streamed bytes before JSON parsing.", + description="The character encoding of the JSON data. Defaults to UTF-8.", title="Encoding", ) @@ -869,32 +851,22 @@ class NoPagination(BaseModel): type: Literal["NoPagination"] -class Scope(BaseModel): - class Config: - extra = Extra.allow - - scope: str = Field(..., description="The OAuth scope string to request from the provider.") - - -class OptionalScope(BaseModel): +class State(BaseModel): class Config: extra = Extra.allow - scope: str = Field(..., description="The OAuth scope string to request from the provider.") - - -class ScopesJoinStrategy(Enum): - space = "space" - comma = "comma" - plus = "plus" + min: int + max: int -class State(BaseModel): +class OAuthScope(BaseModel): class Config: extra = Extra.allow - min: int - max: int + scope: str = Field( + ..., + description="The OAuth scope string to request from the provider.", + ) class OauthConnectorInputSpecification(BaseModel): @@ -916,13 +888,17 @@ class Config: examples=["user:read user:read_orders workspaces:read"], title="Scopes", ) - scopes: Optional[List[Scope]] = Field( + # NOTE: scopes, optional_scopes, and scopes_join_strategy are processed by the + # platform OAuth handler (DeclarativeOAuthSpecHandler.kt), not by the CDK runtime. + # The CDK schema defines the manifest contract; the platform reads these fields + # during the OAuth consent flow to build the authorization URL. + scopes: Optional[List[OAuthScope]] = Field( None, description="List of OAuth scope objects. When present, takes precedence over the `scope` string property.\nThe scope values are joined using the `scopes_join_strategy` (default: space) before being\nsent to the OAuth provider.", examples=[[{"scope": "user:read"}, {"scope": "user:write"}]], title="Scopes", ) - optional_scopes: Optional[List[OptionalScope]] = Field( + optional_scopes: Optional[List[OAuthScope]] = Field( None, description="Optional OAuth scope objects that may or may not be granted.", examples=[[{"scope": "admin:read"}]], @@ -1295,14 +1271,7 @@ class AsyncJobStatusMap(BaseModel): completed: List[str] failed: List[str] timeout: List[str] - skipped: Optional[List[str]] = Field( - None, - description="Statuses that indicate the job was skipped because there is no data to return. Jobs with these statuses will not be retried and no records will be fetched.", - ) - - -class BlockSimultaneousSyncsAction(BaseModel): - type: Literal["BlockSimultaneousSyncsAction"] + skipped: Optional[List[str]] = None class ValueType(Enum): @@ -2149,23 +2118,28 @@ class DefaultPaginator(BaseModel): parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") -class DpathExtractor(BaseModel): - type: Literal["DpathExtractor"] - field_path: List[str] = Field( +class RecordExpander(BaseModel): + type: Literal["RecordExpander"] + expand_records_from_field: List[str] = Field( ..., - description='List of potentially nested fields describing the full path of the field to extract. Use "*" to extract all values from an array. See more info in the [docs](https://docs.airbyte.com/connector-development/config-based/understanding-the-yaml-file/record-selector).', + description="Path to a nested array field within each record. Items from this array will be extracted and emitted as separate records. Supports wildcards (*) for matching multiple arrays.", examples=[ - ["data"], - ["data", "records"], - ["data", "{{ parameters.name }}"], - ["data", "*", "record"], + ["lines", "data"], + ["items"], + ["nested", "array"], + ["sections", "*", "items"], ], - title="Field Path", + title="Expand Records From Field", ) - record_expander: Optional[RecordExpander] = Field( - None, - description="Optional component to expand records by extracting items from nested array fields.", - title="Record Expander", + remain_original_record: Optional[bool] = Field( + False, + description='If true, each expanded record will include the original parent record in an "original_record" field. Defaults to false.', + title="Remain Original Record", + ) + on_no_records: Optional[OnNoRecords] = Field( + OnNoRecords.skip, + description='Behavior when the expansion path is missing, not a list, or an empty list. "skip" (default) emits nothing. "emit_parent" emits the original parent record unchanged.', + title="On No Records", ) parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") @@ -2238,27 +2212,6 @@ class ListPartitionRouter(BaseModel): parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") -class RecordSelector(BaseModel): - type: Literal["RecordSelector"] - extractor: Union[DpathExtractor, CustomRecordExtractor] - record_filter: Optional[Union[RecordFilter, CustomRecordFilter]] = Field( - None, - description="Responsible for filtering records to be emitted by the Source.", - title="Record Filter", - ) - schema_normalization: Optional[Union[SchemaNormalization, CustomSchemaNormalization]] = Field( - None, - description="Responsible for normalization according to the schema.", - title="Schema Normalization", - ) - transform_before_filtering: Optional[bool] = Field( - None, - description="If true, transformation will be applied before record filtering.", - title="Transform Before Filtering", - ) - parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") - - class PaginationReset(BaseModel): type: Literal["PaginationReset"] action: Action1 @@ -2379,6 +2332,27 @@ class Config: ) +class DpathExtractor(BaseModel): + type: Literal["DpathExtractor"] + field_path: List[str] = Field( + ..., + description='List of potentially nested fields describing the full path of the field to extract. Use "*" to extract all values from an array. See more info in the [docs](https://docs.airbyte.com/connector-development/config-based/understanding-the-yaml-file/record-selector).', + examples=[ + ["data"], + ["data", "records"], + ["data", "{{ parameters.name }}"], + ["data", "*", "record"], + ], + title="Field Path", + ) + record_expander: Optional[RecordExpander] = Field( + None, + description="Optional component to expand records by extracting items from nested array fields.", + title="Record Expander", + ) + parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") + + class ZipfileDecoder(BaseModel): class Config: extra = Extra.allow @@ -2391,6 +2365,27 @@ class Config: ) +class RecordSelector(BaseModel): + type: Literal["RecordSelector"] + extractor: Union[DpathExtractor, CustomRecordExtractor] + record_filter: Optional[Union[RecordFilter, CustomRecordFilter]] = Field( + None, + description="Responsible for filtering records to be emitted by the Source.", + title="Record Filter", + ) + schema_normalization: Optional[Union[SchemaNormalization, CustomSchemaNormalization]] = Field( + None, + description="Responsible for normalization according to the schema.", + title="Schema Normalization", + ) + transform_before_filtering: Optional[bool] = Field( + None, + description="If true, transformation will be applied before record filtering.", + title="Transform Before Filtering", + ) + parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") + + class ConfigMigration(BaseModel): type: Literal["ConfigMigration"] description: Optional[str] = Field( @@ -2483,7 +2478,7 @@ class Config: api_budget: Optional[HTTPAPIBudget] = None stream_groups: Optional[Dict[str, StreamGroup]] = Field( None, - description="Groups of streams that share a common resource and should not be read simultaneously. Each group defines a set of stream references and an action that controls how concurrent reads are managed. Only applies to ConcurrentDeclarativeSource.\n", + description="Groups of streams that share a common resource and should not be read simultaneously. Each group defines a set of stream references and an action that controls how concurrent reads are managed. Only applies to ConcurrentDeclarativeSource.", title="Stream Groups", ) max_concurrent_async_job_count: Optional[Union[int, str]] = Field( @@ -2523,7 +2518,7 @@ class Config: api_budget: Optional[HTTPAPIBudget] = None stream_groups: Optional[Dict[str, StreamGroup]] = Field( None, - description="Groups of streams that share a common resource and should not be read simultaneously. Each group defines a set of stream references and an action that controls how concurrent reads are managed. Only applies to ConcurrentDeclarativeSource.\n", + description="Groups of streams that share a common resource and should not be read simultaneously. Each group defines a set of stream references and an action that controls how concurrent reads are managed. Only applies to ConcurrentDeclarativeSource.", title="Stream Groups", ) max_concurrent_async_job_count: Optional[Union[int, str]] = Field( @@ -3019,7 +3014,7 @@ class StateDelegatingStream(BaseModel): ) api_retention_period: Optional[str] = Field( None, - description="The data retention period of the incremental API (ISO8601 duration). If the cursor value is older than this retention period, the connector will automatically fall back to a full refresh to avoid data loss.\nThis is useful for APIs like Stripe Events API which only retain data for 30 days.\n* **PT1H**: 1 hour\n* **P1D**: 1 day\n* **P1W**: 1 week\n* **P1M**: 1 month\n* **P1Y**: 1 year\n* **P30D**: 30 days\n", + description="The data retention period of the incremental API (ISO8601 duration). If the cursor value is older than this retention period, the connector will automatically fall back to a full refresh to avoid data loss.\nThis is useful for APIs like Stripe Events API which only retain data for 30 days.\n * **PT1H**: 1 hour\n * **P1D**: 1 day\n * **P1W**: 1 week\n * **P1M**: 1 month\n * **P1Y**: 1 year\n * **P30D**: 30 days\n", examples=["P30D", "P90D", "P1Y"], title="API Retention Period", ) @@ -3121,9 +3116,10 @@ class AsyncRetriever(BaseModel): None, description="The time in minutes after which the single Async Job should be considered as Timed Out.", ) - failed_retry_wait_time_in_seconds: Optional[Union[conint(ge=1), str]] = Field( + failed_retry_wait_time_in_seconds: Optional[Union[int, str]] = Field( None, description="Time in seconds to wait before retrying a failed async job. Only applies to jobs that ran on the API side and reported a FAILED status (e.g. report generation failed due to a cooldown). Creation failures (HTTP errors when starting a job, such as 429s) and TIMED_OUT jobs are retried immediately and are not affected by this setting. When set, the orchestrator defers retry of real failed jobs until the wait time has elapsed, without blocking other jobs.", + ge=1, ) download_target_requester: Optional[Union[HttpRequester, CustomRequester]] = Field( None, @@ -3204,14 +3200,20 @@ class AsyncRetriever(BaseModel): parameters: Optional[Dict[str, Any]] = Field(None, alias="$parameters") +class BlockSimultaneousSyncsAction(BaseModel): + type: Literal["BlockSimultaneousSyncsAction"] + + class StreamGroup(BaseModel): - streams: List[DeclarativeStream] = Field( + streams: List[str] = Field( ..., - description="List of references to streams that belong to this group.\n", + description='List of references to streams that belong to this group. Use JSON references to stream definitions (e.g., "#/definitions/my_stream").', title="Streams", ) action: BlockSimultaneousSyncsAction = Field( - ..., description="The action to apply to streams in this group.", title="Action" + ..., + description="The action to apply to streams in this group.", + title="Action", ) @@ -3320,3 +3322,4 @@ class DynamicDeclarativeStream(BaseModel): PropertiesFromEndpoint.update_forward_refs() SimpleRetriever.update_forward_refs() AsyncRetriever.update_forward_refs() +UnionPartitionRouter.update_forward_refs() From fbba487cbd59c02dc455c6fe068210cc3b57596d Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 10 Jul 2026 14:27:03 +0000 Subject: [PATCH 3/6] test: add concurrency-graph/caching coverage and legacy-state round-trip for UnionPartitionRouter Co-Authored-By: Daryna Ishchenko --- .../parsers/model_to_component_factory.py | 4 +- .../union_partition_router.py | 8 +- .../test_concurrent_declarative_source.py | 409 ++++++++++++++++++ .../test_concurrent_read_processor.py | 64 +++ 4 files changed, 483 insertions(+), 2 deletions(-) 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 9b464ae50c..20eb2243bb 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -4578,7 +4578,9 @@ def create_union_partition_router( # 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. + # (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( diff --git a/airbyte_cdk/sources/declarative/partition_routers/union_partition_router.py b/airbyte_cdk/sources/declarative/partition_routers/union_partition_router.py index 2c29818b16..5a27077405 100644 --- a/airbyte_cdk/sources/declarative/partition_routers/union_partition_router.py +++ b/airbyte_cdk/sources/declarative/partition_routers/union_partition_router.py @@ -105,7 +105,13 @@ def get_request_body_json( return {} def get_stream_state(self) -> Optional[Mapping[str, StreamState]]: - """Merge the parent stream states of all child partition routers.""" + """ + 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. + """ merged_state: dict[str, StreamState] = {} for router in self.partition_routers: child_state = router.get_stream_state() 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"} From 380963a6e2e3959a312f398ffd1a817bb1db2d40 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 15 Jul 2026 12:17:55 +0000 Subject: [PATCH 4/6] fix(low-code): harden UnionPartitionRouter per review feedback - fail fast at factory time when a built-in child router emits a different partition field - treat partition_field as a plain string in LegacyToPerPartitionStateMigration to match runtime - raise a pointed error for unhashable partition values during dedup - warn on parent stream state key collisions in get_stream_state; return {} instead of None - widen migration type hints and factory error message to cover accepted router types Co-Authored-By: Daryna Ishchenko --- ...legacy_to_per_partition_state_migration.py | 29 ++++---- .../parsers/model_to_component_factory.py | 31 +++++++- .../union_partition_router.py | 26 +++++-- .../test_legacy_to_per_partition_migration.py | 2 +- .../test_model_to_component_factory.py | 71 +++++++++++++++++++ .../test_union_partition_router.py | 34 ++++++++- 6 files changed, 173 insertions(+), 20 deletions(-) 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 43f58cfeac..a9ffaa17c4 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,10 +1,11 @@ # 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, @@ -35,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,20 +47,22 @@ def __init__( self._cursor = cursor self._config = config self._parameters = parameters - self._partition_key_field = InterpolatedString.create( - self._get_partition_field(self._partition_router), parameters=self._parameters - ).eval(self._config) + if isinstance(partition_router, UnionPartitionRouter): + # UnionPartitionRouter treats partition_field as a plain string (it is never + # interpolated at runtime), so the migration must not interpolate it either. + self._partition_key_field = partition_router.partition_field + else: + self._partition_key_field = InterpolatedString.create( + 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: - # UnionPartitionRouter normalizes all its children's partitions to a single - # explicit partition field, so it can be read directly off the component. - if isinstance(partition_router, UnionPartitionRouter): - return partition_router.partition_field - - parent_stream_config = partition_router.parent_stream_configs[0] + def _get_partition_field( + self, partition_router: Union[SubstreamPartitionRouter, CustomPartitionRouter] + ) -> str: + 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 = ( 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 20eb2243bb..8433616f6f 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -1149,7 +1149,7 @@ def create_legacy_to_per_partition_state_migration( ), ): 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 isinstance(partition_router, UnionPartitionRouterModel) and not hasattr( partition_router, "parent_stream_configs" @@ -4575,6 +4575,35 @@ def create_union_partition_router( for child in model.partition_routers ] + # 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(child_model.partition_field) + for child_partition_field in child_partition_fields: + if child_partition_field != model.partition_field: + raise ValueError( + f"UnionPartitionRouter expects all child partition routers to emit the " + f"partition field '{model.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 diff --git a/airbyte_cdk/sources/declarative/partition_routers/union_partition_router.py b/airbyte_cdk/sources/declarative/partition_routers/union_partition_router.py index 5a27077405..66769fceb9 100644 --- a/airbyte_cdk/sources/declarative/partition_routers/union_partition_router.py +++ b/airbyte_cdk/sources/declarative/partition_routers/union_partition_router.py @@ -2,12 +2,15 @@ # 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): @@ -51,9 +54,15 @@ def stream_slices(self) -> Iterable[StreamSlice]: f"partition field '{self.partition_field}'. Got {stream_slice.partition}" ) value = stream_slice.partition[self.partition_field] - if value in seen: - continue - seen.add(value) + 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() @@ -110,11 +119,18 @@ def get_stream_state(self) -> Optional[Mapping[str, StreamState]]: 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. + 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 or None + 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 62afa393b6..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 @@ -359,7 +359,7 @@ def _union_migrator(): 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 6753fc67fa..452857aa14 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 @@ -4815,6 +4815,77 @@ def test_create_union_partition_router_with_request_option(child_router_manifest ) +@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 index eb74d7bc6b..938f2609e6 100644 --- a/unit_tests/sources/declarative/partition_routers/test_union_partition_router.py +++ b/unit_tests/sources/declarative/partition_routers/test_union_partition_router.py @@ -156,7 +156,7 @@ def test_nested_union_partition_router(): @pytest.mark.parametrize( "child_states, expected_state", [ - pytest.param([None, None], None, id="all_children_without_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"}}, @@ -185,6 +185,38 @@ def test_get_stream_state_merges_children_states(child_states, expected_state): 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"])], From ca6cc171ca00b940add3e051f2d891bd0ccce754 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 15 Jul 2026 12:39:17 +0000 Subject: [PATCH 5/6] chore: empty commit to re-trigger CI after infra flake in destination-motherduck check Co-Authored-By: Daryna Ishchenko From 36b06ad16610cdcecdf8d7c7323fbb91efe435b9 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 15 Jul 2026 12:47:40 +0000 Subject: [PATCH 6/6] feat(low-code): interpolate UnionPartitionRouter partition_field at build time partition_field depends only on config/parameters, so the factory evaluates it once and the runtime component always receives a plain string. This lets config-driven child partition fields work with the union router. The migration special case that skipped interpolation for unions is removed since the model's raw partition_field is now interpolated through the shared eval path. Schema documents the interpolation contract with an interpolated example. Co-Authored-By: Daryna Ishchenko --- .../declarative_component_schema.yaml | 6 +++- ...legacy_to_per_partition_state_migration.py | 19 ++++++------ .../models/declarative_component_schema.py | 4 +-- .../parsers/model_to_component_factory.py | 18 ++++++++--- .../test_model_to_component_factory.py | 30 +++++++++++++++++++ 5 files changed, 61 insertions(+), 16 deletions(-) diff --git a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml index 9390ede19a..778ba3bed6 100644 --- a/airbyte_cdk/sources/declarative/declarative_component_schema.yaml +++ b/airbyte_cdk/sources/declarative/declarative_component_schema.yaml @@ -4498,10 +4498,14 @@ definitions: 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. + 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`). 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 a9ffaa17c4..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 @@ -47,21 +47,22 @@ def __init__( self._cursor = cursor self._config = config self._parameters = parameters - if isinstance(partition_router, UnionPartitionRouter): - # UnionPartitionRouter treats partition_field as a plain string (it is never - # interpolated at runtime), so the migration must not interpolate it either. - self._partition_key_field = partition_router.partition_field - else: - self._partition_key_field = InterpolatedString.create( - self._get_partition_field(partition_router), parameters=self._parameters - ).eval(self._config) + self._partition_key_field = InterpolatedString.create( + 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: Union[SubstreamPartitionRouter, CustomPartitionRouter] + 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. diff --git a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py index 646602d042..3e7ce98e05 100644 --- a/airbyte_cdk/sources/declarative/models/declarative_component_schema.py +++ b/airbyte_cdk/sources/declarative/models/declarative_component_schema.py @@ -3254,8 +3254,8 @@ 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.", - examples=["repository"], + 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[ 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 8433616f6f..f78c966930 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -4575,6 +4575,12 @@ def create_union_partition_router( 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. @@ -4595,12 +4601,16 @@ def create_union_partition_router( ).eval(config) ) elif isinstance(child_model, UnionPartitionRouterModel): - child_partition_fields.append(child_model.partition_field) + 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 != model.partition_field: + if child_partition_field != partition_field: raise ValueError( f"UnionPartitionRouter expects all child partition routers to emit the " - f"partition field '{model.partition_field}', but a " + f"partition field '{partition_field}', but a " f"{child_model.type} child emits '{child_partition_field}'." ) @@ -4621,7 +4631,7 @@ def create_union_partition_router( return UnionPartitionRouter( partition_routers=partition_routers, - partition_field=model.partition_field, + partition_field=partition_field, parameters=model.parameters or {}, ) 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 452857aa14..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 @@ -4740,6 +4740,36 @@ def test_create_union_partition_router(): 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", [