def get_limiter(self, table_descriptions): """ Construct a RateLimit object from the throttle declarations """ table_caps = {} for table in table_descriptions: limit = self.tables.get(table.name) or self.default # Add the table limit if limit: table_caps[table.name] = { "read": self._compute_limit(limit["read"], table.read_throughput), "write": self._compute_limit(limit["write"], table.write_throughput), } if table.name not in self.indexes: continue # Add the global index limits for index in itervalues(table.global_indexes): limit = self.indexes[table.name].get( index.name) or self.default if limit: cap = table_caps.setdefault(table.name, {}) cap[index.name] = { "read": self._compute_limit(limit["read"], index.read_throughput), "write": self._compute_limit(limit["write"], index.write_throughput), } kwargs = {"table_caps": table_caps} if self.total: kwargs["total_read"] = float(self.total["read"]) kwargs["total_write"] = float(self.total["write"]) return RateLimit(**kwargs)
def test_throttle_table_default(self): """If no table limit provided, use the default""" limiter = RateLimit(default_read=4, default_write=4) cap = ConsumedCapacity("foobar", Capacity(8, 0), Capacity(8, 0)) with self.inject_capacity(cap, limiter) as sleep: list(self.dynamo.query("foobar", "id = :id", id="a")) sleep.assert_called_with(2)
def test_throttle_total_cap(self): """Sleep if consumed capacity exceeds total""" limiter = RateLimit(total=Capacity(3, 3)) cap = ConsumedCapacity("foobar", Capacity(3, 0)) with self.inject_capacity(cap, limiter) as sleep: list(self.dynamo.query("foobar", "id = :id", id="a")) sleep.assert_called_with(1)
def test_no_throttle(self): """Don't sleep if consumed capacity is within limits""" limiter = RateLimit(3, 3) cap = ConsumedCapacity("foobar", Capacity(0, 2)) with self.inject_capacity(cap, limiter) as sleep: list(self.dynamo.query("foobar", "id = :id", id="a")) sleep.assert_not_called()
def test_throttle_table_default_cap(self): """ If no table limit provided, use the default """ limiter = RateLimit(default=Capacity(4, 4)) cap = ConsumedCapacity('foobar', Capacity(8, 0), Capacity(8, 0)) with self.inject_capacity(cap, limiter) as sleep: list(self.dynamo.query2('foobar', 'id = :id', id='a')) sleep.assert_called_with(2)
def test_throttle_multiply(self): """Seconds to sleep is increades to match limit delta""" limiter = RateLimit(3, 3) cap = ConsumedCapacity("foobar", Capacity(8, 0)) with self.inject_capacity(cap, limiter) as sleep: list(self.dynamo.query("foobar", "id = :id", id="a")) sleep.assert_called_with(3)
def test_throttle_total(self): """ Sleep if consumed capacity exceeds total """ limiter = RateLimit(3, 3) cap = ConsumedCapacity('foobar', Capacity(3, 0)) with self.inject_capacity(cap, limiter) as sleep: list(self.dynamo.query2('foobar', 'id = :id', id='a')) sleep.assert_called_with(1)
def test_throttle_multiple(self): """Sleep if the limit is exceeded by multiple calls""" limiter = RateLimit(4, 4) cap = ConsumedCapacity("foobar", Capacity(3, 0)) with self.inject_capacity(cap, limiter) as sleep: list(self.dynamo.query("foobar", "id = :id", id="a")) list(self.dynamo.query("foobar", "id = :id", id="a")) sleep.assert_called_with(2)
def test_throttle_table(self): """ Sleep if table limit is exceeded """ limiter = RateLimit(3, 3, table_caps={ 'foobar': Capacity(0, 4), }) cap = ConsumedCapacity('foobar', Capacity(8, 0), Capacity(0, 8)) with self.inject_capacity(cap, limiter) as sleep: list(self.dynamo.query2('foobar', 'id = :id', id='a')) sleep.assert_called_with(2)
def test_throttle_callback(self): """Callback is called when a query is throttled""" callback = MagicMock() callback.return_value = True limiter = RateLimit(3, 3, callback=callback) cap = ConsumedCapacity("foobar", Capacity(3, 0)) with self.inject_capacity(cap, limiter) as sleep: list(self.dynamo.query("foobar", "id = :id", id="a")) sleep.assert_not_called() self.assertTrue(callback.called)
def test_global_default(self): """ Global index limit will fall back to table default limit """ limiter = RateLimit(default_read=4, default_write=4) cap = ConsumedCapacity('foobar', Capacity(8, 0), global_index_capacity={ 'baz': Capacity(8, 0), }) with self.inject_capacity(cap, limiter) as sleep: list(self.dynamo.query2('foobar', 'id = :id', id='a')) sleep.assert_called_with(2)
def test_global_default(self): """Global index limit will fall back to table default limit""" limiter = RateLimit(default_read=4, default_write=4) cap = ConsumedCapacity( "foobar", Capacity(8, 0), global_index_capacity={ "baz": Capacity(8, 0), }, ) with self.inject_capacity(cap, limiter) as sleep: list(self.dynamo.query("foobar", "id = :id", id="a")) sleep.assert_called_with(2)
def test_global_index(self): """ Sleep when global index limit is exceeded """ limiter = RateLimit(table_caps={'foobar': { 'baz': Capacity(4, 0), }}) cap = ConsumedCapacity('foobar', Capacity(8, 0), global_index_capacity={ 'baz': Capacity(8, 0), }) with self.inject_capacity(cap, limiter) as sleep: list(self.dynamo.query2('foobar', 'id = :id', id='a')) sleep.assert_called_with(2)
def test_throttle_table(self): """Sleep if table limit is exceeded""" limiter = RateLimit( 3, 3, table_caps={ "foobar": Capacity(0, 4), }, ) cap = ConsumedCapacity("foobar", Capacity(8, 0), Capacity(0, 8)) with self.inject_capacity(cap, limiter) as sleep: list(self.dynamo.query("foobar", "id = :id", id="a")) sleep.assert_called_with(2)
def test_global_default_table(self): """ Global index limit defaults to table limit if not present """ limiter = RateLimit(table_caps={ 'foobar': Capacity(4, 0), }) cap = ConsumedCapacity('foobar', Capacity(8, 0), global_index_capacity={ 'baz': Capacity(8, 0), }) with self.inject_capacity(cap, limiter) as sleep: list(self.dynamo.query2('foobar', 'id = :id', id='a')) sleep.assert_called_with(2)
def test_global_index_by_name(self): """ Global index limit can be specified as tablename:index_name """ limiter = RateLimit(table_caps={ 'foobar:baz': Capacity(4, 0), }) cap = ConsumedCapacity('foobar', Capacity(8, 0), global_index_capacity={ 'baz': Capacity(8, 0), }) with self.inject_capacity(cap, limiter) as sleep: list(self.dynamo.query2('foobar', 'id = :id', id='a')) sleep.assert_called_with(2)
def test_local_index(self): """ Local index capacities count towards the table limit """ limiter = RateLimit(table_caps={ 'foobar': Capacity(4, 0), }) cap = ConsumedCapacity('foobar', Capacity(8, 0), local_index_capacity={ 'local': Capacity(4, 0), }) with self.inject_capacity(cap, limiter) as sleep: list(self.dynamo.query2('foobar', 'id = :id', id='a')) sleep.assert_called_with(1)
def _parse_throttle(self, tablename, throttle): """ Parse a 'throttle' statement and return a RateLimit """ amount = [] desc = self.describe(tablename) throughputs = [desc.read_throughput, desc.write_throughput] for value, throughput in zip(throttle[1:], throughputs): if value == "*": amount.append(0) elif value[-1] == "%": amount.append(throughput * float(value[:-1]) / 100.0) else: amount.append(float(value)) cap = Capacity(*amount) # pylint: disable=E1120 return RateLimit(total=cap, callback=self._on_throttle)
def test_global_index_by_name(self): """Global index limit can be specified as tablename:index_name""" limiter = RateLimit(table_caps={ "foobar:baz": Capacity(4, 0), }) cap = ConsumedCapacity( "foobar", Capacity(8, 0), global_index_capacity={ "baz": Capacity(8, 0), }, ) with self.inject_capacity(cap, limiter) as sleep: list(self.dynamo.query("foobar", "id = :id", id="a")) sleep.assert_called_with(2)
def test_global_index(self): """Sleep when global index limit is exceeded""" limiter = RateLimit(table_caps={"foobar": { "baz": Capacity(4, 0), }}) cap = ConsumedCapacity( "foobar", Capacity(8, 0), global_index_capacity={ "baz": Capacity(8, 0), }, ) with self.inject_capacity(cap, limiter) as sleep: list(self.dynamo.query("foobar", "id = :id", id="a")) sleep.assert_called_with(2)
def test_global_default_table(self): """Global index limit defaults to table limit if not present""" limiter = RateLimit(table_caps={ "foobar": Capacity(4, 0), }) cap = ConsumedCapacity( "foobar", Capacity(8, 0), global_index_capacity={ "baz": Capacity(8, 0), }, ) with self.inject_capacity(cap, limiter) as sleep: list(self.dynamo.query("foobar", "id = :id", id="a")) sleep.assert_called_with(2)
def test_local_index(self): """Local index capacities count towards the table limit""" limiter = RateLimit(table_caps={ "foobar": Capacity(4, 0), }) cap = ConsumedCapacity( "foobar", Capacity(8, 0), local_index_capacity={ "local": Capacity(4, 0), }, ) with self.inject_capacity(cap, limiter) as sleep: list(self.dynamo.query("foobar", "id = :id", id="a")) sleep.assert_called_with(1)
def _parse_throttle(self, tablename: str, throttle: Any) -> RateLimit: """Parse a 'throttle' statement and return a RateLimit""" amount: List[float] = [] desc = self.describe(tablename) if desc.throughput is None: throughputs = (0, 0) else: throughputs = (desc.throughput.read, desc.throughput.write) for value, throughput in zip(throttle[1:], throughputs): if value == "*": amount.append(0) elif value[-1] == "%": amount.append(throughput * float(value[:-1]) / 100.0) else: amount.append(float(value)) cap = Capacity(*amount) # pylint: disable=E1120 return RateLimit(total=cap, callback=self._on_throttle)