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..896d82a0f6 100644 --- a/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py +++ b/airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py @@ -1859,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 add9a1c420..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 @@ -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,90 @@ 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_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)