diff --git a/autopush/db.py b/autopush/db.py index 8e2589c4..946fbf76 100644 --- a/autopush/db.py +++ b/autopush/db.py @@ -163,10 +163,8 @@ def get_rotating_message_table(prefix="message", delta=0, date=None, message_write_throughput=5): # type: (str, int, Optional[datetime.date], int, int) -> Table """Gets the message table for the current month.""" - db = DynamoDBConnection() - dblist = db.list_tables()["TableNames"] tablename = make_rotating_tablename(prefix, delta, date) - if tablename not in dblist: + if not table_exists(DynamoDBConnection(), tablename): return create_rotating_message_table( prefix=prefix, delta=delta, date=date, read_throughput=message_read_throughput, @@ -220,9 +218,7 @@ def create_storage_table(tablename="storage", read_throughput=5, def _make_table(table_func, tablename, read_throughput, write_throughput): # type: (Callable[[str, int, int], Table], str, int, int) -> Table """Private common function to make a table with a table func""" - db = DynamoDBConnection() - dblist = db.list_tables()["TableNames"] - if tablename not in dblist: + if not table_exists(DynamoDBConnection(), tablename): return table_func(tablename, read_throughput, write_throughput) else: return Table(tablename) @@ -356,6 +352,23 @@ def generate_last_connect_values(date): yield int(val) +def list_tables(conn): + """Return a list of the names of all DynamoDB tables.""" + start_table = None + while True: + result = conn.list_tables(exclusive_start_table_name=start_table) + for table in result.get('TableNames', []): + yield table + start_table = result.get('LastEvaluatedTableName', None) + if not start_table: + break + + +def table_exists(conn, tablename): + """Determine if the specified Table exists""" + return any(tablename == name for name in list_tables(conn)) + + class Storage(object): """Create a Storage table abstraction on top of a DynamoDB Table object""" def __init__(self, table, metrics): diff --git a/autopush/tests/test_db.py b/autopush/tests/test_db.py index 9e41c959..2b38a365 100644 --- a/autopush/tests/test_db.py +++ b/autopush/tests/test_db.py @@ -21,6 +21,7 @@ create_router_table, create_storage_table, preflight_check, + table_exists, Storage, Message, Router, @@ -136,12 +137,9 @@ def tearDown(self): def test_custom_tablename(self): db = DynamoDBConnection() db_name = "storage_%s" % uuid.uuid4() - dblist = db.list_tables()["TableNames"] - ok_(db_name not in dblist) - + ok_(not table_exists(db, db_name)) create_storage_table(db_name) - dblist = db.list_tables()["TableNames"] - ok_(db_name in dblist) + ok_(table_exists(db, db_name)) def test_provisioning(self): db_name = "storage_%s" % uuid.uuid4() @@ -390,12 +388,9 @@ def test_drop_old_users(self): def test_custom_tablename(self): db = DynamoDBConnection() db_name = "router_%s" % uuid.uuid4() - dblist = db.list_tables()["TableNames"] - ok_(db_name not in dblist) - + ok_(not table_exists(db, db_name)) create_router_table(db_name) - dblist = db.list_tables()["TableNames"] - ok_(db_name in dblist) + ok_(table_exists(db, db_name)) def test_provisioning(self): db_name = "router_%s" % uuid.uuid4() diff --git a/autopush/web/health.py b/autopush/web/health.py index 973574a9..95e5e8b2 100644 --- a/autopush/web/health.py +++ b/autopush/web/health.py @@ -6,6 +6,7 @@ from twisted.internet.threads import deferToThread from autopush import __version__ +from autopush.db import table_exists from autopush.exceptions import MissingTableException from autopush.web.base import BaseWebHandler @@ -39,14 +40,15 @@ def get(self): def _check_table(self, table): """Checks the tables known about in DynamoDB""" - d = deferToThread(table.connection.list_tables) - d.addCallback(self._check_success, table.table_name) - d.addErrback(self._check_error, table.table_name) + name = table.table_name + d = deferToThread(table_exists, table.connection, name) + d.addCallback(self._check_success, name) + d.addErrback(self._check_error, name) return d - def _check_success(self, result, name): - """Verifies a name is in the list of tables""" - if name not in result.get("TableNames", {}): + def _check_success(self, exists, name): + """Verifies a Table exists""" + if not exists: raise MissingTableException("Nonexistent table") self._health_checks[name] = {"status": "OK"}