fix(cdk): propagate api_budget through CustomRequester HttpClient overrides

Author: agarctfi

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -1865,6 +1865,7 @@ def create_custom_component(self, model: Any, config: Config, **kwargs: Any) ->
         # or an explicit kwarg has already provided one. Without this, custom requesters
         # silently lose the connector-level HTTPAPIBudget and any configured rate-limit
         # policies have no effect at runtime.
+        injected_api_budget = False
         if (
             self._api_budget is not None
             and "api_budget" in component_fields
@@ -1873,8 +1874,35 @@ def create_custom_component(self, model: Any, config: Config, **kwargs: Any) ->
             and issubclass(custom_component_class, HttpRequester)
         ):
             kwargs["api_budget"] = self._api_budget
+            injected_api_budget = True
 
-        return custom_component_class(**kwargs)
+        custom_component = custom_component_class(**kwargs)
+        if injected_api_budget and isinstance(custom_component, HttpRequester):
+            self._sync_injected_api_budget_with_http_client(custom_component)
+
+        return custom_component
+
+    @staticmethod
+    def _sync_injected_api_budget_with_http_client(custom_requester: HttpRequester) -> None:
+        """
+        Custom requesters can replace `_http_client` in `__post_init__` without forwarding `api_budget`.
+        If the factory injected a manifest-level budget and the replacement client kept the default empty
+        budget, sync the active client back to the requester's injected budget.
+        """
+        http_client = getattr(custom_requester, "_http_client", None)
+        http_client_api_budget = getattr(http_client, "_api_budget", None)
+        injected_api_budget = custom_requester.api_budget
+
+        if (
+            http_client is None
+            or http_client_api_budget is None
+            or injected_api_budget is None
+            or http_client_api_budget is injected_api_budget
+        ):
+            return
+
+        if len(getattr(http_client_api_budget, "_policies", [])) == 0:
+            http_client._api_budget = injected_api_budget
 
     @staticmethod
     def _get_class_from_fully_qualified_class_name(

unit_tests/sources/declarative/parsers/test_model_to_component_factory.py

@@ -4368,6 +4368,98 @@ def test_api_budget_propagated_to_custom_requester_subclass_of_http_requester():
     assert custom_requester._http_client._api_budget is custom_requester.api_budget
 
 
+def test_api_budget_propagated_to_custom_requester_that_replaces_http_client():
+    """Injected api budgets must survive requesters that replace `_http_client` in `__post_init__`.
+
+    Some connector requesters call `super().__post_init__()` and then swap in a custom `HttpClient`
+    implementation without forwarding `api_budget`. The budget should still be applied to the active
+    client so manifest-level rate limits continue to work.
+    """
+    manifest_api_budget = {
+        "type": "HTTPAPIBudget",
+        "policies": [
+            {
+                "type": "MovingWindowCallRatePolicy",
+                "rates": [
+                    {
+                        "type": "Rate",
+                        "limit": 60,
+                        "interval": "PT1M",
+                    }
+                ],
+                "matchers": [],
+            }
+        ],
+    }
+
+    custom_requester_definition = {
+        "type": "CustomRequester",
+        "class_name": "unit_tests.sources.declarative.parsers.testing_components.TestingRequesterWithReplacedHttpClient",
+        "url_base": "https://example.org",
+        "path": "/v1/data",
+        "http_method": "GET",
+    }
+
+    config: Mapping[str, Any] = {}
+    local_factory = ModelToComponentFactory()
+    local_factory.set_api_budget(manifest_api_budget, config)
+
+    custom_requester = local_factory.create_component(
+        model_type=CustomRequesterModel,
+        component_definition=custom_requester_definition,
+        config=config,
+        name="custom_stream",
+    )
+
+    assert isinstance(custom_requester, HttpRequester)
+    assert custom_requester.api_budget is not None
+    assert custom_requester._http_client._api_budget is custom_requester.api_budget
+
+
+def test_api_budget_not_overwriting_non_empty_budget_on_replaced_http_client():
+    """A requester that intentionally installs its own budget should keep it."""
+    manifest_api_budget = {
+        "type": "HTTPAPIBudget",
+        "policies": [
+            {
+                "type": "MovingWindowCallRatePolicy",
+                "rates": [
+                    {
+                        "type": "Rate",
+                        "limit": 60,
+                        "interval": "PT1M",
+                    }
+                ],
+                "matchers": [],
+            }
+        ],
+    }
+
+    custom_requester_definition = {
+        "type": "CustomRequester",
+        "class_name": "unit_tests.sources.declarative.parsers.testing_components.TestingRequesterWithReplacedHttpClientAndOwnBudget",
+        "url_base": "https://example.org",
+        "path": "/v1/data",
+        "http_method": "GET",
+    }
+
+    config: Mapping[str, Any] = {}
+    local_factory = ModelToComponentFactory()
+    local_factory.set_api_budget(manifest_api_budget, config)
+
+    custom_requester = local_factory.create_component(
+        model_type=CustomRequesterModel,
+        component_definition=custom_requester_definition,
+        config=config,
+        name="custom_stream",
+    )
+
+    assert isinstance(custom_requester, HttpRequester)
+    assert custom_requester.api_budget is not None
+    assert custom_requester._http_client._api_budget is not custom_requester.api_budget
+    assert len(custom_requester._http_client._api_budget._policies) == 1
+
+
 def test_api_budget_not_propagated_to_non_http_requester_custom_components():
     """Custom components that do NOT subclass `HttpRequester` must not receive `api_budget`.
 

unit_tests/sources/declarative/parsers/testing_components.py

@@ -3,6 +3,7 @@
 #
 
 from dataclasses import dataclass, field
+from datetime import timedelta
 from typing import Any, ClassVar, List, Mapping, Optional
 
 from airbyte_cdk.sources.declarative.extractors import DpathExtractor
@@ -21,6 +22,8 @@
     RequestInput,
 )
 from airbyte_cdk.sources.declarative.retrievers import SimpleRetriever
+from airbyte_cdk.sources.streams.call_rate import APIBudget, MovingWindowCallRatePolicy, Rate
+from airbyte_cdk.sources.streams.http import HttpClient
 
 
 @dataclass
@@ -114,3 +117,43 @@ def __post_init__(self, parameters: Mapping[str, Any]) -> None:
             parameters=parameters or {},
         )
         super().__post_init__(parameters)
+
+
+@dataclass
+class TestingRequesterWithReplacedHttpClient(HttpRequester):
+    def __post_init__(self, parameters: Mapping[str, Any]) -> None:
+        super().__post_init__(parameters)
+        self._http_client = HttpClient(
+            name=self.name,
+            logger=self.logger,
+            error_handler=self.error_handler,
+            authenticator=self._authenticator,
+            use_cache=self.use_cache,
+            backoff_strategy=None,
+            disable_retries=self.disable_retries,
+            message_repository=self.message_repository,
+        )
+
+
+@dataclass
+class TestingRequesterWithReplacedHttpClientAndOwnBudget(HttpRequester):
+    def __post_init__(self, parameters: Mapping[str, Any]) -> None:
+        super().__post_init__(parameters)
+        self._http_client = HttpClient(
+            name=self.name,
+            logger=self.logger,
+            error_handler=self.error_handler,
+            api_budget=APIBudget(
+                policies=[
+                    MovingWindowCallRatePolicy(
+                        rates=[Rate(limit=1, interval=timedelta(seconds=30))],
+                        matchers=[],
+                    )
+                ]
+            ),
+            authenticator=self._authenticator,
+            use_cache=self.use_cache,
+            backoff_strategy=None,
+            disable_retries=self.disable_retries,
+            message_repository=self.message_repository,
+        )