Enable client-side rate limiting on source-stripe #31512

airbyte-integrations/connectors/source-stripe/setup.py

@@ -5,7 +5,7 @@
 
 from setuptools import find_packages, setup
 
-MAIN_REQUIREMENTS = ["airbyte-cdk==0.52.8", "stripe==2.56.0", "pendulum==2.1.2"]
+MAIN_REQUIREMENTS = ["airbyte-cdk==0.53.2", "stripe==2.56.0", "pendulum==2.1.2"]
 
 TEST_REQUIREMENTS = ["pytest-mock~=3.6.1", "pytest~=6.1", "requests-mock", "requests_mock~=1.8", "freezegun==1.2.2"]
 

airbyte-integrations/connectors/source-stripe/source_stripe/source.py

@@ -2,7 +2,9 @@
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
 
+import logging
 import os
+from datetime import datetime, timedelta
 from typing import Any, List, Mapping, MutableMapping, Optional, Tuple
 
 import pendulum
@@ -13,6 +15,14 @@
 from airbyte_cdk.sources import AbstractSource
 from airbyte_cdk.sources.message.repository import InMemoryMessageRepository
 from airbyte_cdk.sources.streams import Stream
+from airbyte_cdk.sources.streams.call_rate import (
+    AbstractAPIBudget,
+    FixedWindowCallRatePolicy,
+    HttpAPIBudget,
+    HttpRequestMatcher,
+    MovingWindowCallRatePolicy,
+    Rate,
+)
 from airbyte_cdk.sources.streams.concurrent.adapters import StreamFacade
 from airbyte_cdk.sources.streams.concurrent.cursor import NoopCursor
 from airbyte_cdk.sources.streams.http.auth import TokenAuthenticator
@@ -33,9 +43,12 @@
     UpdatedCursorIncrementalStripeStream,
 )
 
-_MAX_CONCURRENCY = 3
+logger = logging.getLogger("airbyte")
+
+_MAX_CONCURRENCY = 20
 _CACHE_DISABLED = os.environ.get("CACHE_DISABLED")
 USE_CACHE = not _CACHE_DISABLED
+STRIPE_TEST_ACCOUNT_PREFIX = "sk_test_"
 
 
 class SourceStripe(AbstractSource):
@@ -114,6 +127,73 @@ def customers(**args):
             **args,
         )
 
+    @staticmethod
+    def is_test_account(config: Mapping[str, Any]) -> bool:
+        """Check if configuration uses Stripe test account (https://stripe.com/docs/keys#obtain-api-keys)
+
+        :param config:
+        :return: True if configured to use a test account, False - otherwise
+        """
+
+        return str(config["client_secret"]).startswith(STRIPE_TEST_ACCOUNT_PREFIX)
+
+    def get_api_call_budget(self, config: Mapping[str, Any]) -> AbstractAPIBudget:
+        """Get API call budget which connector is allowed to use.
+
+        :param config:
+        :return:
+        """
+
+        max_call_rate = 25 if self.is_test_account(config) else 100
+        if config.get("call_rate_limit"):
+            call_limit = config["call_rate_limit"]
+            if call_limit > max_call_rate:
+                logger.warning(
+                    "call_rate_limit is larger than maximum allowed %s, fallback to default %s.",
+                    max_call_rate,
+                    max_call_rate,
+                )
+                call_limit = max_call_rate
+        else:
+            call_limit = max_call_rate
+
+        if config.get("call_rate_policy") == "fixed_window":
+            policies = [
+                FixedWindowCallRatePolicy(
+                    next_reset_ts=datetime.now(),
+                    period=timedelta(seconds=1),
+                    call_limit=20,
+                    matchers=[
+                        HttpRequestMatcher(url="https://api.stripe.com/v1/files"),
+                        HttpRequestMatcher(url="https://api.stripe.com/v1/file_links"),
+                    ],
+                ),
+                FixedWindowCallRatePolicy(
+                    next_reset_ts=datetime.now(),
+                    period=timedelta(seconds=1),
+                    call_limit=call_limit,
+                    matchers=[],
+                ),
+            ]
+        elif config.get("call_rate_policy") == "moving_window":
+            policies = [
+                MovingWindowCallRatePolicy(
+                    rates=[Rate(limit=20, interval=timedelta(seconds=1))],
+                    matchers=[
+                        HttpRequestMatcher(url="https://api.stripe.com/v1/files"),
+                        HttpRequestMatcher(url="https://api.stripe.com/v1/file_links"),
+                    ],
+                ),
+                MovingWindowCallRatePolicy(
+                    rates=[Rate(limit=call_limit, interval=timedelta(seconds=1))],
+                    matchers=[],
+                ),
+            ]
+        else:
+            policies = []
+
+        return HttpAPIBudget(policies=policies, maximum_attempts_to_acquire=10000)
+
     def streams(self, config: Mapping[str, Any]) -> List[Stream]:
         config = self.validate_and_fill_with_defaults(config)
         authenticator = TokenAuthenticator(config["client_secret"])
@@ -122,6 +202,7 @@ def streams(self, config: Mapping[str, Any]) -> List[Stream]:
             "account_id": config["account_id"],
             "start_date": config["start_date"],
             "slice_range": config["slice_range"],
+            "api_budget": self.get_api_call_budget(config),
         }
         incremental_args = {**args, "lookback_window_days": config["lookback_window_days"]}
         subscriptions = IncrementalStripeStream(

airbyte-integrations/connectors/source-stripe/source_stripe/spec.yaml

@@ -66,5 +66,19 @@ connectionSpecification:
       examples: [1, 2, 3]
       description: >-
         The number of worker thread to use for the sync. The bigger the value is, the faster the sync will be.
-        Be careful as rate limiting is not implemented.
+        The performance upper boundary depends on call_rate_limit setting and type of account.
       order: 5
+    call_rate_limit:
+      type: integer
+      title: Max number of API calls per second
+      examples: [25, 100]
+      description: >-
+        The number of API calls per second that you allow connector to make. This value can not be bigger than real
+        API call rate limit (https://stripe.com/docs/rate-limits). If not specified the default maximum is 25 and 100
+        calls per second for test and production tokens respectively.
+    call_rate_policy:
+      type: string
+      title: Call rate policy to use
+      examples: [fixed_window, moving_window, none]
+      description: >-
+        The call rate algorithm to use for tracking and limiting calls.

airbyte-integrations/connectors/source-stripe/unit_tests/test_source.py

@@ -1,7 +1,7 @@
 #
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
-
+import datetime
 import logging
 from contextlib import nullcontext as does_not_raise
 from unittest.mock import patch
@@ -10,7 +10,9 @@
 import source_stripe
 import stripe
 from airbyte_cdk.models import ConfiguredAirbyteCatalog, SyncMode
+from airbyte_cdk.sources.streams.call_rate import CachedLimiterSession, LimiterSession, Rate
 from airbyte_cdk.sources.streams.concurrent.adapters import StreamFacade
+from airbyte_cdk.sources.streams.http import HttpStream
 from airbyte_cdk.utils import AirbyteTracedException
 from source_stripe import SourceStripe
 
@@ -92,3 +94,57 @@ def test_when_streams_return_full_refresh_as_concurrent():
     ).streams(_a_valid_config())
 
     assert len(list(filter(lambda stream: isinstance(stream, StreamFacade), streams))) == 1
+
+
+@pytest.mark.parametrize(
+    "input_config, default_call_limit",
+    (
+        ({"account_id": 1, "client_secret": "secret"}, 100),
+        ({"account_id": 1, "client_secret": "secret", "call_rate_limit": 10}, 10),
+        ({"account_id": 1, "client_secret": "secret", "call_rate_limit": 110}, 100),
+        ({"account_id": 1, "client_secret": "sk_test_some_secret"}, 25),
+        ({"account_id": 1, "client_secret": "sk_test_some_secret", "call_rate_limit": 10}, 10),
+        ({"account_id": 1, "client_secret": "sk_test_some_secret", "call_rate_limit": 30}, 25),
+    ),
+)
+def test_call_budget_creation(mocker, input_config, default_call_limit):
+    """Test that call_budget was created with specific config i.e., that first policy has specific matchers."""
+
+    policy_mock = mocker.patch("source_stripe.source.MovingWindowCallRatePolicy")
+    matcher_mock = mocker.patch("source_stripe.source.HttpRequestMatcher")
+    source = SourceStripe()
+
+    source.get_api_call_budget(input_config)
+
+    policy_mock.assert_has_calls(
+        calls=[
+            mocker.call(matchers=[mocker.ANY, mocker.ANY], rates=[Rate(limit=20, interval=datetime.timedelta(seconds=1))]),
+            mocker.call(matchers=[], rates=[Rate(limit=default_call_limit, interval=datetime.timedelta(seconds=1))]),
+        ],
+    )
+
+    matcher_mock.assert_has_calls(
+        calls=[
+            mocker.call(url="https://api.stripe.com/v1/files"),
+            mocker.call(url="https://api.stripe.com/v1/file_links"),
+        ]
+    )
+
+
+def test_call_budget_passed_to_every_stream(mocker):
+    """Test that each stream has call_budget passed and creates a proper session"""
+
+    prod_config = {"account_id": 1, "client_secret": "secret"}
+    source = SourceStripe()
+    get_api_call_budget_mock = mocker.patch.object(source, "get_api_call_budget")
+
+    streams = source.streams(prod_config)
+
+    assert streams
+    get_api_call_budget_mock.assert_called_once()
+
+    for stream in streams:
+        assert isinstance(stream, HttpStream)
+        session = stream.request_session()
+        assert isinstance(session, (CachedLimiterSession, LimiterSession))
+        assert session._api_budget == get_api_call_budget_mock.return_value