wip: feat: add flag to stop table rotation

Closes #1172

autopush/config.py

@@ -183,6 +183,8 @@ class AutopushConfig(object):
     # DynamoDB endpoint override
     aws_ddb_endpoint = attrib(default=None)  # type: str
 
+    allow_table_rotation = attrib(default=True)  # type: bool
+
     def __attrs_post_init__(self):
         """Initialize the Settings object"""
         # Setup hosts/ports/urls

autopush/db.py

@@ -466,6 +466,15 @@ def __init__(self, **kwargs):
     def __getattr__(self, name):
         return getattr(self._resource, name)
 
+    def get_latest_message_tablename(self, prefix="message"):
+        # type: (Optional[str]) -> str  # noqa
+        """Fetches the name of the last message table"""
+        response = filter(
+            lambda name: name.lower().startswith(prefix),
+            self._resource.meta.client.list_tables().get("TableNames"))
+        response.sort()
+        return response[-1]
+
 
 class DynamoDBTable(threading.local):
     def __init__(self, ddb_resource, *args, **kwargs):
@@ -478,8 +487,9 @@ def __getattr__(self, name):
 
 class Message(object):
     """Create a Message table abstraction on top of a DynamoDB Table object"""
-    def __init__(self, tablename, metrics=None, boto_resource=None,
-                 max_ttl=MAX_EXPIRY):
+    def __init__(self, tablename=None, metrics=None, boto_resource=None,
+                 max_ttl=MAX_EXPIRY,
+                 table_base_string="message_"):
         # type: (str, IMetrics, DynamoDBResource, int) -> None
         """Create a new Message object
 
@@ -488,10 +498,13 @@ def __init__(self, tablename, metrics=None, boto_resource=None,
         :param boto_resource: DynamoDBResource for thread
 
         """
-        self.tablename = tablename
         self._max_ttl = max_ttl
         self.resource = boto_resource
+        if tablename is None:
+            tablename = boto_resource.get_latest_message_tablename(
+                table_base_string)
         self.table = DynamoDBTable(self.resource, tablename)
+        self.tablename = tablename
 
     def table_status(self):
         return self.table.table_status
@@ -995,19 +1008,22 @@ class DatabaseManager(object):
 
     router = attrib(default=None)                   # type: Optional[Router]
     message_tables = attrib(default=Factory(list))  # type: List[str]
-    current_msg_month = attrib(init=False)          # type: Optional[str]
+    current_msg_month = attrib(default=None, init=False)  # type: Optional[str]
     current_month = attrib(init=False)              # type: Optional[int]
     _message = attrib(default=None)                 # type: Optional[Message]
+    allow_table_rotation = attrib(default=True)     # type: Optional[bool]
     # for testing:
 
     def __attrs_post_init__(self):
         """Initialize sane defaults"""
-        today = datetime.date.today()
-        self.current_month = today.month
-        self.current_msg_month = make_rotating_tablename(
-            self._message_conf.tablename,
-            date=today
-        )
+        if self.allow_table_rotation:
+            today = datetime.date.today()
+            self.current_month = today.month
+            self.current_msg_month = make_rotating_tablename(
+                self._message_conf.tablename,
+                date=today
+            )
+
         if not self.resource:
             self.resource = DynamoDBResource()
 
@@ -1027,6 +1043,7 @@ def from_config(cls,
             message_conf=conf.message_table,
             metrics=metrics,
             resource=resource,
+            allow_table_rotation=conf.allow_table_rotation,
             **kwargs
         )
 
@@ -1054,7 +1071,8 @@ def setup_tables(self):
         self.create_initial_message_tables()
         self._message = Message(self.current_msg_month,
                                 self.metrics,
-                                boto_resource=self.resource)
+                                boto_resource=self.resource,
+                                table_base_string=self._message_conf.tablename)
 
     @property
     def message(self):
@@ -1078,6 +1096,9 @@ def create_initial_message_tables(self):
         an entry for tomorrow, if tomorrow is a new month.
 
         """
+        if not self.allow_table_rotation:
+            return
+
         mconf = self._message_conf
         today = datetime.date.today()
         last_month = get_rotating_message_tablename(
@@ -1106,6 +1127,7 @@ def create_initial_message_tables(self):
             )
             self.message_tables.append(next_month)
 
+
     @inlineCallbacks
     def update_rotating_tables(self):
         # type: () -> Generator
@@ -1116,6 +1138,8 @@ def update_rotating_tables(self):
         table objects on the settings object.
 
         """
+        if not self.allow_table_rotation:
+            returnValue(False)
         mconf = self._message_conf
         today = datetime.date.today()
         tomorrow = self._tomorrow()

autopush/diagnostic_cli.py

@@ -69,7 +69,10 @@ def run(self):
         print("\n")
 
         if "current_month" in rec:
-            chans = Message(rec["current_month"],
+            month = None
+            if self._conf.allow_table_rotation:
+                month = rec["current_month"]
+            chans = Message(month,
                             boto_resource=self.db.resource).all_channels(uaid)
             print("Channels in message table:")
             self._pp.pprint(chans)

autopush/tests/test_db.py

@@ -84,6 +84,13 @@ def test_init_with_resources(self):
         assert dm.resource is not None
         assert isinstance(dm.resource, DynamoDBResource)
 
+    def test_init_with_no_rotate(self):
+        fake_conf = Mock()
+        fake_conf.allow_table_rotation = False
+        dm = DatabaseManager.from_config(fake_conf)
+        dm.create_initial_message_tables()
+        assert dm.current_msg_month == None
+
 
 class DdbResourceTest(unittest.TestCase):
     @patch("boto3.resource")
@@ -205,10 +212,21 @@ def test_normalize_id(self):
 class MessageTestCase(unittest.TestCase):
     def setUp(self):
         self.resource = autopush.tests.boto_resource
-        table = get_rotating_message_tablename(boto_resource=self.resource)
+        table = get_rotating_message_tablename(
+            prefix="message_int_test",
+            boto_resource=self.resource)
         self.real_table = table
         self.uaid = str(uuid.uuid4())
 
+    def test_non_rotating_tables(self):
+        message_tablename="message_int_test"
+        message = Message(None, SinkMetrics(),
+                          boto_resource=self.resource,
+                          table_base_string=message_tablename)
+        table_name = self.resource.get_latest_message_tablename(
+            prefix=message_tablename)
+        assert message.tablename == table_name
+
     def test_register(self):
         chid = str(uuid.uuid4())
 

autopush/tests/test_webpush_server.py

@@ -17,7 +17,7 @@
     Message,
 )
 from autopush.metrics import SinkMetrics
-from autopush.config import AutopushConfig
+from autopush.config import AutopushConfig, DDBTableConfig
 from autopush.exceptions import ItemNotFound
 from autopush.logging import begin_or_register
 from autopush.tests.support import TestingLogObserver
@@ -466,6 +466,19 @@ def test_migrate_user(self):
         assert item is not None
         assert len(channels) == 3
 
+    def test_no_migrate(self):
+        self.conf.allow_table_rotation = False
+        self.conf.message_table.tablename="message_int_test"
+        self.db = db = DatabaseManager.from_config(
+            self.conf,
+            resource=autopush.tests.boto_resource
+        )
+        assert self.db.allow_table_rotation is False
+        db.setup_tables()
+        tablename = autopush.tests.boto_resource.get_latest_message_tablename(
+            prefix="message_int_test"
+        )
+        assert db.message.tablename == tablename
 
 class TestRegisterProcessor(BaseSetup):
 

autopush/tests/test_z_main.py

@@ -178,6 +178,23 @@ def check_tables(result):
         d.addBoth(lambda x: e.callback(True))
         return e
 
+    def test_no_rotation(self):
+        conf = AutopushConfig(
+            allow_table_rotation=False
+        )
+        db = DatabaseManager.from_config(
+            conf,
+            resource=autopush.tests.boto_resource)
+        db.create_initial_message_tables()
+        # Erase the tables it has on init, and move current month back one
+        # get the table name.
+        db.message_tables = []
+
+        # Get the deferred back
+        e = Deferred()
+        d = db.update_rotating_tables()
+        #  TODO: Finish this test (take from
+        # test_update_rotating_tables_month_end)
 
 class ConnectionMainTestCase(unittest.TestCase):
     def setUp(self):

autopush/webpush_server.py

@@ -363,6 +363,11 @@ def metrics(self):
     def process(self, command):
         raise NotImplementedError()
 
+    def get_month(self, command):
+        if not self.conf.allow_table_rotation:
+            return None
+        return command.message_month
+
 
 class HelloCommand(ProcessorCommand):
     def process(self, hello):
@@ -481,7 +486,7 @@ def process(self, command):
     def _check_storage(self, command):
         timestamp = None
         messages = []
-        message = Message(command.message_month,
+        message = Message(self.get_month(command),
                           boto_resource=self.db.resource)
         if command.include_topic:
             timestamp, messages = message.fetch_messages(
@@ -516,7 +521,7 @@ def _check_storage(self, command):
 class IncrementStorageCommand(ProcessorCommand):
     def process(self, command):
         # type: (IncStoragePosition) -> IncStoragePositionResponse
-        message = Message(command.message_month,
+        message = Message(self.get_month(command),
                           boto_resource=self.db.resource)
         message.update_last_message_read(command.uaid, command.timestamp)
         return IncStoragePositionResponse()
@@ -526,7 +531,7 @@ class DeleteMessageCommand(ProcessorCommand):
     def process(self, command):
         # type: (DeleteMessage) -> DeleteMessageResponse
         notif = command.message.to_WebPushNotification()
-        message = Message(command.message_month,
+        message = Message(self.get_month(command),
                           boto_resource=self.db.resource)
         message.delete_message(notif)
         return DeleteMessageResponse()
@@ -543,7 +548,9 @@ class MigrateUserCommand(ProcessorCommand):
     def process(self, command):
         # type: (MigrateUser) -> MigrateUserResponse
         # Get the current channels for this month
-        message = Message(command.message_month,
+        if not self.conf.allow_table_rotation:
+            return MigrateUserResponse(message_month=None)
+        message = Message(self.get_month(command),
                           boto_resource=self.db.resource)
         _, channels = message.all_channels(command.uaid.hex)
 
@@ -565,7 +572,7 @@ def process(self, command):
 class StoreMessagesUserCommand(ProcessorCommand):
     def process(self, command):
         # type: (StoreMessages) -> StoreMessagesResponse
-        message = Message(command.message_month,
+        message = Message(self.get_month(command),
                           boto_resource=self.db.resource)
         for m in command.messages:
             if "topic" not in m:
@@ -620,7 +627,7 @@ def process(self, command):
             command.channel_id,
             command.key
         )
-        message = self.db.message_table(command.message_month)
+        message = self.db.message_table(self.get_month(command))
         try:
             message.register_channel(command.uaid.hex,
                                      command.channel_id)
@@ -669,7 +676,7 @@ def process(self,
         if not valid:
             return UnregisterErrorResponse(error_msg=msg)
 
-        message = Message(command.message_month,
+        message = Message(self.get_month(command),
                           boto_resource=self.db.resource)
         try:
             message.unregister_channel(command.uaid.hex, command.channel_id)

configs/autopush_shared.ini.sample

@@ -132,3 +132,6 @@ endpoint_port = 8082
 ; e.g {"firefox":{"cert":"certs/main.cert","key":"certs/main.key","topic":"com.mozilla.org.Firefox","max_retry":2},"beta":{"cert":"certs/beta.cert","key":"certs/beta.key","topic":"com.mozilla.org.FirefoxBeta"}}
 #apns_creds =
 
+; With TTL implemented, message table rotation is no longer required.
+; This flag determines if table rotation should be allowed to continue:
+#allow_table_rotation = true