feat: Handle LaunchDarkly rate limiting (#5501)

Author: khvn26

api/app/settings/common.py

@@ -1147,6 +1147,10 @@
 RECURRING_TASK_RUN_RETENTION_DAYS = env.int(
     "RECURRING_TASK_RUN_RETENTION_DAYS", default=30
 )
+TASK_BACKOFF_DEFAULT_DELAY_SECONDS = env.int(
+    "TASK_BACKOFF_DEFAULT_DELAY_SECONDS",
+    default=5,
+)
 
 # Webhook settings
 DISABLE_WEBHOOKS = env.bool("DISABLE_WEBHOOKS", False)

api/integrations/launch_darkly/client.py

@@ -1,17 +1,78 @@
-from typing import Any, Iterator, Optional, TypeVar
+from datetime import timedelta
+from typing import Any, Callable, Generator, Iterator, Optional, TypeVar
 
-from requests import Session
+import backoff
+from backoff.types import Details
+from django.utils.timezone import now as timezone_now
+from requests import RequestException, Session
+from rest_framework.status import HTTP_429_TOO_MANY_REQUESTS
 
 from integrations.launch_darkly import types as ld_types
 from integrations.launch_darkly.constants import (
+    BACKOFF_MAX_RETRIES,
     LAUNCH_DARKLY_API_BASE_URL,
     LAUNCH_DARKLY_API_ITEM_COUNT_LIMIT_PER_PAGE,
     LAUNCH_DARKLY_API_VERSION,
 )
+from integrations.launch_darkly.exceptions import LaunchDarklyRateLimitError
 
 T = TypeVar("T")
 
 
+def launch_darkly_backoff(
+    _get_json_response: Callable[..., T],
+) -> Callable[..., T]:
+    # Handle LaunchDarkly rate limiting according to their API documentation:
+    # https://launchdarkly.com/docs/api
+    #
+    # 1. If a request returns a 429 Too Many Requests status code, the client should back off.
+    # 2. When backing off, the client retries the request after the time specified in the `Retry-After` header
+    #   or the `X-Ratelimit-Reset` header, if present, or a default of `BACKOFF_DEFAULT_RETRY_SECONDS`.
+    # 3. After `BACKOFF_MAX_RETRIES` retries, we give up.
+    #   If the last error was a 429 and contained retry information,
+    #   signal for the import request to be retried later by raising a `LaunchDarklyRateLimitError`.
+
+    def _get_retry_after(exc: RequestException) -> float | None:
+        if (
+            (response := exc.response) is not None
+        ) and response.status_code == HTTP_429_TOO_MANY_REQUESTS:
+            headers = response.headers
+            if retry_after := headers.get("Retry-After"):
+                return float(retry_after)
+            if ratelimit_reset := headers.get("X-Ratelimit-Reset"):
+                return float(ratelimit_reset) - timezone_now().timestamp()
+        return None
+
+    def _wait_gen() -> Generator[float, None, None]:
+        exc: RequestException = yield  # type: ignore[misc,assignment]
+
+        while True:
+            if retry_after := _get_retry_after(exc):
+                yield retry_after
+            else:
+                return
+
+    def _handle_giveup(
+        details: Details,
+    ) -> None:
+        exc: RequestException = details["exception"]  # type: ignore[typeddict-item]
+
+        if retry_after := _get_retry_after(exc):
+            raise LaunchDarklyRateLimitError(
+                retry_at=timezone_now() + timedelta(seconds=retry_after)
+            )
+
+        raise exc
+
+    return backoff.on_exception(
+        wait_gen=_wait_gen,
+        exception=RequestException,
+        jitter=backoff.random_jitter,
+        on_giveup=_handle_giveup,
+        max_tries=BACKOFF_MAX_RETRIES,
+    )(_get_json_response)
+
+
 class LaunchDarklyClient:
     def __init__(self, token: str) -> None:
         client_session = Session()
@@ -23,6 +84,7 @@ def __init__(self, token: str) -> None:
         )
         self.client_session = client_session
 
+    @launch_darkly_backoff
     def _get_json_response(
         self,
         endpoint: str,
@@ -53,7 +115,7 @@ def _iter_paginated_items(
         if additional_params:
             params.update(additional_params)
 
-        response_json = self._get_json_response(  # type: ignore[var-annotated]
+        response_json: dict[str, Any] = self._get_json_response(
             endpoint=collection_endpoint,
             params=params,
         )

api/integrations/launch_darkly/constants.py

@@ -6,3 +6,5 @@
 
 LAUNCH_DARKLY_IMPORTED_TAG_COLOR = "#3d4db6"
 LAUNCH_DARKLY_IMPORTED_DEFAULT_TAG_LABEL = "Imported"
+
+BACKOFF_MAX_RETRIES = 5

api/integrations/launch_darkly/exceptions.py

@@ -0,0 +1,13 @@
+from datetime import datetime
+
+
+class LaunchDarklyAPIError(Exception):
+    """Base exception for LaunchDarkly integration errors."""
+
+
+class LaunchDarklyRateLimitError(LaunchDarklyAPIError):
+    """Exception raised when the LaunchDarkly API rate limit is exceeded."""
+
+    def __init__(self, retry_at: datetime):
+        super().__init__()
+        self.retry_at = retry_at

api/integrations/launch_darkly/services.py

@@ -1,7 +1,7 @@
 import logging
 import re
 from contextlib import contextmanager
-from typing import Callable, Optional, Tuple
+from typing import Callable, Generator, Optional, Tuple
 
 from django.core import signing
 from django.utils import timezone
@@ -59,10 +59,10 @@ def _log_error(
     import_request.status["error_messages"] += [error_message]
 
 
-@contextmanager  # type: ignore[arg-type]
-def _complete_import_request(  # type: ignore[misc]
+@contextmanager
+def _complete_import_request(
     import_request: LaunchDarklyImportRequest,
-) -> None:
+) -> Generator[None, None, None]:
     """
     Wrap code so the import request always ends up completed.
 

api/integrations/launch_darkly/tasks.py

@@ -1,12 +1,31 @@
+import structlog
 from task_processor.decorators import (
     register_task_handler,
 )
+from task_processor.exceptions import TaskBackoffError
 
+from integrations.launch_darkly.exceptions import LaunchDarklyRateLimitError
 from integrations.launch_darkly.models import LaunchDarklyImportRequest
 from integrations.launch_darkly.services import process_import_request
 
+logger = structlog.get_logger("launch_darkly")
+
 
 @register_task_handler()
 def process_launch_darkly_import_request(import_request_id: int) -> None:
     import_request = LaunchDarklyImportRequest.objects.get(id=import_request_id)
-    process_import_request(import_request)
+    log = logger.bind(
+        import_request_id=import_request.id,
+        ld_project_key=import_request.ld_project_key,
+        organisation_id=import_request.project.organisation_id,
+        project_id=import_request.project_id,
+    )
+    try:
+        process_import_request(import_request)
+    except LaunchDarklyRateLimitError as exc:
+        log.warning(
+            "import-rate-limit-reached",
+            retry_at=exc.retry_at,
+            error_message=str(exc),
+        )
+        raise TaskBackoffError(delay_until=exc.retry_at) from exc

api/poetry.lock

@@ -160,7 +160,7 @@ pygithub = "2.1.1"
 hubspot-api-client = "^8.2.1"
 djangorestframework-dataclasses = "^1.3.1"
 pyotp = "^2.9.0"
-flagsmith-common = "^1.13.0"
+flagsmith-common = "^1.14.0"
 django-stubs = "^5.1.3"
 tzdata = "^2024.1"
 djangorestframework-simplejwt = "^5.3.1"
@@ -236,7 +236,7 @@ types-psycopg2 = "^2.9.21.20250121"
 types-python-dateutil = "^2.9.0.20241206"
 types-pytz = "^2025.1.0.20250204"
 ruff = "^0.9.7"
-flagsmith-common = { version = "^1.13.0", extras = ["test-tools"] }
+flagsmith-common = { version = "*", extras = ["test-tools"] }
 
 [build-system]
 requires = ["poetry>=2.0.0"]

api/pyproject.toml

@@ -24,7 +24,7 @@ def ld_token() -> str:
 
 @pytest.fixture
 def ld_client_mock(mocker: MockerFixture) -> MagicMock:
-    ld_client_mock = mocker.MagicMock(spec=LaunchDarklyClient)
+    ld_client_mock: MagicMock = mocker.MagicMock(spec=LaunchDarklyClient)
 
     for method_name, response_data_path in {
         "get_project": "client_responses/get_project.json",
@@ -39,7 +39,7 @@ def ld_client_mock(mocker: MockerFixture) -> MagicMock:
     ld_client_mock.get_flag_count.return_value = 9
     ld_client_mock.get_flag_tags.return_value = ["testtag", "testtag2"]
 
-    return ld_client_mock  # type: ignore[no-any-return]
+    return ld_client_mock
 
 
 @pytest.fixture