From 1f9e3df6b6b9e31685c0e7a3ce25258c4b811528 Mon Sep 17 00:00:00 2001 From: jrconlin Date: Fri, 13 Apr 2018 18:04:43 -0700 Subject: [PATCH] wip: feat: add flag to stop table rotation Closes #1172 --- autopush/config.py | 2 ++ autopush/db.py | 46 ++++++++++++++++++++------- autopush/diagnostic_cli.py | 5 ++- autopush/tests/test_db.py | 20 +++++++++++- autopush/tests/test_webpush_server.py | 15 ++++++++- autopush/tests/test_z_main.py | 17 ++++++++++ autopush/webpush_server.py | 21 ++++++++---- configs/autopush_shared.ini.sample | 3 ++ 8 files changed, 108 insertions(+), 21 deletions(-) diff --git a/autopush/config.py b/autopush/config.py index 08e0e80b..67a19435 100644 --- a/autopush/config.py +++ b/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 diff --git a/autopush/db.py b/autopush/db.py index f8df639f..62bdd3bf 100644 --- a/autopush/db.py +++ b/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() diff --git a/autopush/diagnostic_cli.py b/autopush/diagnostic_cli.py index 24debb1f..4d47cfea 100644 --- a/autopush/diagnostic_cli.py +++ b/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) diff --git a/autopush/tests/test_db.py b/autopush/tests/test_db.py index 4dbace48..b65040a5 100644 --- a/autopush/tests/test_db.py +++ b/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()) diff --git a/autopush/tests/test_webpush_server.py b/autopush/tests/test_webpush_server.py index 543022ea..316c7254 100644 --- a/autopush/tests/test_webpush_server.py +++ b/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): diff --git a/autopush/tests/test_z_main.py b/autopush/tests/test_z_main.py index 967c4292..42606f46 100644 --- a/autopush/tests/test_z_main.py +++ b/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): diff --git a/autopush/webpush_server.py b/autopush/webpush_server.py index c13ee02e..3ff41da5 100644 --- a/autopush/webpush_server.py +++ b/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) diff --git a/configs/autopush_shared.ini.sample b/configs/autopush_shared.ini.sample index 6994d1e6..687eff84 100644 --- a/configs/autopush_shared.ini.sample +++ b/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