From 66bf33cfab7cbdbbb26c4c9e64d052846d9cfa78 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 25 May 2026 14:50:07 +0000 Subject: [PATCH 1/2] fix: pass API budget to custom requesters --- .../parsers/model_to_component_factory.py | 2 + .../test_model_to_component_factory.py | 48 ++++++++++++++++++- 2 files changed, 49 insertions(+), 1 deletion(-) 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 bc9b4f82c0..ea4613d755 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -1801,6 +1801,8 @@ def create_custom_component(self, model: Any, config: Config, **kwargs: Any) -> component_fields = get_type_hints(custom_component_class) model_args = model.dict() model_args["config"] = config + if "api_budget" in component_fields: + model_args["api_budget"] = self._api_budget # There are cases where a parent component will pass arguments to a child component via kwargs. When there are field collisions # we defer to these arguments over the component's definition 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 add9a1c420..77f27781ac 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.declarative_component_schema import ( + CustomRequester as CustomRequesterModel, +) from airbyte_cdk.sources.declarative.models.declarative_component_schema import ( OffsetIncrement as OffsetIncrementModel, ) @@ -173,7 +176,7 @@ ) from airbyte_cdk.sources.declarative.yaml_declarative_source import YamlDeclarativeSource from airbyte_cdk.sources.message.repository import StateFilteringMessageRepository -from airbyte_cdk.sources.streams.call_rate import MovingWindowCallRatePolicy +from airbyte_cdk.sources.streams.call_rate import HttpAPIBudget, MovingWindowCallRatePolicy from airbyte_cdk.sources.streams.concurrent.clamping import ( ClampingEndProvider, DayClampingStrategy, @@ -4239,6 +4242,49 @@ def test_api_budget(): assert policy._bucket.rates[0].interval == 100 # 100 ms +def test_api_budget_passed_to_custom_requester(): + manifest = { + "type": "DeclarativeSource", + "api_budget": { + "type": "HTTPAPIBudget", + "policies": [ + { + "type": "MovingWindowCallRatePolicy", + "rates": [ + { + "limit": 3, + "interval": "PT0.1S", + } + ], + "matchers": [], + } + ], + }, + "my_requester": { + "type": "CustomRequester", + "class_name": "unit_tests.sources.declarative.parsers.testing_components.TestingRequester", + "path": "/v3/marketing/lists", + "url_base": "https://api.sendgrid.com", + "http_method": "GET", + }, + } + + factory = ModelToComponentFactory() + factory.set_api_budget(manifest["api_budget"], input_config) + + custom_requester = factory.create_component( + model_type=CustomRequesterModel, + component_definition=manifest["my_requester"], + config=input_config, + name="lists_stream", + decoder=None, + ) + + assert isinstance(custom_requester.api_budget, HttpAPIBudget) + assert custom_requester._http_client._api_budget is custom_requester.api_budget + assert len(custom_requester._http_client._api_budget._policies) == 1 + + def test_api_budget_fixed_window_policy(): manifest = { "type": "DeclarativeSource", From 5fa728fbba5d734f860230167b32c364a56cd32e Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 25 May 2026 15:10:42 +0000 Subject: [PATCH 2/2] fix: preserve custom requester API budget defaults --- .../parsers/model_to_component_factory.py | 6 ++- .../test_model_to_component_factory.py | 41 +++++++++++++++++++ .../declarative/parsers/testing_components.py | 7 ++++ 3 files changed, 52 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 ea4613d755..896d82a0f6 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -1801,8 +1801,6 @@ def create_custom_component(self, model: Any, config: Config, **kwargs: Any) -> component_fields = get_type_hints(custom_component_class) model_args = model.dict() model_args["config"] = config - if "api_budget" in component_fields: - model_args["api_budget"] = self._api_budget # There are cases where a parent component will pass arguments to a child component via kwargs. When there are field collisions # we defer to these arguments over the component's definition @@ -1861,6 +1859,10 @@ def create_custom_component(self, model: Any, config: Config, **kwargs: Any) -> for class_field in component_fields.keys() if class_field in model_args } + + if "api_budget" in component_fields and kwargs.get("api_budget") is None: + kwargs["api_budget"] = self._api_budget + return custom_component_class(**kwargs) @staticmethod 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 77f27781ac..9fd1cb58b0 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 @@ -4285,6 +4285,47 @@ def test_api_budget_passed_to_custom_requester(): assert len(custom_requester._http_client._api_budget._policies) == 1 +def test_api_budget_does_not_override_custom_requester_default_value(): + manifest = { + "type": "DeclarativeSource", + "api_budget": { + "type": "HTTPAPIBudget", + "policies": [ + { + "type": "MovingWindowCallRatePolicy", + "rates": [ + { + "limit": 3, + "interval": "PT0.1S", + } + ], + "matchers": [], + } + ], + }, + "my_requester": { + "type": "CustomRequester", + "class_name": "unit_tests.sources.declarative.parsers.testing_components.TestingRequesterWithDefaultBudget", + "path": "/v3/marketing/lists", + "url_base": "https://api.sendgrid.com", + "http_method": "GET", + }, + } + + factory = ModelToComponentFactory() + factory.set_api_budget(manifest["api_budget"], input_config) + + custom_requester = factory.create_component( + model_type=CustomRequesterModel, + component_definition=manifest["my_requester"], + config=input_config, + name="lists_stream", + decoder=None, + ) + + assert custom_requester.api_budget is None + + def test_api_budget_fixed_window_policy(): manifest = { "type": "DeclarativeSource", diff --git a/unit_tests/sources/declarative/parsers/testing_components.py b/unit_tests/sources/declarative/parsers/testing_components.py index 6269acb8e3..b21b63efd6 100644 --- a/unit_tests/sources/declarative/parsers/testing_components.py +++ b/unit_tests/sources/declarative/parsers/testing_components.py @@ -114,3 +114,10 @@ def __post_init__(self, parameters: Mapping[str, Any]) -> None: parameters=parameters or {}, ) super().__post_init__(parameters) + + +@dataclass +class TestingRequesterWithDefaultBudget(TestingRequester): + def __post_init__(self, parameters: Mapping[str, Any]) -> None: + self.api_budget = None + super().__post_init__(parameters)