fix: ensure we paginate through all table names

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):

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()

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"}