diff --git a/airbyte_cdk/sources/declarative/checks/check_dynamic_stream.py b/airbyte_cdk/sources/declarative/checks/check_dynamic_stream.py index 2dac2ecc39..9b38b90257 100644 --- a/airbyte_cdk/sources/declarative/checks/check_dynamic_stream.py +++ b/airbyte_cdk/sources/declarative/checks/check_dynamic_stream.py @@ -4,7 +4,7 @@ import logging from dataclasses import InitVar, dataclass -from typing import Any, List, Mapping, Tuple, Union +from typing import Any, List, Mapping, Optional, Tuple, Union, cast from airbyte_cdk.sources import Source from airbyte_cdk.sources.declarative.checks.check_stream import evaluate_availability @@ -38,7 +38,22 @@ def check_connection( logger: logging.Logger, config: Mapping[str, Any], ) -> Tuple[bool, Any]: - streams: List[Union[Stream, AbstractStream]] = source.streams(config=config) # type: ignore # this is a migration step and we expect the declarative CDK to migrate off of ConnectionChecker + from airbyte_cdk.sources.declarative.concurrent_declarative_source import ( + ConcurrentDeclarativeSource, + ) + + streams: Optional[List[Union[Stream, AbstractStream]]] = None + if ( + isinstance(source, ConcurrentDeclarativeSource) + and self.use_check_availability + and self.stream_count >= 1 + ): + streams = cast( + Optional[List[Union[Stream, AbstractStream]]], + source._first_streams(self.stream_count), + ) + if streams is None: + streams = source.streams(config=config) # type: ignore # this is a migration step and we expect the declarative CDK to migrate off of ConnectionChecker if len(streams) == 0: return False, f"No streams to connect to from source {source}" diff --git a/airbyte_cdk/sources/declarative/checks/check_stream.py b/airbyte_cdk/sources/declarative/checks/check_stream.py index 618511f8d7..a9d6bb41e6 100644 --- a/airbyte_cdk/sources/declarative/checks/check_stream.py +++ b/airbyte_cdk/sources/declarative/checks/check_stream.py @@ -5,7 +5,7 @@ import logging import traceback from dataclasses import InitVar, dataclass -from typing import Any, Dict, List, Mapping, Optional, Tuple, Union +from typing import Any, Dict, List, Mapping, Optional, Tuple, Union, cast from airbyte_cdk.sources import Source from airbyte_cdk.sources.declarative.checks.connection_checker import ConnectionChecker @@ -74,7 +74,9 @@ def check_connection( """Checks the connection to the source and its streams.""" stream_names = self._get_stream_names(config) try: - streams: List[Union[Stream, AbstractStream]] = source.streams(config=config) # type: ignore # this is a migration step and we expect the declarative CDK to migrate off of ConnectionChecker + streams: List[Union[Stream, AbstractStream]] = self._build_streams( + source, config, stream_names + ) if not streams: return False, f"No streams to connect to from source {source}" except Exception as error: @@ -94,9 +96,9 @@ def check_connection( return stream_availability, message should_check_dynamic_streams = ( - hasattr(source, "resolved_manifest") + self.dynamic_streams_check_configs + and hasattr(source, "resolved_manifest") and hasattr(source, "dynamic_streams") - and self.dynamic_streams_check_configs ) if should_check_dynamic_streams: @@ -104,6 +106,26 @@ def check_connection( return True, None + def _build_streams( + self, + source: Source, + config: Mapping[str, Any], + stream_names: List[str], + ) -> List[Union[Stream, AbstractStream]]: + from airbyte_cdk.sources.declarative.concurrent_declarative_source import ( + ConcurrentDeclarativeSource, + ) + + if ( + isinstance(source, ConcurrentDeclarativeSource) + and not self.dynamic_streams_check_configs + ): + streams = source._streams_by_name(stream_names) + if streams is not None: + return cast(List[Union[Stream, AbstractStream]], streams) + + return source.streams(config=config) # type: ignore # this is a migration step and we expect the declarative CDK to migrate off of ConnectionChecker + def _get_stream_names(self, config: Mapping[str, Any]) -> List[str]: if ( CHECK_STREAM_NAMES_CONFIG_KEY not in config diff --git a/airbyte_cdk/sources/declarative/concurrent_declarative_source.py b/airbyte_cdk/sources/declarative/concurrent_declarative_source.py index 24803a72b9..7e7e6e3876 100644 --- a/airbyte_cdk/sources/declarative/concurrent_declarative_source.py +++ b/airbyte_cdk/sources/declarative/concurrent_declarative_source.py @@ -442,6 +442,11 @@ def streams(self, config: Mapping[str, Any]) -> List[AbstractStream]: # type: i In both case, we will assume that calling the DeclarativeStream is perfectly fine as the result for these is the same regardless of if it is a DeclarativeStream or a DefaultStream (concurrent). This should simply be removed once we have moved away from the mentioned code paths above. """ + self._prepare_for_stream_build() + stream_configs = self._stream_configs(self._source_config) + self.dynamic_streams + return self._create_streams(stream_configs, stream_configs) + + def _prepare_for_stream_build(self) -> None: if self._spec_component: self._spec_component.validate_config(self._config_for_validation) @@ -449,9 +454,14 @@ def streams(self, config: Mapping[str, Any]) -> List[AbstractStream]: # type: i if api_budget_model: self._constructor.set_api_budget(api_budget_model, self._config) - stream_configs = self._stream_configs(self._source_config) + self.dynamic_streams - - prepared_configs = self._initialize_cache_for_parent_streams(deepcopy(stream_configs)) + def _create_streams( + self, + selected_configs: List[Dict[str, Any]], + all_configs: List[Dict[str, Any]], + ) -> List[AbstractStream]: + prepared_configs = self._initialize_cache_for_parent_streams( + deepcopy(selected_configs), all_stream_configs=all_configs + ) source_streams = [ self._constructor.create_component( @@ -467,11 +477,62 @@ def streams(self, config: Mapping[str, Any]) -> List[AbstractStream]: # type: i for stream_config in prepared_configs ] - self._apply_stream_groups(source_streams) + top_level_stream_names: Set[str] = set() + for stream_config in all_configs: + stream_name = stream_config.get("name") + if isinstance(stream_name, str): + top_level_stream_names.add(stream_name) + self._apply_stream_groups(source_streams, top_level_stream_names) return source_streams - def _apply_stream_groups(self, streams: List[AbstractStream]) -> None: + def _uses_default_streams(self) -> bool: + """A subclass that overrides `streams()` must keep going through it during check.""" + return type(self).streams is ConcurrentDeclarativeSource.streams + + def _streams_by_name(self, stream_names: List[str]) -> Optional[List[AbstractStream]]: + """Build the named static streams plus every dynamic stream. + + Return None when the names are empty, any name isn't static, or `streams()` is overridden. + """ + if not self._uses_default_streams(): + return None + try: + static_configs = self._stream_configs(self._source_config) + except Exception: + return None + static_stream_names = {stream_config.get("name") for stream_config in static_configs} + if not stream_names or any(name not in static_stream_names for name in stream_names): + return None + + self._prepare_for_stream_build() + dynamic_configs = self.dynamic_streams + selected_names = set(stream_names) + selected_configs = [ + stream_config + for stream_config in static_configs + if stream_config.get("name") in selected_names + ] + dynamic_configs + return self._create_streams(selected_configs, static_configs + dynamic_configs) + + def _first_streams(self, count: int) -> Optional[List[AbstractStream]]: + """Build the first `count` streams, static then dynamic, in `streams()` order. + + `count` must be at least 1. Return None when `streams()` is overridden. + """ + if not self._uses_default_streams(): + return None + if count < 1: + raise ValueError("count must be >= 1") + self._prepare_for_stream_build() + stream_configs = self._stream_configs(self._source_config) + self.dynamic_streams + return self._create_streams(stream_configs[:count], stream_configs) + + def _apply_stream_groups( + self, + streams: List[AbstractStream], + top_level_stream_names: Optional[Set[str]] = None, + ) -> None: """Set block_simultaneous_read on streams based on the manifest's stream_groups config. Iterates over the resolved manifest's stream_groups and matches group membership @@ -494,10 +555,15 @@ def _apply_stream_groups(self, streams: List[AbstractStream]) -> None: # Validate no stream shares a group with any of its ancestor streams stream_name_to_instance: Dict[str, AbstractStream] = {s.name: s for s in streams} - def _collect_all_ancestor_names(stream_name: str) -> Set[str]: + def _collect_all_ancestor_names( + stream_name: str, inline_instance: Optional[AbstractStream] = None + ) -> Set[str]: """Recursively collect all ancestor stream names.""" ancestors: Set[str] = set() inst = stream_name_to_instance.get(stream_name) + if inst is None and stream_name in (top_level_stream_names or set()): + # A top-level stream that this call did not build: walk the copy inlined in the child + inst = inline_instance if not isinstance(inst, DefaultStream): return ancestors partition_router = inst.get_partition_router() @@ -514,7 +580,9 @@ def _collect_all_ancestor_names(stream_name: str) -> Set[str]: 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)) + ancestors.update( + _collect_all_ancestor_names(parent_name, parent_config.stream) + ) return ancestors for stream in streams: @@ -534,9 +602,37 @@ def _collect_all_ancestor_names(stream_name: str) -> Set[str]: if isinstance(stream, DefaultStream) and stream.name in stream_name_to_group: stream.block_simultaneous_read = stream_name_to_group[stream.name] + @staticmethod + def _parent_stream_names(stream_configs: List[Dict[str, Any]]) -> Set[str]: + parent_stream_names: Set[str] = set() + for stream_config in stream_configs: + parent_stream = stream_config.get("incremental_sync", {}).get("parent_stream") + if parent_stream: + parent_stream_names.add(parent_stream["name"]) + elif stream_config.get("retriever", {}).get("partition_router", {}): + partition_router = stream_config["retriever"]["partition_router"] + routers = ( + list(partition_router) + if isinstance(partition_router, list) + else [partition_router] + ) + while routers: + router = routers.pop() + if not isinstance(router, dict): + continue + parent_stream_names.update( + parent_config["stream"]["name"] + for parent_config in router.get("parent_stream_configs") or [] + ) + routers.extend(router.get("partition_routers") or []) + if router.get("underlying_partition_router"): + routers.append(router["underlying_partition_router"]) + return parent_stream_names + @staticmethod def _initialize_cache_for_parent_streams( stream_configs: List[Dict[str, Any]], + all_stream_configs: Optional[List[Dict[str, Any]]] = None, ) -> List[Dict[str, Any]]: """Enable caching for parent streams unless explicitly disabled. @@ -545,7 +641,10 @@ def _initialize_cache_for_parent_streams( settings are respected for streams that cannot use caching (e.g., scroll-based pagination APIs where caching causes duplicate records). """ - parent_streams = set() + parent_streams = ConcurrentDeclarativeSource._parent_stream_names(stream_configs) + parent_streams.update( + ConcurrentDeclarativeSource._parent_stream_names(all_stream_configs or []) + ) def _set_cache_if_not_disabled(requester: Dict[str, Any]) -> None: """Set use_cache to True only if not explicitly disabled.""" @@ -556,7 +655,6 @@ def update_with_cache_parent_configs( parent_configs: list[dict[str, Any]], ) -> None: for parent_config in parent_configs: - parent_streams.add(parent_config["stream"]["name"]) if parent_config["stream"]["type"] == "StateDelegatingStream": _set_cache_if_not_disabled( parent_config["stream"]["full_refresh_stream"]["retriever"]["requester"] @@ -569,7 +667,6 @@ def update_with_cache_parent_configs( for stream_config in stream_configs: if stream_config.get("incremental_sync", {}).get("parent_stream"): - parent_streams.add(stream_config["incremental_sync"]["parent_stream"]["name"]) _set_cache_if_not_disabled( stream_config["incremental_sync"]["parent_stream"]["retriever"]["requester"] ) diff --git a/unit_tests/sources/declarative/checks/test_check_dynamic_stream.py b/unit_tests/sources/declarative/checks/test_check_dynamic_stream.py index 6a41d945cd..31618008ab 100644 --- a/unit_tests/sources/declarative/checks/test_check_dynamic_stream.py +++ b/unit_tests/sources/declarative/checks/test_check_dynamic_stream.py @@ -5,10 +5,12 @@ import json import logging from copy import deepcopy +from unittest.mock import MagicMock import pytest from airbyte_cdk.models import Status +from airbyte_cdk.sources.declarative.checks.check_dynamic_stream import CheckDynamicStream from airbyte_cdk.sources.declarative.concurrent_declarative_source import ( ConcurrentDeclarativeSource, ) @@ -168,3 +170,157 @@ def test_check_dynamic_stream( assert connection_status.status == available_expectation for message in expected_messages: assert message in connection_status.message + + +def _static_stream(name: str) -> dict: + return { + "type": "DeclarativeStream", + "name": name, + "primary_key": [], + "schema_loader": { + "type": "InlineSchemaLoader", + "schema": {"type": "object", "properties": {}}, + }, + "retriever": { + "type": "SimpleRetriever", + "requester": { + "type": "HttpRequester", + "url_base": "https://api.test.com", + "path": f"/{name}", + "http_method": "GET", + }, + "record_selector": { + "type": "RecordSelector", + "extractor": {"type": "DpathExtractor", "field_path": []}, + }, + "paginator": {"type": "NoPagination"}, + }, + } + + +@pytest.mark.parametrize( + "stream_count, expected_stream_names", + [ + pytest.param(1, ["first_static"], id="one-stream"), + pytest.param(3, ["first_static", "second_static", "item_1"], id="three-streams"), + ], +) +def test_check_dynamic_stream_builds_only_the_first_stream_count_in_order( + mocker, stream_count, expected_stream_names +): + manifest = deepcopy(_MANIFEST) + manifest["streams"] = [_static_stream("first_static"), _static_stream("second_static")] + manifest["check"]["stream_count"] = stream_count + source = ConcurrentDeclarativeSource( + source_config=manifest, + config={**_CONFIG, "api_key": "test"}, + catalog=None, + state=None, + ) + create_component_spy = mocker.spy(source._constructor, "create_component") + + with HttpMocker() as http_mocker: + http_mocker.get( + HttpRequest(url="https://api.test.com/items"), + HttpResponse( + body=json.dumps([{"id": 1, "name": "item_1"}, {"id": 2, "name": "item_2"}]) + ), + ) + for stream_name in ("first_static", "second_static"): + http_mocker.get( + HttpRequest(url=f"https://api.test.com/{stream_name}"), + HttpResponse(body=json.dumps([])), + ) + http_mocker.get( + HttpRequest(url="https://api.test.com/items/1"), + HttpResponse(body=json.dumps([])), + ) + + assert source.check(logger, source._config).status == Status.SUCCEEDED + + created_stream_names = [ + component_call.args[1]["name"] + for component_call in create_component_spy.call_args_list + if len(component_call.args) > 1 + and getattr(component_call.args[0], "__name__", "") + in {"DeclarativeStream", "StateDelegatingStream"} + ] + assert created_stream_names == expected_stream_names + + +def test_check_dynamic_stream_uses_streams_override_for_concurrent_source(): + class StreamsOverrideSource(ConcurrentDeclarativeSource): + def __init__(self, *args, **kwargs): + self.streams_call_count = 0 + super().__init__(*args, **kwargs) + + def streams(self, config): + self.streams_call_count += 1 + if config.get("end_date") == "": + config.pop("end_date") + return super().streams(config=config) + + stream = _static_stream("selected") + stream["incremental_sync"] = { + "type": "DatetimeBasedCursor", + "start_datetime": "2024-01-01", + "end_datetime": "{{ config.get('end_date', today_utc()) }}", + "datetime_format": "%Y-%m-%d", + "cursor_datetime_formats": ["%Y-%m-%d"], + "cursor_field": "updated_at", + } + source = StreamsOverrideSource( + source_config={ + "version": "6.7.0", + "type": "DeclarativeSource", + "streams": [stream], + "check": {"type": "CheckDynamicStream", "stream_count": 1}, + }, + config={"end_date": ""}, + catalog=None, + state=None, + ) + + with HttpMocker() as http_mocker: + http_mocker.get( + HttpRequest(url="https://api.test.com/selected"), + HttpResponse(body=json.dumps([])), + ) + assert source.check(logger, source._config).status == Status.SUCCEEDED + + assert source.streams_call_count == 1 + + +@pytest.mark.parametrize( + "stream_count,use_check_availability", + [(0, True), (1, False)], +) +def test_check_dynamic_stream_uses_streams_for_non_positive_count_or_disabled_availability( + stream_count, use_check_availability +): + source = MagicMock(spec=ConcurrentDeclarativeSource) + source.streams.return_value = [] + checker = CheckDynamicStream( + stream_count=stream_count, + parameters={}, + use_check_availability=use_check_availability, + ) + + checker.check_connection(source, logger, {}) + + source.streams.assert_called_once_with(config={}) + source._first_streams.assert_not_called() + + +def test_check_dynamic_stream_uses_streams_for_legacy_manifest_source(): + from airbyte_cdk.legacy.sources.declarative.manifest_declarative_source import ( + ManifestDeclarativeSource, + ) + + source = MagicMock(spec=ManifestDeclarativeSource) + source.streams.return_value = [] + checker = CheckDynamicStream(stream_count=1, parameters={}) + + checker.check_connection(source, logger, {}) + + source.streams.assert_called_once_with(config={}) diff --git a/unit_tests/sources/declarative/checks/test_check_stream.py b/unit_tests/sources/declarative/checks/test_check_stream.py index 8bc9571d1a..a05de813df 100644 --- a/unit_tests/sources/declarative/checks/test_check_stream.py +++ b/unit_tests/sources/declarative/checks/test_check_stream.py @@ -5,7 +5,7 @@ import json import logging from copy import deepcopy -from typing import Any, Iterable, Mapping, Optional +from typing import Any, Iterable, Mapping, Optional, cast from unittest.mock import MagicMock import pytest @@ -907,3 +907,587 @@ def test_check_stream_only_type_provided(): ) with pytest.raises(ValueError): source.check(logger, _CONFIG) + + +def _stream_config(name: str, path: Optional[str] = None) -> dict: + return { + "type": "DeclarativeStream", + "name": name, + "primary_key": [], + "schema_loader": { + "type": "InlineSchemaLoader", + "schema": {"type": "object", "properties": {}}, + }, + "retriever": { + "type": "SimpleRetriever", + "requester": { + "type": "HttpRequester", + "url_base": "https://api.test.com", + "path": path or f"/{name}", + "http_method": "GET", + }, + "record_selector": { + "type": "RecordSelector", + "extractor": {"type": "DpathExtractor", "field_path": []}, + }, + "paginator": {"type": "NoPagination"}, + }, + } + + +def _incremental_substream(name: str, parent_stream: Optional[dict] = None) -> dict: + stream = _stream_config(name) + stream["incremental_sync"] = { + "type": "DatetimeBasedCursor", + "start_datetime": "2024-01-01", + "end_datetime": "2024-01-02", + "datetime_format": "%Y-%m-%d", + "cursor_datetime_formats": ["%Y-%m-%d"], + "cursor_field": "updated_at", + } + if parent_stream is not None: + stream["retriever"]["partition_router"] = { + "type": "SubstreamPartitionRouter", + "parent_stream_configs": [ + { + "type": "ParentStreamConfig", + "stream": deepcopy(parent_stream), + "parent_key": "id", + "partition_field": f"{name}_parent_id", + } + ], + } + return stream + + +def _source_with_streams( + stream_configs: list[dict], + stream_names: list[str], + config: Optional[dict] = None, +) -> ConcurrentDeclarativeSource: + return ConcurrentDeclarativeSource( + source_config={ + "version": "6.7.0", + "type": "DeclarativeSource", + "streams": deepcopy(stream_configs), + "check": {"type": "CheckStream", "stream_names": stream_names}, + }, + config=config or {}, + catalog=None, + state=None, + ) + + +def _created_stream_names(create_component_spy) -> list[str]: + return [ + call.args[1]["name"] + for call in create_component_spy.call_args_list + if len(call.args) > 1 + and getattr(call.args[0], "__name__", "") in {"DeclarativeStream", "StateDelegatingStream"} + ] + + +def test_check_builds_only_the_named_static_stream(mocker): + source = _source_with_streams( + [_stream_config("first"), _stream_config("selected"), _stream_config("third")], + ["selected"], + ) + create_component_spy = mocker.spy(source._constructor, "create_component") + + with HttpMocker() as http_mocker: + http_mocker.get( + HttpRequest(url="https://api.test.com/selected"), + HttpResponse(body=json.dumps([{"id": 1}])), + ) + assert source.check(logger, {}).status == Status.SUCCEEDED + + assert _created_stream_names(create_component_spy) == ["selected"] + + +def test_check_does_not_build_an_unread_stream_that_fails_to_build(mocker): + streams = [_stream_config("checked"), _stream_config("unread_bad")] + source = _source_with_streams(streams, ["checked"]) + + def patch_bad_stream(source): + original_create_component = source._constructor.create_component + + def create_component(model, component_config, *args, **kwargs): + if ( + getattr(model, "__name__", "") in {"DeclarativeStream", "StateDelegatingStream"} + and component_config["name"] == "unread_bad" + ): + raise ValueError("bad unread stream") + return original_create_component(model, component_config, *args, **kwargs) + + mocker.patch.object(source._constructor, "create_component", side_effect=create_component) + + patch_bad_stream(source) + with HttpMocker() as http_mocker: + http_mocker.get( + HttpRequest(url="https://api.test.com/checked"), + HttpResponse(body=json.dumps([{"id": 1}])), + ) + assert source.check(logger, {}).status == Status.SUCCEEDED + + failing_source = _source_with_streams(streams, ["unread_bad"]) + patch_bad_stream(failing_source) + status = failing_source.check(logger, {}) + assert status.status == Status.FAILED + assert "Encountered an error while discovering streams" in status.message + + +@pytest.mark.parametrize("condition", [True, False]) +def test_check_stream_name_inside_conditional_streams(condition, mocker): + manifest = { + "version": "6.7.0", + "type": "DeclarativeSource", + "streams": [ + { + "type": "ConditionalStreams", + "condition": "{{ config['include_conditional'] }}", + "streams": [_stream_config("conditional_stream")], + }, + _stream_config("other_stream"), + ], + "check": {"type": "CheckStream", "stream_names": ["conditional_stream"]}, + } + source = ConcurrentDeclarativeSource( + source_config=manifest, + config={"include_conditional": condition}, + catalog=None, + state=None, + ) + + if condition: + create_component_spy = mocker.spy(source._constructor, "create_component") + with HttpMocker() as http_mocker: + http_mocker.get( + HttpRequest(url="https://api.test.com/conditional_stream"), + HttpResponse(body=json.dumps([{"id": 1}])), + ) + assert source.check(logger, source._config).status == Status.SUCCEEDED + assert _created_stream_names(create_component_spy) == ["conditional_stream"] + else: + with pytest.raises(ValueError, match="conditional_stream is not part of the catalog"): + source.check(logger, source._config) + + +def test_check_stream_names_can_be_overridden_from_config_for_concurrent_source(mocker): + source = _source_with_streams( + [_stream_config("default_stream"), _stream_config("selected_stream")], + ["default_stream"], + config={"__airbyte_check_stream_names": ["selected_stream"]}, + ) + create_component_spy = mocker.spy(source._constructor, "create_component") + + with HttpMocker() as http_mocker: + http_mocker.get( + HttpRequest(url="https://api.test.com/selected_stream"), + HttpResponse(body=json.dumps([{"id": 1}])), + ) + assert source.check(logger, source._config).status == Status.SUCCEEDED + + assert _created_stream_names(create_component_spy) == ["selected_stream"] + + +def test_check_with_misspelled_stream_name_uses_catalog_error(): + source = _source_with_streams( + [_stream_config("first"), _stream_config("second")], ["misspelled"] + ) + + with pytest.raises( + ValueError, + match=r"misspelled is not part of the catalog\. Expected one of \['first', 'second'\]\.", + ): + source.check(logger, {}) + + +def test_check_for_dynamic_stream_name_falls_back_to_streams(mocker): + manifest = deepcopy(_MANIFEST_WITHOUT_CHECK_COMPONENT) + manifest["dynamic_streams"] = [manifest["dynamic_streams"][1]] + manifest["check"] = {"type": "CheckStream", "stream_names": ["item_3"]} + source = ConcurrentDeclarativeSource( + source_config=manifest, + config={**_CONFIG, "api_key": "test"}, + catalog=None, + state=None, + ) + streams_spy = mocker.spy(source, "streams") + + with HttpMocker() as http_mocker: + http_mocker.get( + HttpRequest(url="https://api.test.com/static"), + HttpResponse(body=json.dumps([])), + ) + http_mocker.get( + HttpRequest(url="https://api.test.com/items/3"), + HttpResponse(body=json.dumps([])), + ) + assert source.check(logger, source._config).status == Status.SUCCEEDED + + streams_spy.assert_called_once() + + +def test_check_dynamic_stream_configs_are_evaluated_once_when_not_checked(mocker): + manifest = deepcopy(_MANIFEST_WITHOUT_CHECK_COMPONENT) + manifest["dynamic_streams"] = [manifest["dynamic_streams"][1]] + manifest["check"] = {"type": "CheckStream", "stream_names": ["static_stream"]} + source = ConcurrentDeclarativeSource( + source_config=manifest, + config={**_CONFIG, "api_key": "test"}, + catalog=None, + state=None, + ) + + with HttpMocker() as http_mocker: + http_mocker.get( + HttpRequest(url="https://api.test.com/static"), + HttpResponse(body=json.dumps([{"id": 1}])), + ) + dynamic_configs_spy = mocker.spy(source, "_dynamic_stream_configs") + assert source.check(logger, source._config).status == Status.SUCCEEDED + + dynamic_configs_spy.assert_called_once() + + +def test_check_builds_a_named_state_delegating_stream(mocker): + stream = _stream_config("state_stream") + incremental_stream = deepcopy(stream) + incremental_stream["incremental_sync"] = { + "type": "DatetimeBasedCursor", + "start_datetime": "2024-01-01T00:00:00Z", + "end_datetime": "2024-01-02T00:00:00Z", + "datetime_format": "%Y-%m-%dT%H:%M:%SZ", + "cursor_datetime_formats": ["%Y-%m-%dT%H:%M:%SZ"], + "cursor_field": "updated_at", + } + state_delegating_stream = { + "type": "StateDelegatingStream", + "name": "state_stream", + "full_refresh_stream": stream, + "incremental_stream": incremental_stream, + } + source = _source_with_streams([state_delegating_stream], ["state_stream"]) + create_component_spy = mocker.spy(source._constructor, "create_component") + + streams = source._streams_by_name(["state_stream"]) + + assert streams is not None + assert [stream.name for stream in streams] == ["state_stream"] + assert _created_stream_names(create_component_spy) == ["state_stream"] + assert any( + component_call.args + and getattr(component_call.args[0], "__name__", "") == "StateDelegatingStream" + for component_call in create_component_spy.call_args_list + ) + + +@pytest.mark.parametrize( + "check_component", + [ + pytest.param( + {"type": "CheckStream", "stream_names": ["selected"]}, + id="check-stream", + ), + pytest.param( + {"type": "CheckDynamicStream", "stream_count": 1}, + id="check-dynamic-stream", + ), + ], +) +def test_check_still_fails_when_spec_config_validations_fail(check_component): + source_config = { + "version": "6.7.0", + "type": "DeclarativeSource", + "streams": [_stream_config("selected")], + "check": check_component, + "spec": { + "type": "Spec", + "connection_specification": { + "$schema": "http://json-schema.org/draft-07/schema#", + "type": "object", + "properties": {"settings": {"type": "object"}}, + }, + "config_normalization_rules": { + "type": "ConfigNormalizationRules", + "validations": [ + { + "type": "DpathValidator", + "field_path": ["settings"], + "validation_strategy": { + "type": "ValidateAdheresToSchema", + "base_schema": { + "type": "object", + "properties": {"mode": {"type": "string", "enum": ["sync"]}}, + "required": ["mode"], + "additionalProperties": False, + }, + }, + } + ], + }, + }, + } + source = ConcurrentDeclarativeSource( + source_config=source_config, + config={"settings": {"mode": "invalid"}}, + catalog=None, + state=None, + ) + + with HttpMocker(): + if check_component["type"] == "CheckDynamicStream": + with pytest.raises(ValueError, match="JSON schema validation error"): + source.check(logger, source._config) + return + status = source.check(logger, source._config) + + assert status.status == Status.FAILED + assert "Encountered an error while discovering streams" in status.message + + +def test_check_with_dynamic_stream_config_still_builds_unread_static_streams(mocker): + source = ConcurrentDeclarativeSource( + source_config={ + "version": "6.7.0", + "type": "DeclarativeSource", + "streams": [_stream_config("selected"), _stream_config("unread_bad")], + "check": { + "type": "CheckStream", + "stream_names": ["selected"], + "dynamic_streams_check_configs": [ + { + "type": "DynamicStreamCheckConfig", + "dynamic_stream_name": "configured_dynamic_stream", + } + ], + }, + }, + config={}, + catalog=None, + state=None, + ) + original_create_component = source._constructor.create_component + + def create_component(model, component_config, *args, **kwargs): + if ( + getattr(model, "__name__", "") in {"DeclarativeStream", "StateDelegatingStream"} + and component_config["name"] == "unread_bad" + ): + raise ValueError("bad unread stream") + return original_create_component(model, component_config, *args, **kwargs) + + mocker.patch.object(source._constructor, "create_component", side_effect=create_component) + + with HttpMocker() as http_mocker: + http_mocker.get( + HttpRequest(url="https://api.test.com/selected"), + HttpResponse(body=json.dumps([])), + ) + status = source.check(logger, source._config) + + assert status.status == Status.FAILED + assert "Encountered an error while discovering streams" in status.message + + +def test_check_fallback_prepares_streams_once_for_unknown_name(mocker): + source = _source_with_streams([_stream_config("selected")], ["unknown"]) + prepare_spy = mocker.spy(source, "_prepare_for_stream_build") + + with pytest.raises(ValueError, match="unknown is not part of the catalog"): + source.check(logger, {}) + + prepare_spy.assert_called_once() + + +@pytest.mark.parametrize( + "stream_names", + [ + pytest.param([], id="empty-names"), + pytest.param(["selected", "unknown"], id="mixed-known-and-unknown"), + ], +) +def test_streams_by_name_returns_none_for_empty_or_non_static_names(stream_names): + source = _source_with_streams([_stream_config("selected")], ["selected"]) + + assert source._streams_by_name(stream_names) is None + + +def test_streams_by_name_enables_cache_for_checked_parent_without_mutating_manifest(): + parent = _stream_config("parent") + child = _incremental_substream("child", parent) + source = _source_with_streams([parent, child], ["parent"]) + manifest_before = deepcopy(source.resolved_manifest) + + streams = source._streams_by_name(["parent"]) + + assert streams is not None + assert [stream.name for stream in streams] == ["parent"] + parent_stream = streams[0] + assert ( + parent_stream._stream_partition_generator._partition_factory._retriever.requester.use_cache + ) + assert source.resolved_manifest == manifest_before + + +def test_check_still_fails_when_a_dynamic_stream_cannot_be_built(mocker): + manifest = deepcopy(_MANIFEST_WITHOUT_CHECK_COMPONENT) + manifest["dynamic_streams"] = [manifest["dynamic_streams"][1]] + manifest["check"] = {"type": "CheckStream", "stream_names": ["static_stream"]} + source = ConcurrentDeclarativeSource( + source_config=manifest, + config={**_CONFIG, "api_key": "test"}, + catalog=None, + state=None, + ) + original_create_component = source._constructor.create_component + + def create_component(*args, **kwargs): + model = kwargs.get("model_type", args[0] if args else None) + component_config = kwargs.get("component_definition", args[1] if len(args) > 1 else {}) + if ( + getattr(model, "__name__", "") in {"DeclarativeStream", "StateDelegatingStream"} + and component_config["name"] == "item_3" + ): + raise ValueError("bad dynamic stream") + return original_create_component(*args, **kwargs) + + mocker.patch.object(source._constructor, "create_component", side_effect=create_component) + + with HttpMocker() as http_mocker: + http_mocker.get( + HttpRequest(url="https://api.test.com/static"), + HttpResponse(body=json.dumps([])), + ) + status = source.check(logger, source._config) + + assert status.status == Status.FAILED + assert "Encountered an error while discovering streams" in status.message + + +def _stream_group_manifest(streams: list[dict], check_stream_name: str) -> dict: + streams_by_name = {stream["name"]: stream for stream in streams} + return { + "version": "6.7.0", + "type": "DeclarativeSource", + "streams": streams, + "check": {"type": "CheckStream", "stream_names": [check_stream_name]}, + "stream_groups": { + "g1": { + "streams": [deepcopy(streams_by_name[name]) for name in ("A", "C")], + "action": {"type": "BlockSimultaneousSyncsAction"}, + } + }, + } + + +def test_check_detects_group_conflict_with_a_checked_streams_grandparent(): + stream_a = _incremental_substream("A") + stream_b = _incremental_substream("B", stream_a) + stream_c = _incremental_substream("C", stream_b) + source = ConcurrentDeclarativeSource( + source_config=_stream_group_manifest([stream_a, stream_b, stream_c], "C"), + config={}, + catalog=None, + state=None, + ) + + status = source.check(logger, {}) + + assert status.status == Status.FAILED + assert "Stream 'C' and its parent stream 'A' are both in group" in status.message + + +def test_stream_groups_ignore_inline_non_top_level_grandparents(): + stream_a = _incremental_substream("A") + stream_b = _incremental_substream("B", stream_a) + stream_c = _incremental_substream("C", stream_b) + source = ConcurrentDeclarativeSource( + source_config=_stream_group_manifest([stream_a, stream_c], "C"), + config={}, + catalog=None, + state=None, + ) + + assert [stream.name for stream in source.streams(config={})] == ["A", "C"] + selected_streams = source._streams_by_name(["C"]) + + assert selected_streams is not None + assert [stream.name for stream in selected_streams] == ["C"] + + +def test_check_does_not_build_unread_grouped_child(mocker): + stream_a = _incremental_substream("A") + stream_c = _incremental_substream("C", _incremental_substream("B", stream_a)) + source = ConcurrentDeclarativeSource( + source_config=_stream_group_manifest([stream_a, stream_c], "A"), + config={}, + catalog=None, + state=None, + ) + mocker.patch( + "airbyte_cdk.sources.declarative.checks.check_stream.evaluate_availability", + return_value=(True, None), + ) + + assert source.check(logger, {}).status == Status.SUCCEEDED + + +def test_check_uses_streams_override_for_concurrent_source(): + class StreamsOverrideSource(ConcurrentDeclarativeSource): + def __init__(self, *args, **kwargs): + self.streams_call_count = 0 + super().__init__(*args, **kwargs) + + def streams(self, config: Mapping[str, Any]): + self.streams_call_count += 1 + mutable_config = cast(dict[str, Any], config) + if mutable_config.get("end_date") == "": + mutable_config.pop("end_date") + return super().streams(config=config) + + stream = _stream_config("selected") + stream["incremental_sync"] = { + "type": "DatetimeBasedCursor", + "start_datetime": "2024-01-01", + "end_datetime": "{{ config.get('end_date', today_utc()) }}", + "datetime_format": "%Y-%m-%d", + "cursor_datetime_formats": ["%Y-%m-%d"], + "cursor_field": "updated_at", + } + source = StreamsOverrideSource( + source_config={ + "version": "6.7.0", + "type": "DeclarativeSource", + "streams": [stream], + "check": {"type": "CheckStream", "stream_names": ["selected"]}, + }, + config={"end_date": ""}, + catalog=None, + state=None, + ) + + with HttpMocker() as http_mocker: + http_mocker.get( + HttpRequest(url="https://api.test.com/selected"), + HttpResponse(body=json.dumps([])), + ) + assert source.check(logger, source._config).status == Status.SUCCEEDED + + assert source.streams_call_count == 1 + + +def test_check_stream_uses_streams_for_legacy_manifest_source(): + from airbyte_cdk.legacy.sources.declarative.manifest_declarative_source import ( + ManifestDeclarativeSource, + ) + + source = MagicMock(spec=ManifestDeclarativeSource) + stream = MagicMock(spec=Stream) + stream.name = "selected" + source.streams.return_value = [stream] + checker = CheckStream(stream_names=["selected"], parameters={}) + + assert checker._build_streams(source, {}, ["selected"]) == [stream] + + source.streams.assert_called_once_with(config={}) diff --git a/unit_tests/sources/declarative/test_concurrent_declarative_source.py b/unit_tests/sources/declarative/test_concurrent_declarative_source.py index dc82ca4d2b..327bfa0e3f 100644 --- a/unit_tests/sources/declarative/test_concurrent_declarative_source.py +++ b/unit_tests/sources/declarative/test_concurrent_declarative_source.py @@ -6649,13 +6649,16 @@ def test_get_partition_router(stream_factory, expected_type): assert isinstance(router, GroupingPartitionRouter) -def test_api_budget_is_set_before_dynamic_streams_evaluated(): - """Verify that set_api_budget is called before dynamic_streams is accessed in streams(). - - This is a regression test for https://github.com/airbytehq/oncall/issues/11954 - where dynamic stream discovery HTTP requests bypassed the configured rate limiter - because set_api_budget was called after self.dynamic_streams was evaluated. - """ +@pytest.mark.parametrize( + "build_path", + [ + pytest.param("streams", id="streams"), + pytest.param("_streams_by_name", id="streams_by_name"), + pytest.param("_first_streams", id="first_streams"), + ], +) +def test_api_budget_is_set_before_dynamic_streams_evaluated(build_path): + """Verify all stream-building paths set the API budget before dynamic discovery.""" source = ConcurrentDeclarativeSource( source_config=_MANIFEST, config=_CONFIG, catalog=None, state=None ) @@ -6673,7 +6676,6 @@ def tracking_dynamic_stream_configs(*args, **kwargs): call_order.append("dynamic_stream_configs") return original_dynamic_stream_configs(*args, **kwargs) - # Add an api_budget to the source config so set_api_budget is actually called source._source_config["api_budget"] = { "type": "HTTPAPIBudget", "policies": [ @@ -6691,7 +6693,13 @@ def tracking_dynamic_stream_configs(*args, **kwargs): source, "_dynamic_stream_configs", side_effect=tracking_dynamic_stream_configs ), ): - source.streams(config=_CONFIG) + if build_path == "streams": + source.streams(config=_CONFIG) + elif build_path == "_streams_by_name": + stream_name = source._stream_configs(source._source_config)[0]["name"] + source._streams_by_name([stream_name]) + else: + source._first_streams(1) assert "set_api_budget" in call_order, "set_api_budget was never called" assert "dynamic_stream_configs" in call_order, "dynamic_stream_configs was never called" @@ -6700,6 +6708,131 @@ def tracking_dynamic_stream_configs(*args, **kwargs): ) +@pytest.mark.parametrize("count", [0, -1]) +def test_first_streams_rejects_non_positive_count(count): + source = ConcurrentDeclarativeSource( + source_config=_MANIFEST, + config=_CONFIG, + catalog=None, + state=None, + ) + + with pytest.raises(ValueError, match="count must be >= 1"): + source._first_streams(count) + + +def test_initialize_cache_for_parent_streams_reads_all_configs_without_mutating_them(): + def stream_config(name: str) -> dict: + return { + "type": "DeclarativeStream", + "name": name, + "retriever": {"requester": {}}, + } + + partition_parent = stream_config("partition_parent") + partition_child = stream_config("partition_child") + partition_child["retriever"]["partition_router"] = [ + { + "type": "UnionPartitionRouter", + "partition_routers": [ + { + "type": "GroupingPartitionRouter", + "underlying_partition_router": { + "type": "SubstreamPartitionRouter", + "parent_stream_configs": [ + { + "stream": partition_parent, + "parent_key": "id", + "partition_field": "parent_id", + } + ], + }, + } + ], + } + ] + incremental_child = stream_config("incremental_child") + incremental_child["incremental_sync"] = { + "parent_stream": { + "name": "incremental_parent", + "retriever": {"requester": {}}, + } + } + all_stream_configs = [partition_parent, partition_child, incremental_child] + unchanged_configs = deepcopy(all_stream_configs) + selected_configs = deepcopy([partition_parent]) + + initialized = ConcurrentDeclarativeSource._initialize_cache_for_parent_streams( + selected_configs, all_stream_configs=all_stream_configs + ) + + assert initialized[0]["retriever"]["requester"]["use_cache"] is True + assert all_stream_configs == unchanged_configs + assert ConcurrentDeclarativeSource._parent_stream_names(all_stream_configs) == { + "partition_parent", + "incremental_parent", + } + + +def test_streams_by_name_builds_a_checked_substream_without_returning_its_parent(): + def stream_config(name: str, path: str) -> dict: + return { + "type": "DeclarativeStream", + "name": name, + "primary_key": [], + "schema_loader": { + "type": "InlineSchemaLoader", + "schema": {"type": "object", "properties": {}}, + }, + "retriever": { + "type": "SimpleRetriever", + "requester": { + "type": "HttpRequester", + "url_base": "https://api.test.com", + "path": path, + "http_method": "GET", + }, + "record_selector": { + "type": "RecordSelector", + "extractor": {"type": "DpathExtractor", "field_path": []}, + }, + "paginator": {"type": "NoPagination"}, + }, + } + + parent = stream_config("parent", "/parent") + child = stream_config("child", "/child") + child["retriever"]["partition_router"] = { + "type": "SubstreamPartitionRouter", + "parent_stream_configs": [ + { + "stream": deepcopy(parent), + "parent_key": "id", + "partition_field": "parent_id", + } + ], + } + source = ConcurrentDeclarativeSource( + source_config={ + "version": "6.7.0", + "type": "DeclarativeSource", + "check": {"type": "CheckStream", "stream_names": ["child"]}, + "streams": [parent, child], + }, + config={}, + catalog=None, + state=None, + ) + + streams = source._streams_by_name(["child"]) + + assert streams is not None + assert [stream.name for stream in streams] == ["child"] + child_stream = streams[0] + parent_streams = child_stream._stream_partition_generator._stream_slicer.parent_stream_configs + assert [parent_config.stream.name for parent_config in parent_streams] == ["parent"] + + def test_dynamic_stream_discovery_http_requests_use_api_budget(): """Verify that HttpComponentsResolver's requester receives the configured api_budget.