feat: retain date when checking message tables

autopush/db.py

@@ -128,11 +128,12 @@ def make_rotating_tablename(prefix, delta=0, date=None):
     return "{}_{:04d}_{:02d}".format(prefix, date.year, date.month)
 
 
-def create_rotating_message_table(prefix="message", read_throughput=5,
-                                  write_throughput=5, delta=0):
-    # type: (str, int, int, int) -> Table
+def create_rotating_message_table(prefix="message", delta=0, date=None,
+                                  read_throughput=5,
+                                  write_throughput=5):
+    # type: (str, int, Optional[datetime.date], int, int) -> Table
     """Create a new message table for webpush style message storage"""
-    tablename = make_rotating_tablename(prefix, delta)
+    tablename = make_rotating_tablename(prefix, delta, date)
     return Table.create(tablename,
                         schema=[HashKey("uaid"),
                                 RangeKey("chidmessageid")],
@@ -151,9 +152,10 @@ def get_rotating_message_table(prefix="message", delta=0, date=None,
     tablename = make_rotating_tablename(prefix, delta, date)
     if tablename not in dblist:
         return create_rotating_message_table(
-            prefix=prefix, delta=delta,
+            prefix=prefix, delta=delta, date=date,
             read_throughput=message_read_throughput,
-            write_throughput=message_write_throughput)
+            write_throughput=message_write_throughput,
+        )
     else:
         return Table(tablename)
 

autopush/settings.py

@@ -264,8 +264,10 @@ def update_rotating_tables(self):
         if ((tomorrow.month != today.month) and
                 sorted(self.message_tables.keys())[-1] !=
                 tomorrow.month):
-            next_month = get_rotating_message_table(
-                self._message_prefix, 0, tomorrow)
+            next_month = yield deferToThread(
+                get_rotating_message_table,
+                self._message_prefix, 0, tomorrow
+            )
             self.message_tables[next_month.table_name] = Message(
                 next_month, self.metrics)
 

autopush/tests/test_db.py

@@ -1,6 +1,6 @@
 import unittest
 import uuid
-from datetime import datetime
+from datetime import datetime, timedelta
 
 from autopush.websocket import ms_time
 from boto.dynamodb2.exceptions import (
@@ -25,7 +25,7 @@
     Message,
     Router,
     generate_last_connect,
-)
+    make_rotating_tablename)
 from autopush.exceptions import AutopushException
 from autopush.metrics import SinkMetrics
 from autopush.utils import WebPushNotification
@@ -360,6 +360,14 @@ def raise_condition(*args, **kwargs):
         result = message.delete_message(notif)
         eq_(result, False)
 
+    def test_message_rotate_table_with_date(self):
+        prefix = "message" + uuid.uuid4().hex
+        future = datetime.today() + timedelta(days=32)
+        tbl_name = make_rotating_tablename(prefix, date=future)
+
+        m = get_rotating_message_table(prefix=prefix, date=future)
+        eq_(m.table_name, tbl_name)
+
 
 class RouterTestCase(unittest.TestCase):
     def setUp(self):