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_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_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_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 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(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 _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_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_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_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)
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 six.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)