Rate limited Page Iterator (pynamodb#481)

betamoo · harleyk · commit 8ad3ad1ad963 · 2018-04-04T10:40:36.000-07:00
diff --git a/docs/rate_limited_operations.rst b/docs/rate_limited_operations.rst
@@ -0,0 +1,65 @@
+Rate-Limited Operation
+================
+
+`Scan`, `Query` and `Count` operations can be rate-limited based on the consumed capacities returned from DynamoDB.
+Simply specify the `rate_limit` argument when calling these methods.
+
+.. note::
+
+    Rate-limiting is only meant to slow operations down to conform to capacity limitations.
+    Rate-limiting can not be used to speed operations up. Specifying a higher rate-limit that exceeds the possible
+    writing speed allowed by the environment, will not have any effect.
+
+Example Usage
+^^^^^^^^^^^^
+
+Suppose that you have defined a `User` Model for the examples below.
+
+.. code-block:: python
+
+    from pynamodb.models import Model
+    from pynamodb.attributes import (
+        UnicodeAttribute
+    )
+
+
+    class User(Model):
+        class Meta:
+            table_name = 'Users'
+
+        id = UnicodeAttribute(hash_key=True)
+        name = UnicodeAttribute(range_key=True)
+
+
+Here is an example using `rate-limit` in while scaning the `User` model
+
+.. code-block:: python
+
+    # Using only 5 RCU per second
+    for user in User.scan(rate_limit = 5):
+        print("User id: {}, name: {}".format(user.id, user.name))
+
+
+Query
+^^^^^^^^^^^^^
+
+You can use `rate-limit` when querying items from your table:
+
+.. code-block:: python
+
+    # Using only 15 RCU per second
+    for user in User.query('id1', User.name.startswith('re'), rate_limit = 15):
+        print("Query returned user {0}".format(user))
+
+
+Count
+^^^^^^^^^^^^^
+
+You can use `rate-limit` when counting items in your table:
+
+.. code-block:: python
+
+    # Using only 15 RCU per second
+    count = User.count(rate_limit = 15):
+    print("Count : {}".format(count))
+
diff --git a/pynamodb/models.py b/pynamodb/models.py
@@ -536,6 +536,7 @@ def count(cls,
               consistent_read=False,
               index_name=None,
               limit=None,
+              rate_limit=None,
               **filters):
         """
         Provides a filtered count
@@ -589,7 +590,8 @@ def count(cls,
             cls._get_connection().query,
             query_args,
             query_kwargs,
-            limit=limit
+            limit=limit,
+            rate_limit=rate_limit
         )
 
         # iterate through results
@@ -610,6 +612,7 @@ def query(cls,
               last_evaluated_key=None,
               attributes_to_get=None,
               page_size=None,
+              rate_limit=None,
               **filters):
         """
         Provides a high level query API
@@ -674,7 +677,8 @@ def query(cls,
             query_args,
             query_kwargs,
             map_fn=cls.from_raw_data,
-            limit=limit
+            limit=limit,
+            rate_limit=rate_limit
         )
 
     @classmethod
@@ -763,6 +767,7 @@ def scan(cls,
              page_size=None,
              consistent_read=None,
              index_name=None,
+             rate_limit=None,
              **filters):
         """
         Iterates through all items in the table
@@ -807,7 +812,8 @@ def scan(cls,
             scan_args,
             scan_kwargs,
             map_fn=cls.from_raw_data,
-            limit=limit
+            limit=limit,
+            rate_limit=rate_limit,
         )
 
     @classmethod
diff --git a/pynamodb/pagination.py b/pynamodb/pagination.py
@@ -1,4 +1,70 @@
-from pynamodb.constants import CAMEL_COUNT, ITEMS, LAST_EVALUATED_KEY, SCANNED_COUNT
+import time
+from pynamodb.constants import (CAMEL_COUNT, ITEMS, LAST_EVALUATED_KEY, SCANNED_COUNT,
+                                CONSUMED_CAPACITY, TOTAL, CAPACITY_UNITS)
+
+
+class RateLimiter(object):
+    """
+    RateLimiter limits operations to a pre-set rate of units/seconds
+
+    Example:
+        Initialize a RateLimiter with the desired rate
+            rate_limiter = RateLimiter(rate_limit)
+
+        Now, every time before calling an operation, call acquire()
+            rate_limiter.acquire()
+
+        And after an operation, update the number of units consumed
+            rate_limiter.consume(units)
+
+    """
+    def __init__(self, rate_limit, time_module = None):
+        """
+        Initializes a RateLimiter object
+
+        :param rate_limit: The desired rate
+        :param time_module: Optional: the module responsible for calculating time. Intended to be used for testing purposes.
+        """
+        if rate_limit <= 0:
+            raise ValueError("rate_limit must be greater than zero")
+        self._rate_limit = rate_limit
+        self._consumed = 0
+        self._time_of_last_acquire = 0.0
+        self._time_module = time_module or time
+
+    def consume(self, units):
+        """
+        Records the amount of units consumed.
+
+        :param units: Number of units consumed
+
+        :return: None
+        """
+        self._consumed += units
+
+    def acquire(self):
+        """
+        Sleeps the appropriate amount of time to follow the rate limit restriction
+
+        :return: None
+        """
+
+        self._time_module.sleep(max(0, self._consumed/float(self.rate_limit) - (self._time_module.time()-self._time_of_last_acquire)))
+        self._consumed = 0
+        self._time_of_last_acquire = self._time_module.time()
+
+    @property
+    def rate_limit(self):
+        """
+        A limit of units per seconds
+        """
+        return self._rate_limit
+
+    @rate_limit.setter
+    def rate_limit(self, rate_limit):
+        if rate_limit <= 0:
+            raise ValueError("rate_limit must be greater than zero")
+        self._rate_limit = rate_limit
 
 
 class PageIterator(object):
@@ -8,13 +74,16 @@ class PageIterator(object):
     http://docs.aws.amazon.com/amazondynamodb/latest/developerguide/Query.html#Query.Pagination
     http://docs.aws.amazon.com/amazondynamodb/latest/developerguide/Scan.html#Scan.Pagination
     """
-    def __init__(self, operation, args, kwargs):
+    def __init__(self, operation, args, kwargs, rate_limit = None):
         self._operation = operation
         self._args = args
         self._kwargs = kwargs
         self._first_iteration = True
         self._last_evaluated_key = kwargs.get('exclusive_start_key')
         self._total_scanned_count = 0
+        self._rate_limiter = None
+        if rate_limit:
+            self._rate_limiter = RateLimiter(rate_limit)
 
     def __iter__(self):
         return self
@@ -26,10 +95,18 @@ def __next__(self):
         self._first_iteration = False
 
         self._kwargs['exclusive_start_key'] = self._last_evaluated_key
+
+        if self._rate_limiter:
+            self._rate_limiter.acquire()
+            self._kwargs['return_consumed_capacity'] = TOTAL
         page = self._operation(*self._args, **self._kwargs)
         self._last_evaluated_key = page.get(LAST_EVALUATED_KEY)
         self._total_scanned_count += page[SCANNED_COUNT]
 
+        if self._rate_limiter:
+            consumed_capacity = page.get(CONSUMED_CAPACITY, {}).get(CAPACITY_UNITS, 0)
+            self._rate_limiter.consume(consumed_capacity)
+
         return page
 
     def next(self):
@@ -69,8 +146,8 @@ class ResultIterator(object):
     http://docs.aws.amazon.com/amazondynamodb/latest/developerguide/Query.html#Query.Pagination
     http://docs.aws.amazon.com/amazondynamodb/latest/developerguide/Scan.html#Scan.Pagination
     """
-    def __init__(self, operation, args, kwargs, map_fn=None, limit=None):
-        self.page_iter = PageIterator(operation, args, kwargs)
+    def __init__(self, operation, args, kwargs, map_fn=None, limit=None, rate_limit = None):
+        self.page_iter = PageIterator(operation, args, kwargs, rate_limit)
         self._first_iteration = True
         self._map_fn = map_fn
         self._limit = limit
diff --git a/pynamodb/tests/test_pagination.py b/pynamodb/tests/test_pagination.py
@@ -0,0 +1,79 @@
+import pytest
+from pynamodb.pagination import RateLimiter
+
+
+class MockTime():
+    def __init__(self):
+        self.current_time = 0.0
+
+    def sleep(self, amount):
+        self.current_time += amount
+
+    def time(self):
+        return self.current_time
+
+    def increment_time(self, amount):
+        self.current_time += amount
+
+
+def test_rate_limiter_exceptions():
+    with pytest.raises(ValueError):
+        r = RateLimiter(0)
+
+    with pytest.raises(ValueError):
+        r = RateLimiter(-1)
+
+    with pytest.raises(ValueError):
+        r = RateLimiter(10)
+        r.rate_limit = 0
+
+    with pytest.raises(ValueError):
+        r = RateLimiter(10)
+        r.rate_limit = -1
+
+
+def test_basic_rate_limiting():
+    mock_time = MockTime()
+    r = RateLimiter(0.1, mock_time)
+
+    # 100 operations
+    for i in range(0, 100):
+        r.acquire()
+        # Simulates an operation that takes 1 second
+        mock_time.increment_time(1)
+        r.consume(1)
+
+    # Since the first acquire doesn't take time, thus we should be expecting (100-1) * 10 seconds = 990 delay
+    # plus 1 for the last increment_time(1) operation
+    assert mock_time.time() == 991.0
+
+
+def test_basic_rate_limiting_small_increment():
+    mock_time = MockTime()
+    r = RateLimiter(0.1, mock_time)
+
+    # 100 operations
+    for i in range(0, 100):
+        r.acquire()
+        # Simulates an operation that takes 2 second
+        mock_time.increment_time(2)
+        r.consume(1)
+
+    # Since the first acquire doesn't take time, thus we should be expecting (100-1) * 10 seconds = 990 delay
+    # plus 2 for the last increment_time(2) operation
+    assert mock_time.time() == 992.0
+
+
+def test_basic_rate_limiting_large_increment():
+    mock_time = MockTime()
+    r = RateLimiter(0.1, mock_time)
+
+    # 100 operations
+    for i in range(0, 100):
+        r.acquire()
+        # Simulates an operation that takes 2 second
+        mock_time.increment_time(11)
+        r.consume(1)
+
+    # The operation takes longer than the minimum wait, so rate limiting should have no effect
+    assert mock_time.time() == 1100.0
