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..e610c634 100644 --- a/autopush/db.py +++ b/autopush/db.py @@ -145,9 +145,28 @@ def create_rotating_message_table( boto_resource=None # type: DynamoDBResource ): # type: (...) -> Any # noqa - """Create a new message table for webpush style message storage""" + """Create a new message table for webpush style message storage with a + rotating name. + + """ + tablename = make_rotating_tablename(prefix, delta, date) + return create_message_table( + tablename=tablename, + read_throughput=read_throughput, + write_throughput=write_throughput, + boto_resource=boto_resource + ) + +def create_message_table( + tablename, # type: str + read_throughput=5, # type: int + write_throughput=5, # type: int + boto_resource=None, # type: DynamoDBResource + ): + # type: (...) -> Any # noqa + """Create a new message table for webpush style message storage""" try: table = boto_resource.Table(tablename) if table.table_status == 'ACTIVE': @@ -458,6 +477,7 @@ def __init__(self, **kwargs): if "region_name" in conf: del(conf["region_name"]) self.conf = conf + self._latest_message_table = None self._resource = boto3.resource( "dynamodb", config=botocore.config.Config(region_name=region), @@ -466,6 +486,31 @@ def __init__(self, **kwargs): def __getattr__(self, name): return getattr(self._resource, name) + def get_latest_message_tablenames(self, prefix="message", previous=1): + # type: (Optional[str], int) -> [str] # noqa + """Fetches the name of the last message table""" + client = self._resource.meta.client + paginator = client.get_paginator("list_tables") + tables = [] + for table in paginator.paginate().search( + "TableNames[?contains(@,'{}')==`true`]|sort(@)[-1]".format( + prefix)): + tables.append(table) + if not len(tables) or tables[0] is None: + return [prefix] + tables.sort() + return tables[0-previous:] + + def get_latest_message_tablename(self, prefix="message"): + # type: (Optional[str]) -> str # noqa + """Fetches the name of the last message table""" + if not self._latest_message_table: + self._latest_message_table = self.get_latest_message_tablenames( + prefix=prefix, + previous=1 + )[0] + return self._latest_message_table + class DynamoDBTable(threading.local): def __init__(self, ddb_resource, *args, **kwargs): @@ -479,7 +524,8 @@ 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): + max_ttl=MAX_EXPIRY, + table_base_string="message_"): # type: (str, IMetrics, DynamoDBResource, int) -> None """Create a new Message object @@ -488,10 +534,10 @@ 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 self.table = DynamoDBTable(self.resource, tablename) + self.tablename = tablename def table_status(self): return self.table.table_status @@ -998,16 +1044,27 @@ class DatabaseManager(object): current_msg_month = attrib(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 + ) + else: + # fetch out the last message table as the "current_msg_month" + # Message may still init to this table if it recv's None, but + # this makes the value explicit. + resource = self.resource + self.current_msg_month = resource.get_latest_message_tablename( + prefix=self._message_conf.tablename + ) + if not self.resource: self.resource = DynamoDBResource() @@ -1027,6 +1084,7 @@ def from_config(cls, message_conf=conf.message_table, metrics=metrics, resource=resource, + allow_table_rotation=conf.allow_table_rotation, **kwargs ) @@ -1054,7 +1112,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 +1137,24 @@ def create_initial_message_tables(self): an entry for tomorrow, if tomorrow is a new month. """ + if not self.allow_table_rotation: + tablenames = self.resource.get_latest_message_tablenames( + prefix=self._message_conf.tablename, + previous=3 + ) + # Create the most recent table if it's not there. + tablename = tablenames[-1:][0] + if not table_exists(tablename, + boto_resource=self.resource): + create_message_table( + tablename=tablename, + read_throughput=self._message_conf.read_throughput, + write_throughput=self._message_conf.write_throughput, + boto_resource=self.resource + ) + self.message_tables.extend(tablenames) + return + mconf = self._message_conf today = datetime.date.today() last_month = get_rotating_message_tablename( @@ -1116,6 +1193,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/main_argparse.py b/autopush/main_argparse.py index 7572963c..5688c6ca 100644 --- a/autopush/main_argparse.py +++ b/autopush/main_argparse.py @@ -105,6 +105,10 @@ def add_shared_args(parser): help="AWS DynamoDB endpoint override", type=str, default=None, env_var="AWS_LOCAL_DYNAMODB") + parser.add_argument('--allow_table_rotation', + help="Allow monthly message table rotation", + action="store_true", default=False, + env_var="ALLOW_TABLE_ROTATION") # No ENV because this is for humans _add_external_router_args(parser) _obsolete_args(parser) diff --git a/autopush/tests/test_db.py b/autopush/tests/test_db.py index 4dbace48..f7b5d662 100644 --- a/autopush/tests/test_db.py +++ b/autopush/tests/test_db.py @@ -84,6 +84,47 @@ 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 + fake_conf.message_table = Mock() + fake_conf.message_table.tablename = "message_int_test" + fake_conf.message_table.read_throughput = 5 + fake_conf.message_table.write_throughput = 5 + dm = DatabaseManager.from_config( + fake_conf, + resource=autopush.tests.boto_resource) + dm.create_initial_message_tables() + assert dm.current_msg_month == \ + autopush.tests.boto_resource.get_latest_message_tablename( + prefix=fake_conf.message_table.tablename + ) + + def test_init_with_no_rotate_create_table(self): + fake_conf = Mock() + fake_conf.allow_table_rotation = False + fake_conf.message_table = Mock() + fake_conf.message_table.tablename = "message_bogus" + fake_conf.message_table.read_throughput = 5 + fake_conf.message_table.write_throughput = 5 + safe = autopush.tests.boto_resource._latest_message_table + try: + autopush.tests.boto_resource._latest_message_table = None + dm = DatabaseManager.from_config( + fake_conf, + resource=autopush.tests.boto_resource) + dm.create_initial_message_tables() + assert dm.current_msg_month == \ + autopush.tests.boto_resource.get_latest_message_tablename( + prefix=fake_conf.message_table.tablename + ) + assert dm.message_tables == [fake_conf.message_table.tablename] + finally: + # clean up the bogus table. + autopush.tests.boto_resource._latest_message_table = safe + dm.resource._resource.meta.client.delete_table( + TableName=fake_conf.message_table.tablename) + class DdbResourceTest(unittest.TestCase): @patch("boto3.resource") @@ -205,10 +246,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" + table_name = self.resource.get_latest_message_tablename( + prefix=message_tablename) + message = Message(table_name, SinkMetrics(), + boto_resource=self.resource, + table_base_string=message_tablename) + assert message.tablename == table_name + def test_register(self): chid = str(uuid.uuid4()) diff --git a/autopush/tests/test_rs_integration.py b/autopush/tests/test_rs_integration.py index 68965ad8..a2fb557c 100644 --- a/autopush/tests/test_rs_integration.py +++ b/autopush/tests/test_rs_integration.py @@ -13,11 +13,12 @@ import re import socket import time +import datetime import uuid from contextlib import contextmanager from http.server import BaseHTTPRequestHandler, HTTPServer from httplib import HTTPResponse # noqa -from mock import Mock, call +from mock import Mock, call, patch from threading import Thread, Event from unittest.case import SkipTest @@ -114,6 +115,18 @@ class TestRustWebPush(unittest.TestCase): use_cryptography=True, ) + def start_ep(self, ep_conf): + self._ep_conf = ep_conf + + # Endpoint HTTP router + self.ep = ep = EndpointApplication( + ep_conf, + resource=autopush.tests.boto_resource + ) + ep.setup(rotate_tables=False) + ep.startService() + self.addCleanup(ep.stopService) + def setUp(self): self.logs = TestingLogObserver() begin_or_register(self.logs) @@ -133,16 +146,7 @@ def setUp(self): human_logs=False, **self.conn_kwargs() ) - - # Endpoint HTTP router - self.ep = ep = EndpointApplication( - ep_conf, - resource=autopush.tests.boto_resource - ) - ep.setup(rotate_tables=False) - ep.startService() - self.addCleanup(ep.stopService) - + self.start_ep(ep_conf) # Websocket server self.conn = conn = RustConnectionApplication( conn_conf, @@ -182,6 +186,36 @@ def legacy_endpoint(self): def _ws_url(self): return "ws://localhost:{}/".format(self.connection_port) + @inlineCallbacks + def test_no_rotation(self): + # override autopush settings + safe = self._ep_conf.allow_table_rotation + self._ep_conf.allow_table_rotation = False + yield self.ep.stopService() + try: + self.start_ep(self._ep_conf) + data = str(uuid.uuid4()) + client = yield self.quick_register() + result = yield client.send_notification(data=data) + assert result["headers"]["encryption"] == client._crypto_key + assert result["data"] == base64url_encode(data) + assert result["messageType"] == "notification" + + assert len(self.ep.db.message_tables) == 1 + table_name = self.ep.db.message_tables[0] + target_day = datetime.date(2016, 2, 29) + with patch.object(datetime, 'date', + Mock(wraps=datetime.date)) as patched: + patched.today.return_value = target_day + yield self.ep.db.update_rotating_tables() + assert len(self.ep.db.message_tables) == 1 + assert table_name == self.ep.db.message_tables[0] + finally: + yield self.ep.stopService() + self._ep_conf.allow_table_rotation = safe + self.start_ep(self._ep_conf) + yield self.shut_down(client) + @inlineCallbacks def test_hello_only_has_three_calls(self): db.TRACK_DB_CALLS = True diff --git a/autopush/tests/test_webpush_server.py b/autopush/tests/test_webpush_server.py index 063cc9fc..cc2afdf3 100644 --- a/autopush/tests/test_webpush_server.py +++ b/autopush/tests/test_webpush_server.py @@ -431,6 +431,20 @@ 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..fcfbbdb8 100644 --- a/autopush/tests/test_z_main.py +++ b/autopush/tests/test_z_main.py @@ -178,6 +178,42 @@ def check_tables(result): d.addBoth(lambda x: e.callback(True)) return e + def test_no_rotation(self): + today = datetime.date.today() + next_month = today.month + 1 + next_year = today.year + if next_month > 12: # pragma: nocover + next_month = 1 + next_year += 1 + tomorrow = datetime.datetime(year=next_year, + month=next_month, + day=1) + conf = AutopushConfig( + hostname="example.com", + resolve_hostname=True, + allow_table_rotation=False + ) + resource = autopush.tests.boto_resource + db = DatabaseManager.from_config( + conf, + resource=resource) + db._tomorrow = Mock(return_value=tomorrow) + db.create_initial_message_tables() + assert len(db.message_tables) == 1 + assert db.message_tables[0] == resource.get_latest_message_tablename( + prefix=conf.message_table.tablename + ) + + def check_tables(result): + assert len(db.message_tables) == 1 + assert db.message_tables[0] == \ + resource.get_latest_message_tablename( + prefix=conf.message_table.tablename + ) + dd = db.update_rotating_tables() + dd.addCallback(check_tables) + return dd + class ConnectionMainTestCase(unittest.TestCase): def setUp(self): 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