airbyte_cdk.sources.streams.call_rate

  1#
  2# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
  3#
  4
  5import abc
  6import dataclasses
  7import datetime
  8import logging
  9import re
 10import time
 11from datetime import timedelta
 12from threading import RLock
 13from typing import TYPE_CHECKING, Any, Mapping, Optional
 14from urllib import parse
 15
 16import requests
 17import requests_cache
 18from pyrate_limiter import InMemoryBucket, Limiter, RateItem, TimeClock
 19from pyrate_limiter import Rate as PyRateRate
 20from pyrate_limiter.exceptions import BucketFullException
 21
 22# prevents mypy from complaining about missing session attributes in LimiterMixin
 23if TYPE_CHECKING:
 24    MIXIN_BASE = requests.Session
 25else:
 26    MIXIN_BASE = object
 27
 28logger = logging.getLogger("airbyte")
 29logging.getLogger("pyrate_limiter").setLevel(logging.WARNING)
 30
 31
 32@dataclasses.dataclass
 33class Rate:
 34    """Call rate limit"""
 35
 36    limit: int
 37    interval: timedelta
 38
 39
 40class CallRateLimitHit(Exception):
 41    def __init__(self, error: str, item: Any, weight: int, rate: str, time_to_wait: timedelta):
 42        """Constructor
 43
 44        :param error: error message
 45        :param item: object passed into acquire_call
 46        :param weight: how many credits were requested
 47        :param rate: string representation of the rate violated
 48        :param time_to_wait: how long should wait util more call will be available
 49        """
 50        self.item = item
 51        self.weight = weight
 52        self.rate = rate
 53        self.time_to_wait = time_to_wait
 54        super().__init__(error)
 55
 56
 57class AbstractCallRatePolicy(abc.ABC):
 58    """Call rate policy interface.
 59    Should be configurable with different rules, like N per M for endpoint X. Endpoint X is matched with APIBudget.
 60    """
 61
 62    @abc.abstractmethod
 63    def matches(self, request: Any) -> bool:
 64        """Tells if this policy matches specific request and should apply to it
 65
 66        :param request:
 67        :return: True if policy should apply to this request, False - otherwise
 68        """
 69
 70    @abc.abstractmethod
 71    def try_acquire(self, request: Any, weight: int) -> None:
 72        """Try to acquire request
 73
 74        :param request: a request object representing a single call to API
 75        :param weight: number of requests to deduct from credit
 76        :return:
 77        """
 78
 79    @abc.abstractmethod
 80    def update(
 81        self, available_calls: Optional[int], call_reset_ts: Optional[datetime.datetime]
 82    ) -> None:
 83        """Update call rate counting with current values
 84
 85        :param available_calls:
 86        :param call_reset_ts:
 87        """
 88
 89
 90class RequestMatcher(abc.ABC):
 91    """Callable that help to match a request object with call rate policies."""
 92
 93    @abc.abstractmethod
 94    def __call__(self, request: Any) -> bool:
 95        """
 96
 97        :param request:
 98        :return: True if matches the provided request object, False - otherwise
 99        """
100
101
102class HttpRequestMatcher(RequestMatcher):
103    """Simple implementation of RequestMatcher for HTTP requests using HttpRequestRegexMatcher under the hood."""
104
105    def __init__(
106        self,
107        method: Optional[str] = None,
108        url: Optional[str] = None,
109        params: Optional[Mapping[str, Any]] = None,
110        headers: Optional[Mapping[str, Any]] = None,
111    ):
112        """Constructor
113
114        :param method: HTTP method (e.g., "GET", "POST").
115        :param url: Full URL to match.
116        :param params: Dictionary of query parameters to match.
117        :param headers: Dictionary of headers to match.
118        """
119        # Parse the URL to extract the base and path
120        if url:
121            parsed_url = parse.urlsplit(url)
122            url_base = f"{parsed_url.scheme}://{parsed_url.netloc}"
123            url_path = parsed_url.path if parsed_url.path != "/" else None
124        else:
125            url_base = None
126            url_path = None
127
128        # Use HttpRequestRegexMatcher under the hood
129        self._regex_matcher = HttpRequestRegexMatcher(
130            method=method,
131            url_base=url_base,
132            url_path_pattern=re.escape(url_path) if url_path else None,
133            params=params,
134            headers=headers,
135        )
136
137    def __call__(self, request: Any) -> bool:
138        """
139        :param request: A requests.Request or requests.PreparedRequest instance.
140        :return: True if the request matches all provided criteria; False otherwise.
141        """
142        return self._regex_matcher(request)
143
144    def __str__(self) -> str:
145        return (
146            f"HttpRequestMatcher(method={self._regex_matcher._method}, "
147            f"url={self._regex_matcher._url_base}{self._regex_matcher._url_path_pattern.pattern if self._regex_matcher._url_path_pattern else ''}, "
148            f"params={self._regex_matcher._params}, headers={self._regex_matcher._headers})"
149        )
150
151
152class HttpRequestRegexMatcher(RequestMatcher):
153    """
154    Extended RequestMatcher for HTTP requests that supports matching on:
155      - HTTP method (case-insensitive)
156      - URL base (scheme + netloc) optionally
157      - URL path pattern (a regex applied to the path portion of the URL)
158      - Query parameters (must be present)
159      - Headers (header names compared case-insensitively)
160    """
161
162    def __init__(
163        self,
164        method: Optional[str] = None,
165        url_base: Optional[str] = None,
166        url_path_pattern: Optional[str] = None,
167        params: Optional[Mapping[str, Any]] = None,
168        headers: Optional[Mapping[str, Any]] = None,
169        weight: Optional[int] = None,
170    ):
171        """
172        :param method: HTTP method (e.g. "GET", "POST"); compared case-insensitively.
173        :param url_base: Base URL (scheme://host) that must match.
174        :param url_path_pattern: A regex pattern that will be applied to the path portion of the URL.
175        :param params: Dictionary of query parameters that must be present in the request.
176        :param headers: Dictionary of headers that must be present (header keys are compared case-insensitively).
177        :param weight: The weight of a request matching this matcher. If set, this value is used
178            when acquiring a call from the rate limiter, enabling cost-based rate limiting
179            where different endpoints consume different amounts from a shared budget.
180            If not set, each request counts as 1.
181        """
182        if weight is not None and weight < 1:
183            raise ValueError(f"weight must be >= 1, got {weight}")
184        self._weight = weight
185        self._method = method.upper() if method else None
186
187        # Normalize the url_base if provided: remove trailing slash.
188        self._url_base = url_base.rstrip("/") if url_base else None
189
190        # Compile the URL path pattern if provided.
191        self._url_path_pattern = re.compile(url_path_pattern) if url_path_pattern else None
192
193        # Normalize query parameters to strings.
194        self._params = {str(k): str(v) for k, v in (params or {}).items()}
195
196        # Normalize header keys to lowercase.
197        self._headers = {str(k).lower(): str(v) for k, v in (headers or {}).items()}
198
199    @staticmethod
200    def _match_dict(obj: Mapping[str, Any], pattern: Mapping[str, Any]) -> bool:
201        """Check that every key/value in the pattern exists in the object."""
202        return pattern.items() <= obj.items()
203
204    def __call__(self, request: Any) -> bool:
205        """
206        :param request: A requests.Request or requests.PreparedRequest instance.
207        :return: True if the request matches all provided criteria; False otherwise.
208        """
209        # Prepare the request (if needed) and extract the URL details.
210        if isinstance(request, requests.Request):
211            prepared_request = request.prepare()
212        elif isinstance(request, requests.PreparedRequest):
213            prepared_request = request
214        else:
215            return False
216
217        # Check HTTP method.
218        if self._method is not None:
219            if prepared_request.method != self._method:
220                return False
221
222        # Parse the URL.
223        parsed_url = parse.urlsplit(prepared_request.url)
224        # Reconstruct the base: scheme://netloc
225        request_url_base = f"{str(parsed_url.scheme)}://{str(parsed_url.netloc)}"
226        # The path (without query parameters)
227        request_path = str(parsed_url.path).rstrip("/")
228
229        # If a base URL is provided, check that it matches.
230        if self._url_base is not None:
231            if request_url_base != self._url_base:
232                return False
233
234        # If a URL path pattern is provided, ensure the path matches the regex.
235        if self._url_path_pattern is not None:
236            if not self._url_path_pattern.search(request_path):
237                return False
238
239        # Check query parameters.
240        if self._params:
241            query_params = dict(parse.parse_qsl(str(parsed_url.query)))
242            if not self._match_dict(query_params, self._params):
243                return False
244
245        # Check headers (normalize keys to lower-case).
246        if self._headers:
247            req_headers = {k.lower(): v for k, v in prepared_request.headers.items()}
248            if not self._match_dict(req_headers, self._headers):
249                return False
250
251        return True
252
253    @property
254    def weight(self) -> Optional[int]:
255        """The weight of a request matching this matcher, or None if not set."""
256        return self._weight
257
258    def __str__(self) -> str:
259        regex = self._url_path_pattern.pattern if self._url_path_pattern else None
260        return (
261            f"HttpRequestRegexMatcher(method={self._method}, url_base={self._url_base}, "
262            f"url_path_pattern={regex}, params={self._params}, headers={self._headers}, weight={self._weight})"
263        )
264
265
266class BaseCallRatePolicy(AbstractCallRatePolicy, abc.ABC):
267    def __init__(self, matchers: list[RequestMatcher]):
268        self._matchers = matchers
269
270    def matches(self, request: Any) -> bool:
271        """Tell if this policy matches specific request and should apply to it
272
273        :param request:
274        :return: True if policy should apply to this request, False - otherwise
275        """
276
277        if not self._matchers:
278            return True
279        return any(matcher(request) for matcher in self._matchers)
280
281    def get_weight(self, request: Any) -> int:
282        """Get the weight for a request based on the first matching matcher.
283
284        If a matcher has a weight configured, that weight is used.
285        Otherwise, defaults to 1.
286
287        :param request: a request object
288        :return: the weight for this request
289        """
290        for matcher in self._matchers:
291            if matcher(request):
292                if isinstance(matcher, HttpRequestRegexMatcher) and matcher.weight is not None:
293                    return matcher.weight
294                return 1
295        return 1
296
297
298class UnlimitedCallRatePolicy(BaseCallRatePolicy):
299    """
300    This policy is for explicit unlimited call rates.
301    It can be used when we want to match a specific group of requests and don't apply any limits.
302
303    Example:
304
305    APICallBudget(
306        [
307            UnlimitedCallRatePolicy(
308                matchers=[HttpRequestMatcher(url="/some/method", headers={"sandbox": true})],
309            ),
310            FixedWindowCallRatePolicy(
311                matchers=[HttpRequestMatcher(url="/some/method")],
312                next_reset_ts=datetime.now(),
313                period=timedelta(hours=1)
314                call_limit=1000,
315            ),
316        ]
317    )
318
319    The code above will limit all calls to /some/method except calls that have header sandbox=True
320    """
321
322    def try_acquire(self, request: Any, weight: int) -> None:
323        """Do nothing"""
324
325    def update(
326        self, available_calls: Optional[int], call_reset_ts: Optional[datetime.datetime]
327    ) -> None:
328        """Do nothing"""
329
330
331class FixedWindowCallRatePolicy(BaseCallRatePolicy):
332    def __init__(
333        self,
334        next_reset_ts: datetime.datetime,
335        period: timedelta,
336        call_limit: int,
337        matchers: list[RequestMatcher],
338    ):
339        """A policy that allows {call_limit} calls within a {period} time interval
340
341        :param next_reset_ts: next call rate reset time point
342        :param period: call rate reset period
343        :param call_limit:
344        :param matchers:
345        """
346
347        self._next_reset_ts = next_reset_ts
348        self._offset = period
349        self._call_limit = call_limit
350        self._calls_num = 0
351        self._lock = RLock()
352        super().__init__(matchers=matchers)
353
354    def try_acquire(self, request: Any, weight: int) -> None:
355        if weight > self._call_limit:
356            raise ValueError("Weight can not exceed the call limit")
357        if not self.matches(request):
358            raise ValueError("Request does not match the policy")
359
360        with self._lock:
361            self._update_current_window()
362
363            if self._calls_num + weight > self._call_limit:
364                reset_in = self._next_reset_ts - datetime.datetime.now()
365                error_message = (
366                    f"reached maximum number of allowed calls {self._call_limit} "
367                    f"per {self._offset} interval, next reset in {reset_in}."
368                )
369                raise CallRateLimitHit(
370                    error=error_message,
371                    item=request,
372                    weight=weight,
373                    rate=f"{self._call_limit} per {self._offset}",
374                    time_to_wait=reset_in,
375                )
376
377            self._calls_num += weight
378
379    def __str__(self) -> str:
380        matcher_str = ", ".join(f"{matcher}" for matcher in self._matchers)
381        return (
382            f"FixedWindowCallRatePolicy(call_limit={self._call_limit}, period={self._offset}, "
383            f"calls_used={self._calls_num}, next_reset={self._next_reset_ts}, "
384            f"matchers=[{matcher_str}])"
385        )
386
387    def update(
388        self, available_calls: Optional[int], call_reset_ts: Optional[datetime.datetime]
389    ) -> None:
390        """Update call rate counters, by default, only reacts to decreasing updates of available_calls and changes to call_reset_ts.
391        We ignore updates with available_calls > current_available_calls to support call rate limits that are lower than API limits.
392
393        :param available_calls:
394        :param call_reset_ts:
395        """
396        with self._lock:
397            self._update_current_window()
398            current_available_calls = self._call_limit - self._calls_num
399
400            if available_calls is not None and current_available_calls > available_calls:
401                logger.debug(
402                    "got rate limit update from api, adjusting available calls from %s to %s",
403                    current_available_calls,
404                    available_calls,
405                )
406                self._calls_num = self._call_limit - available_calls
407
408            if call_reset_ts is not None and call_reset_ts != self._next_reset_ts:
409                logger.debug(
410                    "got rate limit update from api, adjusting reset time from %s to %s",
411                    self._next_reset_ts,
412                    call_reset_ts,
413                )
414                self._next_reset_ts = call_reset_ts
415
416    def _update_current_window(self) -> None:
417        now = datetime.datetime.now()
418        if now > self._next_reset_ts:
419            logger.debug("started new window, %s calls available now", self._call_limit)
420            self._next_reset_ts = self._next_reset_ts + self._offset
421            self._calls_num = 0
422
423
424class MovingWindowCallRatePolicy(BaseCallRatePolicy):
425    """
426    Policy to control requests rate implemented on top of PyRateLimiter lib.
427    The main difference between this policy and FixedWindowCallRatePolicy is that the rate-limiting window
428    is moving along requests that we made, and there is no moment when we reset an available number of calls.
429    This strategy requires saving of timestamps of all requests within a window.
430    """
431
432    def __init__(self, rates: list[Rate], matchers: list[RequestMatcher]):
433        """Constructor
434
435        :param rates: list of rates, the order is important and must be ascending
436        :param matchers:
437        """
438        if not rates:
439            raise ValueError("The list of rates can not be empty")
440        pyrate_rates = [
441            PyRateRate(limit=rate.limit, interval=int(rate.interval.total_seconds() * 1000))
442            for rate in rates
443        ]
444        self._bucket = InMemoryBucket(pyrate_rates)
445        # Limiter will create the background task that clears old requests in the bucket
446        self._limiter = Limiter(self._bucket)
447        super().__init__(matchers=matchers)
448
449    def try_acquire(self, request: Any, weight: int) -> None:
450        if not self.matches(request):
451            raise ValueError("Request does not match the policy")
452        lowest_limit = min(rate.limit for rate in self._bucket.rates)
453        if weight > lowest_limit:
454            raise ValueError(
455                f"Weight can not exceed the lowest configured rate limit ({lowest_limit})"
456            )
457
458        try:
459            self._limiter.try_acquire(request, weight=weight)
460        except BucketFullException as exc:
461            item = self._limiter.bucket_factory.wrap_item(request, weight)
462            assert isinstance(item, RateItem)
463
464            with self._limiter.lock:
465                time_to_wait = self._bucket.waiting(item)
466                assert isinstance(time_to_wait, int)
467
468                raise CallRateLimitHit(
469                    error=str(exc.meta_info["error"]),
470                    item=request,
471                    weight=int(exc.meta_info["weight"]),
472                    rate=str(exc.meta_info["rate"]),
473                    time_to_wait=timedelta(milliseconds=time_to_wait),
474                )
475
476    def update(
477        self, available_calls: Optional[int], call_reset_ts: Optional[datetime.datetime]
478    ) -> None:
479        """Adjust call bucket to reflect the state of the API server
480
481        :param available_calls:
482        :param call_reset_ts:
483        :return:
484        """
485        if (
486            available_calls is not None and call_reset_ts is None
487        ):  # we do our best to sync buckets with API
488            if available_calls == 0:
489                with self._limiter.lock:
490                    now: int = TimeClock().now()  # type: ignore[no-untyped-call]
491                    items_to_add = self._get_available_calls(now)
492                    if items_to_add > 0:
493                        added = self._bucket.put(
494                            RateItem(name="dummy", timestamp=now, weight=items_to_add)
495                        )
496                        if not added:
497                            logger.warning(
498                                "Failed to sync rate limit bucket with API: could not add %s items",
499                                items_to_add,
500                            )
501        # TODO: add support if needed, it might be that it is not possible to make a good solution for this case
502        # if available_calls is not None and call_reset_ts is not None:
503        #     ts = call_reset_ts.timestamp()
504
505    def _get_available_calls(self, now: int) -> int:
506        """Return the number of calls still allowed by the strictest configured rate at the given timestamp (ms)."""
507        return min(
508            rate.limit
509            - sum(1 for item in self._bucket.items if item.timestamp >= now - rate.interval)
510            for rate in self._bucket.rates
511        )
512
513    def __str__(self) -> str:
514        """Return a human-friendly description of the moving window rate policy for logging purposes."""
515        rates_info = ", ".join(
516            f"{rate.limit} per {timedelta(milliseconds=rate.interval)}"
517            for rate in self._bucket.rates
518        )
519        current_bucket_count = self._bucket.count()
520        matcher_str = ", ".join(f"{matcher}" for matcher in self._matchers)
521        return (
522            f"MovingWindowCallRatePolicy(rates=[{rates_info}], current_bucket_count={current_bucket_count}, "
523            f"matchers=[{matcher_str}])"
524        )
525
526
527class AbstractAPIBudget(abc.ABC):
528    """Interface to some API where a client allowed to have N calls per T interval.
529
530    Important: APIBudget is not doing any API calls, the end user code is responsible to call this interface
531        to respect call rate limitation of the API.
532
533    It supports multiple policies applied to different group of requests. To distinct these groups we use RequestMatchers.
534    Individual policy represented by MovingWindowCallRatePolicy and currently supports only moving window strategy.
535    """
536
537    @abc.abstractmethod
538    def acquire_call(
539        self, request: Any, block: bool = True, timeout: Optional[float] = None
540    ) -> None:
541        """Try to get a call from budget, will block by default
542
543        :param request:
544        :param block: when true (default) will block the current thread until call credit is available
545        :param timeout: if set will limit maximum time in block, otherwise will wait until credit is available
546        :raises: CallRateLimitHit - when no credits left and if timeout was set the waiting time exceed the timeout
547        """
548
549    @abc.abstractmethod
550    def get_matching_policy(self, request: Any) -> Optional[AbstractCallRatePolicy]:
551        """Find matching call rate policy for specific request"""
552
553    @abc.abstractmethod
554    def update_from_response(self, request: Any, response: Any) -> None:
555        """Update budget information based on response from API
556
557        :param request: the initial request that triggered this response
558        :param response: response from the API
559        """
560
561
562class APIBudget(AbstractAPIBudget):
563    """Default APIBudget implementation"""
564
565    def __init__(
566        self, policies: list[AbstractCallRatePolicy], maximum_attempts_to_acquire: int = 100000
567    ) -> None:
568        """Constructor
569
570        :param policies: list of policies in this budget
571        :param maximum_attempts_to_acquire: number of attempts before throwing hit ratelimit exception, we put some big number here
572         to avoid situations when many threads compete with each other for a few lots over a significant amount of time
573        """
574
575        self._policies = policies
576        self._maximum_attempts_to_acquire = maximum_attempts_to_acquire
577
578    def _extract_endpoint(self, request: Any) -> str:
579        """Extract the endpoint URL from the request if available."""
580        endpoint = None
581        try:
582            # If the request is already a PreparedRequest, it should have a URL.
583            if isinstance(request, requests.PreparedRequest):
584                endpoint = request.url
585            # If it's a requests.Request, we call prepare() to extract the URL.
586            elif isinstance(request, requests.Request):
587                prepared = request.prepare()
588                endpoint = prepared.url
589        except Exception as e:
590            logger.debug(f"Error extracting endpoint: {e}")
591        if endpoint:
592            return endpoint
593        return "unknown endpoint"
594
595    def get_matching_policy(self, request: Any) -> Optional[AbstractCallRatePolicy]:
596        for policy in self._policies:
597            if policy.matches(request):
598                return policy
599        return None
600
601    def acquire_call(
602        self, request: Any, block: bool = True, timeout: Optional[float] = None
603    ) -> None:
604        """Try to get a call from budget, will block by default.
605        Matchers will be called sequentially in the same order they were added.
606        The first matcher that returns True will
607
608        :param request: the API request
609        :param block: when True (default) will block until a call credit is available
610        :param timeout: if provided, limits maximum waiting time; otherwise, waits indefinitely
611        :raises: CallRateLimitHit if the call credit cannot be acquired within the timeout
612        """
613
614        policy = self.get_matching_policy(request)
615        endpoint = self._extract_endpoint(request)
616        if policy:
617            logger.debug(f"Acquiring call for endpoint {endpoint} using policy: {policy}")
618            self._do_acquire(request=request, policy=policy, block=block, timeout=timeout)
619        elif self._policies:
620            logger.debug(
621                f"No policies matched for endpoint {endpoint} (request: {request}). Allowing call by default."
622            )
623
624    def update_from_response(self, request: Any, response: Any) -> None:
625        """Update budget information based on the API response.
626
627        :param request: the initial request that triggered this response
628        :param response: response from the API
629        """
630        pass
631
632    def _do_acquire(
633        self, request: Any, policy: AbstractCallRatePolicy, block: bool, timeout: Optional[float]
634    ) -> None:
635        """Internal method to try to acquire a call credit.
636
637        :param request: the API request
638        :param policy: the matching rate-limiting policy
639        :param block: indicates whether to block until a call credit is available
640        :param timeout: maximum time to wait if blocking
641        :raises: CallRateLimitHit if unable to acquire a call credit
642        """
643        last_exception = None
644        endpoint = self._extract_endpoint(request)
645        # sometimes we spend all budget before a second attempt, so we have a few more attempts
646        for attempt in range(1, self._maximum_attempts_to_acquire):
647            try:
648                weight = policy.get_weight(request) if isinstance(policy, BaseCallRatePolicy) else 1
649                policy.try_acquire(request, weight=weight)
650                return
651            except CallRateLimitHit as exc:
652                last_exception = exc
653                if block:
654                    if timeout is not None:
655                        time_to_wait = min(timedelta(seconds=timeout), exc.time_to_wait)
656                    else:
657                        time_to_wait = exc.time_to_wait
658                    # Ensure we never sleep for a negative duration.
659                    time_to_wait = max(timedelta(0), time_to_wait)
660                    logger.debug(
661                        f"Policy {policy} reached call limit for endpoint {endpoint} ({exc.rate}). "
662                        f"Sleeping for {time_to_wait} on attempt {attempt}."
663                    )
664                    time.sleep(time_to_wait.total_seconds())
665                else:
666                    logger.debug(
667                        f"Policy {policy} reached call limit for endpoint {endpoint} ({exc.rate}) "
668                        f"and blocking is disabled."
669                    )
670                    raise
671
672        if last_exception:
673            logger.debug(
674                f"Exhausted all {self._maximum_attempts_to_acquire} attempts to acquire a call for endpoint {endpoint} "
675                f"using policy: {policy}"
676            )
677            raise last_exception
678
679
680class HttpAPIBudget(APIBudget):
681    """Implementation of AbstractAPIBudget for HTTP"""
682
683    def __init__(
684        self,
685        ratelimit_reset_header: str = "ratelimit-reset",
686        ratelimit_remaining_header: str = "ratelimit-remaining",
687        status_codes_for_ratelimit_hit: list[int] = [429],
688        **kwargs: Any,
689    ):
690        """Constructor
691
692        :param ratelimit_reset_header: name of the header that has a timestamp of the next reset of call budget
693        :param ratelimit_remaining_header: name of the header that has the number of calls left
694        :param status_codes_for_ratelimit_hit: list of HTTP status codes that signal about rate limit being hit
695        """
696        self._ratelimit_reset_header = ratelimit_reset_header
697        self._ratelimit_remaining_header = ratelimit_remaining_header
698        self._status_codes_for_ratelimit_hit = status_codes_for_ratelimit_hit
699        super().__init__(**kwargs)
700
701    def update_from_response(self, request: Any, response: Any) -> None:
702        policy = self.get_matching_policy(request)
703        if not policy:
704            return
705
706        if isinstance(response, requests.Response):
707            available_calls = self.get_calls_left_from_response(response)
708            reset_ts = self.get_reset_ts_from_response(response)
709            policy.update(available_calls=available_calls, call_reset_ts=reset_ts)
710
711    def get_reset_ts_from_response(
712        self, response: requests.Response
713    ) -> Optional[datetime.datetime]:
714        if response.headers.get(self._ratelimit_reset_header):
715            return datetime.datetime.fromtimestamp(
716                int(response.headers[self._ratelimit_reset_header])
717            )
718        return None
719
720    def get_calls_left_from_response(self, response: requests.Response) -> Optional[int]:
721        if response.headers.get(self._ratelimit_remaining_header):
722            return int(response.headers[self._ratelimit_remaining_header])
723
724        if response.status_code in self._status_codes_for_ratelimit_hit:
725            return 0
726
727        return None
728
729
730class LimiterMixin(MIXIN_BASE):
731    """Mixin class that adds rate-limiting behavior to requests."""
732
733    def __init__(
734        self,
735        api_budget: AbstractAPIBudget,
736        **kwargs: Any,
737    ):
738        self._api_budget = api_budget
739        super().__init__(**kwargs)  # type: ignore # Base Session doesn't take any kwargs
740
741    def send(self, request: requests.PreparedRequest, **kwargs: Any) -> requests.Response:
742        """Send a request with rate-limiting."""
743        self._api_budget.acquire_call(request)
744        response = super().send(request, **kwargs)
745        self._api_budget.update_from_response(request, response)
746        return response
747
748
749class LimiterSession(LimiterMixin, requests.Session):
750    """Session that adds rate-limiting behavior to requests."""
751
752
753class CachedLimiterSession(requests_cache.CacheMixin, LimiterMixin, requests.Session):
754    """Session class with caching and rate-limiting behavior."""