diff --git a/airbyte_cdk/legacy/sources/declarative/manifest_declarative_source.py b/airbyte_cdk/legacy/sources/declarative/manifest_declarative_source.py index 818ce0179d..dea33879ac 100644 --- a/airbyte_cdk/legacy/sources/declarative/manifest_declarative_source.py +++ b/airbyte_cdk/legacy/sources/declarative/manifest_declarative_source.py @@ -314,14 +314,14 @@ def streams(self, config: Mapping[str, Any]) -> List[Union[Stream, AbstractStrea } ) - stream_configs = ( - self._stream_configs(self._source_config, config=config) + self.dynamic_streams - ) - api_budget_model = self._source_config.get("api_budget") if api_budget_model: self._constructor.set_api_budget(api_budget_model, config) + stream_configs = ( + self._stream_configs(self._source_config, config=config) + self.dynamic_streams + ) + source_streams = [ self._constructor.create_component( ( diff --git a/airbyte_cdk/sources/declarative/concurrent_declarative_source.py b/airbyte_cdk/sources/declarative/concurrent_declarative_source.py index 2926156924..fb9c0ea0bb 100644 --- a/airbyte_cdk/sources/declarative/concurrent_declarative_source.py +++ b/airbyte_cdk/sources/declarative/concurrent_declarative_source.py @@ -406,12 +406,12 @@ def streams(self, config: Mapping[str, Any]) -> List[AbstractStream]: # type: i if self._spec_component: self._spec_component.validate_config(self._config) - stream_configs = self._stream_configs(self._source_config) + self.dynamic_streams - api_budget_model = self._source_config.get("api_budget") 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)) source_streams = [ diff --git a/unit_tests/legacy/sources/declarative/test_manifest_declarative_source.py b/unit_tests/legacy/sources/declarative/test_manifest_declarative_source.py index 8bc130a35f..b565874323 100644 --- a/unit_tests/legacy/sources/declarative/test_manifest_declarative_source.py +++ b/unit_tests/legacy/sources/declarative/test_manifest_declarative_source.py @@ -37,6 +37,9 @@ from airbyte_cdk.sources.declarative.parsers.model_to_component_factory import ( ModelToComponentFactory, ) +from airbyte_cdk.sources.declarative.resolvers.http_components_resolver import ( + HttpComponentsResolver, +) from airbyte_cdk.sources.declarative.retrievers.simple_retriever import SimpleRetriever from airbyte_cdk.sources.streams.concurrent.default_stream import DefaultStream from unit_tests.sources.declarative.parsers.test_model_to_component_factory import get_retriever @@ -2348,3 +2351,192 @@ def test_given_invalid_config_streams_validates_config_and_raises(): with pytest.raises(ValueError): source.streams(input_config) + + +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. + """ + manifest = { + "version": "3.8.2", + "type": "DeclarativeSource", + "check": {"type": "CheckStream", "stream_names": ["lists"]}, + "streams": [ + { + "type": "DeclarativeStream", + "name": "lists", + "schema_loader": { + "type": "InlineSchemaLoader", + "schema": {"type": "object"}, + }, + "retriever": { + "type": "SimpleRetriever", + "requester": { + "type": "HttpRequester", + "url_base": "https://example.org", + "path": "/test", + "authenticator": {"type": "NoAuth"}, + }, + "record_selector": { + "type": "RecordSelector", + "extractor": {"type": "DpathExtractor", "field_path": []}, + }, + }, + } + ], + } + + source = ManifestDeclarativeSource(source_config=manifest) + + call_order: list[str] = [] + original_set_api_budget = source._constructor.set_api_budget + + def tracking_set_api_budget(*args, **kwargs): + call_order.append("set_api_budget") + return original_set_api_budget(*args, **kwargs) + + original_dynamic_stream_configs = source._dynamic_stream_configs + + 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": [ + { + "type": "MovingWindowCallRatePolicy", + "rates": [{"type": "Rate", "limit": 5, "interval": "PT1S"}], + "matchers": [], + } + ], + } + + with ( + patch.object(source._constructor, "set_api_budget", side_effect=tracking_set_api_budget), + patch.object( + source, "_dynamic_stream_configs", side_effect=tracking_dynamic_stream_configs + ), + ): + source.streams(config={}) + + 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" + assert call_order.index("set_api_budget") < call_order.index("dynamic_stream_configs"), ( + f"set_api_budget must be called before dynamic_stream_configs, but call order was: {call_order}" + ) + + +def test_dynamic_stream_discovery_http_requests_use_api_budget(): + """Verify that HttpComponentsResolver's requester receives the configured api_budget. + + Regression test for https://github.com/airbytehq/oncall/issues/11954 + The discovery HTTP requests made by HttpComponentsResolver must be rate-limited + by the api_budget configured in the manifest. + """ + manifest = { + "version": "3.8.2", + "type": "DeclarativeSource", + "check": {"type": "CheckStream", "stream_names": []}, + "dynamic_streams": [ + { + "type": "DynamicDeclarativeStream", + "stream_template": { + "type": "DeclarativeStream", + "$parameters": { + "name": "dynamic_items", + "primary_key": "id", + "url_base": "https://api.test.com", + }, + "schema_loader": { + "type": "InlineSchemaLoader", + "schema": { + "$schema": "https://json-schema.org/draft-07/schema#", + "type": "object", + "properties": {"id": {"type": "string"}}, + }, + }, + "retriever": { + "type": "SimpleRetriever", + "record_selector": { + "type": "RecordSelector", + "extractor": {"type": "DpathExtractor", "field_path": []}, + }, + "paginator": {"type": "NoPagination"}, + "requester": { + "type": "HttpRequester", + "url_base": "https://api.test.com", + "path": "/items", + "http_method": "GET", + "authenticator": {"type": "NoAuth"}, + }, + }, + }, + "components_resolver": { + "type": "HttpComponentsResolver", + "$parameters": { + "name": "resolver", + "primary_key": "id", + "url_base": "https://api.test.com", + }, + "retriever": { + "type": "SimpleRetriever", + "record_selector": { + "type": "RecordSelector", + "extractor": {"type": "DpathExtractor", "field_path": []}, + }, + "paginator": {"type": "NoPagination"}, + "requester": { + "type": "HttpRequester", + "url_base": "https://api.test.com", + "path": "/components", + "http_method": "GET", + "authenticator": {"type": "NoAuth"}, + }, + }, + "components_mapping": [ + { + "type": "ComponentMappingDefinition", + "field_path": ["name"], + "value": "{{ components_values.name }}", + } + ], + }, + } + ], + "api_budget": { + "type": "HTTPAPIBudget", + "policies": [ + { + "type": "MovingWindowCallRatePolicy", + "rates": [{"type": "Rate", "limit": 5, "interval": "PT1S"}], + "matchers": [], + } + ], + }, + } + + source = ManifestDeclarativeSource(source_config=manifest) + + captured_resolvers: list[Any] = [] + + def capturing_resolve(resolver_self, *args, **kwargs): + captured_resolvers.append(resolver_self) + return iter([]) + + with patch.object(HttpComponentsResolver, "resolve_components", capturing_resolve): + source.streams(config={}) + + assert len(captured_resolvers) == 1, ( + f"Expected exactly one HttpComponentsResolver, got {len(captured_resolvers)}" + ) + resolver = captured_resolvers[0] + requester = resolver.retriever.requester + assert requester.api_budget is not None, ( + "HttpComponentsResolver's requester should have api_budget set during dynamic stream " + "discovery, but it was None. This means discovery HTTP requests are not rate-limited." + ) diff --git a/unit_tests/sources/declarative/test_concurrent_declarative_source.py b/unit_tests/sources/declarative/test_concurrent_declarative_source.py index bf1f61610a..5ffda84156 100644 --- a/unit_tests/sources/declarative/test_concurrent_declarative_source.py +++ b/unit_tests/sources/declarative/test_concurrent_declarative_source.py @@ -57,6 +57,9 @@ ClientSideIncrementalRecordFilterDecorator, ) from airbyte_cdk.sources.declarative.partition_routers import AsyncJobPartitionRouter +from airbyte_cdk.sources.declarative.resolvers.http_components_resolver import ( + HttpComponentsResolver, +) from airbyte_cdk.sources.declarative.retrievers.simple_retriever import SimpleRetriever from airbyte_cdk.sources.declarative.stream_slicers.declarative_partition_generator import ( StreamSlicerPartitionGenerator, @@ -5518,3 +5521,166 @@ def test_get_partition_router(stream_factory, expected_type): assert isinstance(router, SubstreamPartitionRouter) elif expected_type == "GroupingPartitionRouter": 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. + """ + source = ConcurrentDeclarativeSource( + source_config=_MANIFEST, config=_CONFIG, catalog=None, state=None + ) + + call_order: list[str] = [] + original_set_api_budget = source._constructor.set_api_budget + + def tracking_set_api_budget(*args, **kwargs): + call_order.append("set_api_budget") + return original_set_api_budget(*args, **kwargs) + + original_dynamic_stream_configs = source._dynamic_stream_configs + + 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": [ + { + "type": "MovingWindowCallRatePolicy", + "rates": [{"type": "Rate", "limit": 5, "interval": "PT1S"}], + "matchers": [], + } + ], + } + + with ( + patch.object(source._constructor, "set_api_budget", side_effect=tracking_set_api_budget), + patch.object( + source, "_dynamic_stream_configs", side_effect=tracking_dynamic_stream_configs + ), + ): + source.streams(config=_CONFIG) + + 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" + assert call_order.index("set_api_budget") < call_order.index("dynamic_stream_configs"), ( + f"set_api_budget must be called before dynamic_stream_configs, but call order was: {call_order}" + ) + + +def test_dynamic_stream_discovery_http_requests_use_api_budget(): + """Verify that HttpComponentsResolver's requester receives the configured api_budget. + + Regression test for https://github.com/airbytehq/oncall/issues/11954 + The discovery HTTP requests made by HttpComponentsResolver must be rate-limited + by the api_budget configured in the manifest. + """ + manifest = { + "version": "5.0.0", + "definitions": { + "selector": { + "type": "RecordSelector", + "extractor": {"type": "DpathExtractor", "field_path": []}, + }, + "requester": { + "type": "HttpRequester", + "url_base": "https://api.test.com", + "http_method": "GET", + "authenticator": {"type": "NoAuth"}, + }, + }, + "dynamic_streams": [ + { + "type": "DynamicDeclarativeStream", + "stream_template": { + "type": "DeclarativeStream", + "$parameters": { + "name": "dynamic_items", + "primary_key": "id", + "url_base": "https://api.test.com", + }, + "schema_loader": { + "type": "InlineSchemaLoader", + "schema": { + "$schema": "https://json-schema.org/draft-07/schema#", + "type": "object", + "properties": {"id": {"type": "string"}}, + }, + }, + "retriever": { + "type": "SimpleRetriever", + "record_selector": {"$ref": "#/definitions/selector"}, + "paginator": {"type": "NoPagination"}, + "requester": { + "$ref": "#/definitions/requester", + "path": "/items", + }, + }, + }, + "components_resolver": { + "type": "HttpComponentsResolver", + "$parameters": { + "name": "resolver", + "primary_key": "id", + "url_base": "https://api.test.com", + }, + "retriever": { + "type": "SimpleRetriever", + "record_selector": {"$ref": "#/definitions/selector"}, + "paginator": {"type": "NoPagination"}, + "requester": { + "$ref": "#/definitions/requester", + "path": "/components", + }, + }, + "components_mapping": [ + { + "type": "ComponentMappingDefinition", + "field_path": ["name"], + "value": "{{ components_values.name }}", + } + ], + }, + } + ], + "api_budget": { + "type": "HTTPAPIBudget", + "policies": [ + { + "type": "MovingWindowCallRatePolicy", + "rates": [{"type": "Rate", "limit": 5, "interval": "PT1S"}], + "matchers": [], + } + ], + }, + "check": {"type": "CheckStream", "stream_names": []}, + } + + source = ConcurrentDeclarativeSource( + source_config=manifest, config={"api_key": "test"}, catalog=None, state=None + ) + + captured_resolvers: list[Any] = [] + + def capturing_resolve(resolver_self, *args, **kwargs): + captured_resolvers.append(resolver_self) + return iter([]) + + with patch.object(HttpComponentsResolver, "resolve_components", capturing_resolve): + source.streams(config={"api_key": "test"}) + + assert len(captured_resolvers) == 1, ( + f"Expected exactly one HttpComponentsResolver, got {len(captured_resolvers)}" + ) + resolver = captured_resolvers[0] + requester = resolver.retriever.requester + assert requester.api_budget is not None, ( + "HttpComponentsResolver's requester should have api_budget set during dynamic stream " + "discovery, but it was None. This means discovery HTTP requests are not rate-limited." + )