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)
