Skip to content
This repository has been archived by the owner on Jul 13, 2023. It is now read-only.

Commit

Permalink
fix: ensure we paginate through all table names
Browse files Browse the repository at this point in the history
Closes #1000
  • Loading branch information
pjenvey committed Sep 8, 2017
1 parent ab67f3a commit f9f0d0b
Show file tree
Hide file tree
Showing 3 changed files with 32 additions and 22 deletions.
25 changes: 19 additions & 6 deletions autopush/db.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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):
Expand Down
15 changes: 5 additions & 10 deletions autopush/tests/test_db.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
create_router_table,
create_storage_table,
preflight_check,
table_exists,
Storage,
Message,
Router,
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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()
Expand Down
14 changes: 8 additions & 6 deletions autopush/web/health.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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"}

Expand Down

0 comments on commit f9f0d0b

Please sign in to comment.