Replaced ratelimit and tenacity dependencies with simpler implementations (#195)

This PR simplifies the handling of rate-limiting to a single
`@rate_limited` decorator that could be added to methods requiring
throttling.

Author: nfx

pyproject.toml

@@ -28,10 +28,6 @@ classifiers = [
 dependencies = [
     "databricks-sdk~=0.8.0",
     "PyYAML>=6.0.0,<7.0.0",
-
-    # TODO: remove later
-    "ratelimit>=2.2.1,<3.0.0",
-    "tenacity>=8.2.2,<9.0.0",
 ]
 
 [project.optional-dependencies]

src/databricks/labs/ucx/config.py

@@ -41,6 +41,8 @@ class ConnectConfig:
     profile: str | None = None
     debug_headers: bool | None = False
     rate_limit: int | None = None
+    max_connections_per_pool: int | None = None
+    max_connection_pools: int | None = None
 
     @staticmethod
     def from_databricks_config(cfg: Config) -> "ConnectConfig":
@@ -57,6 +59,8 @@ def from_databricks_config(cfg: Config) -> "ConnectConfig":
             profile=cfg.profile,
             debug_headers=cfg.debug_headers,
             rate_limit=cfg.rate_limit,
+            max_connection_pools=cfg.max_connection_pools,
+            max_connections_per_pool=cfg.max_connections_per_pool,
         )
 
     @classmethod
@@ -136,7 +140,7 @@ def from_dict(cls, raw: dict) -> "MigrationConfig":
             groups=GroupsConfig.from_dict(raw.get("groups", {})),
             connect=ConnectConfig.from_dict(raw.get("connect", {})),
             instance_pool_id=raw.get("instance_pool_id", None),
-            num_threads=raw.get("num_threads", 4),
+            num_threads=raw.get("num_threads", 8),
             log_level=raw.get("log_level", "INFO"),
         )
 
@@ -170,6 +174,8 @@ def to_databricks_config(self) -> Config:
             profile=connect.profile,
             debug_headers=connect.debug_headers,
             rate_limit=connect.rate_limit,
+            max_connection_pools=connect.max_connection_pools,
+            max_connections_per_pool=connect.max_connections_per_pool,
             product="ucx",
             product_version=__version__,
         )

src/databricks/labs/ucx/managers/group.py

@@ -5,14 +5,14 @@
 
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.service import iam
-from ratelimit import limits, sleep_and_retry
 
 from databricks.labs.ucx.config import GroupsConfig
 from databricks.labs.ucx.generic import StrEnum
 from databricks.labs.ucx.providers.groups_info import (
     GroupMigrationState,
     MigrationGroupInfo,
 )
+from databricks.labs.ucx.providers.mixins.hardening import rate_limited
 from databricks.labs.ucx.utils import ThreadedExecution
 
 logger = logging.getLogger(__name__)
@@ -117,8 +117,7 @@ def _replace_group(self, migration_info: MigrationGroupInfo):
 
         self._reflect_account_group_to_workspace(migration_info.account)
 
-    @sleep_and_retry
-    @limits(calls=5, period=1)  # assumption
+    @rate_limited(max_requests=5)  # assumption
     def _reflect_account_group_to_workspace(self, acc_group: iam.Group) -> None:
         logger.info(f"Reflecting group {acc_group.display_name} to workspace")
 

src/databricks/labs/ucx/providers/mixins/hardening.py

@@ -0,0 +1,61 @@
+import logging
+import threading
+import time
+from functools import wraps
+
+logger = logging.getLogger(__name__)
+
+
+class RateLimiter:
+    """Used to limit the rate of requests to a certain maximum rate within a specified burst period,
+    defaulting to one second. It is nice to control and manage the flow of requests to prevent
+    overloading a service.
+
+    It uses a variation of the Leaky Bucket algorithm (https://en.wikipedia.org/wiki/Leaky_bucket)
+    for its simplicity.
+    """
+
+    def __init__(self, *, max_requests: int = 30, burst_period_seconds: int = 1):
+        self._capacity = max_requests
+        self._burst_period_seconds = burst_period_seconds
+        self._lock = threading.RLock()
+        self._last = time.time()
+        self._bucket = 0
+
+    def throttle(self):
+        """The throttle method is used to check and control the rate of requests. It's intended to be called
+        before making a new request and is thread-safe.
+
+        It calculates the time elapsed since the last request and subtracts this from the _burst_period_seconds
+        to determine how much time remains in the current burst period. If the calculated delay is less than zero,
+        it means that the burst period has elapsed, so it resets the request count and updates the last request
+        timestamp. If the request count exceeds the allowed maximum (_max_requests), it sleeps for the remaining
+        time in the burst period (delay) to enforce the rate limit.
+        """
+        with self._lock:
+            now = time.time()
+            delay = self._burst_period_seconds - (now - self._last)
+            if delay < 0:
+                # If the bucket is empty, it stops leaking.
+                self._bucket = 0
+                self._last = now
+            self._bucket += 1
+            if self._bucket > self._capacity:
+                # If the bucket is over capacity, start leaking
+                # by sleeping on the current thread.
+                logger.debug(f"Throttled for {delay}s")
+                time.sleep(delay)
+
+
+def rate_limited(*, max_requests: int = 30, burst_period_seconds: int = 1):
+    def decorator(func):
+        rate_limiter = RateLimiter(max_requests=max_requests, burst_period_seconds=burst_period_seconds)
+
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            rate_limiter.throttle()
+            return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator

src/databricks/labs/ucx/support/group_level.py

@@ -2,10 +2,10 @@
 from functools import partial
 
 from databricks.sdk.service import iam
-from ratelimit import limits, sleep_and_retry
 
 from databricks.labs.ucx.inventory.types import Destination, PermissionsInventoryItem
 from databricks.labs.ucx.providers.groups_info import GroupMigrationState
+from databricks.labs.ucx.providers.mixins.hardening import rate_limited
 from databricks.labs.ucx.support.base import BaseSupport
 
 
@@ -17,8 +17,7 @@ def _crawler_task(self, group: iam.Group, property_name: str):
             raw_object_permissions=json.dumps([e.as_dict() for e in getattr(group, property_name)]),
         )
 
-    @sleep_and_retry
-    @limits(calls=10, period=1)
+    @rate_limited(max_requests=10)
     def _applier_task(self, group_id: str, value: list[iam.ComplexValue], property_name: str):
         operations = [iam.Patch(op=iam.PatchOp.ADD, path=property_name, value=[e.as_dict() for e in value])]
         schemas = [iam.PatchSchema.URN_IETF_PARAMS_SCIM_API_MESSAGES_2_0_PATCH_OP]

src/databricks/labs/ucx/support/listing.py

@@ -7,9 +7,9 @@
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.service import ml, workspace
 from databricks.sdk.service.workspace import ObjectInfo, ObjectType
-from ratelimit import limits, sleep_and_retry
 
 from databricks.labs.ucx.inventory.types import RequestObjectType
+from databricks.labs.ucx.providers.mixins.hardening import rate_limited
 from databricks.labs.ucx.support.permissions import GenericPermissionsInfo
 
 logger = logging.getLogger(__name__)
@@ -44,8 +44,7 @@ def _progress_report(self, _):
                 f" rps: {rps:.3f}/sec"
             )
 
-    @sleep_and_retry
-    @limits(calls=45, period=1)  # safety value, can be 50 actually
+    @rate_limited(max_requests=45)  # safety value, can be 50 actually
     def _list_workspace(self, path: str) -> Iterator[ObjectType]:
         # TODO: remove, use SDK
         return self._ws.workspace.list(path=path, recursive=False)

src/databricks/labs/ucx/support/permissions.py

@@ -6,14 +6,14 @@
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.core import DatabricksError
 from databricks.sdk.service import iam
-from ratelimit import limits, sleep_and_retry
 
 from databricks.labs.ucx.inventory.types import (
     Destination,
     PermissionsInventoryItem,
     RequestObjectType,
 )
 from databricks.labs.ucx.providers.groups_info import GroupMigrationState
+from databricks.labs.ucx.providers.mixins.hardening import rate_limited
 from databricks.labs.ucx.support.base import BaseSupport, logger
 
 
@@ -79,15 +79,13 @@ def _prepare_new_acl(
 
         return acl_requests
 
-    @sleep_and_retry
-    @limits(calls=30, period=1)
+    @rate_limited(max_requests=30)
     def _applier_task(
         self, ws: WorkspaceClient, object_id: str, acl: list[iam.AccessControlRequest], request_type: RequestObjectType
     ):
         ws.permissions.update(request_object_type=request_type, request_object_id=object_id, access_control_list=acl)
 
-    @sleep_and_retry
-    @limits(calls=100, period=1)
+    @rate_limited(max_requests=100)
     def _crawler_task(
         self,
         ws: WorkspaceClient,

src/databricks/labs/ucx/support/secrets.py

@@ -5,10 +5,10 @@
 
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.service import iam, workspace
-from ratelimit import limits, sleep_and_retry
 
 from databricks.labs.ucx.inventory.types import Destination, PermissionsInventoryItem
 from databricks.labs.ucx.providers.groups_info import GroupMigrationState
+from databricks.labs.ucx.providers.mixins.hardening import rate_limited
 from databricks.labs.ucx.support.base import BaseSupport
 
 
@@ -66,8 +66,7 @@ def _inflight_check(
         msg = f"Failed to apply permissions for {group_name} on scope {scope_name} in {num_retries} retries"
         raise ValueError(msg)
 
-    @sleep_and_retry
-    @limits(calls=30, period=1)
+    @rate_limited(max_requests=30)
     def _rate_limited_put_acl(self, object_id: str, principal: str, permission: workspace.AclPermission):
         self._ws.secrets.put_acl(object_id, principal, permission)
         self._inflight_check(scope_name=object_id, group_name=principal, expected_permission=permission)

src/databricks/labs/ucx/support/sql.py

@@ -7,10 +7,10 @@
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.core import DatabricksError
 from databricks.sdk.service import iam, sql
-from ratelimit import limits, sleep_and_retry
 
 from databricks.labs.ucx.inventory.types import Destination, PermissionsInventoryItem
 from databricks.labs.ucx.providers.groups_info import GroupMigrationState
+from databricks.labs.ucx.providers.mixins.hardening import rate_limited
 from databricks.labs.ucx.support.base import BaseSupport, logger
 
 
@@ -46,8 +46,7 @@ def _safe_get_dbsql_permissions(self, object_type: sql.ObjectTypePlural, object_
             else:
                 raise e
 
-    @sleep_and_retry
-    @limits(calls=100, period=1)
+    @rate_limited(max_requests=100)
     def _crawler_task(self, object_id: str, object_type: sql.ObjectTypePlural) -> PermissionsInventoryItem | None:
         permissions = self._safe_get_dbsql_permissions(object_type=object_type, object_id=object_id)
         if permissions:
@@ -57,8 +56,7 @@ def _crawler_task(self, object_id: str, object_type: sql.ObjectTypePlural) -> Pe
                 raw_object_permissions=json.dumps(permissions.as_dict()),
             )
 
-    @sleep_and_retry
-    @limits(calls=30, period=1)
+    @rate_limited(max_requests=30)
     def _applier_task(self, object_type: sql.ObjectTypePlural, object_id: str, acl: list[sql.AccessControl]):
         """
         Please note that we only have SET option (DBSQL Permissions API doesn't support UPDATE operation).

tests/unit/test_ratelimit.py

@@ -0,0 +1,25 @@
+import pytest
+
+from databricks.labs.ucx.providers.mixins.hardening import rate_limited
+
+
+def test_ratelimiting(mocker):
+    _x = 0
+
+    def _now():
+        nonlocal _x
+        _x += 0.2
+        return _x
+
+    mocker.patch("time.time", side_effect=_now)
+    sleep = mocker.patch("time.sleep")
+
+    @rate_limited(max_requests=5)
+    def try_something():
+        return 1
+
+    for _ in range(0, 20):
+        try_something()
+
+    sleep.assert_called()
+    assert 4.2 == pytest.approx(_x)