diff --git a/autopush/db.py b/autopush/db.py index ba8c2fc4..7b427063 100644 --- a/autopush/db.py +++ b/autopush/db.py @@ -58,10 +58,11 @@ def normalize_id(id): return '-'.join((raw[:8], raw[8:12], raw[12:16], raw[16:20], raw[20:])) -def make_rotating_tablename(prefix, delta=0): +def make_rotating_tablename(prefix, delta=0, date=None): """Creates a tablename for table rotation based on a prefix with a given month delta.""" - date = get_month(delta=delta) + if not date: + date = get_month(delta=delta) return "{}_{}_{}".format(prefix, date.year, date.month) @@ -77,11 +78,11 @@ def create_rotating_message_table(prefix="message", read_throughput=5, ) -def get_rotating_message_table(prefix="message", delta=0): +def get_rotating_message_table(prefix="message", delta=0, date=None): """Gets the message table for the current month.""" db = DynamoDBConnection() dblist = db.list_tables()["TableNames"] - tablename = make_rotating_tablename(prefix, delta) + tablename = make_rotating_tablename(prefix, delta, date) if tablename not in dblist: return create_rotating_message_table(prefix=prefix, delta=delta) else: diff --git a/autopush/settings.py b/autopush/settings.py index 669cfe7e..e717b227 100644 --- a/autopush/settings.py +++ b/autopush/settings.py @@ -19,11 +19,10 @@ get_router_table, get_storage_table, get_rotating_message_table, - make_rotating_tablename, preflight_check, Storage, Router, - Message + Message, ) from autopush.exceptions import InvalidTokenException from autopush.metrics import ( @@ -163,9 +162,12 @@ def __init__(self, self.router = Router(self.router_table, self.metrics) # Used to determine whether a connection is out of date with current - # db objects - self.current_msg_month = make_rotating_tablename(self._message_prefix) - self.current_month = datetime.date.today().month + # db objects. There are three noteworty cases: + # 1 "Last Month" the table requires a rollover. + # 2 "This Month" the most common case. + # 3 "Tomorrow" in the case where the system is actively rolling over + # and one node is not yet using the latest clock, no rollover + # required. self.create_initial_message_tables() # Run preflight check @@ -204,18 +206,35 @@ def message(self, value): """Setter to set the current message table""" self.message_tables[self.current_msg_month] = value + def add_tomorrow(self, today, this_month_table): + """check to see if tomorrow is a new month, in case we're a node that + has a slow clock.""" + tomorrow = get_rotating_message_table( + self._message_prefix, 0, today + datetime.timedelta(days=1)) + if tomorrow.table_name != this_month_table: + self.message_tables[tomorrow.table_name] = Message( + tomorrow, self.metrics) + def create_initial_message_tables(self): """Initializes a dict of the initial rotating messages tables. - An entry for last months table, and an entry for this months table. + An entry for last months table, an entry for this months table, + and an entry for tomorrow, if it starts a new month. """ + today = datetime.date.today() last_month = get_rotating_message_table(self._message_prefix, -1) this_month = get_rotating_message_table(self._message_prefix) + self.current_month = today.month + self.current_msg_month = this_month.table_name self.message_tables = { last_month.table_name: Message(last_month, self.metrics), - this_month.table_name: Message(this_month, self.metrics), + this_month.table_name: Message(this_month, self.metrics) } + # Due to issues around mocking "datetime" and desire to + # reduce excessive hotpath function calls, we have to skip this + # line + self.add_tomorrow(today, this_month.table_name) @inlineCallbacks def update_rotating_tables(self): @@ -227,6 +246,7 @@ def update_rotating_tables(self): """ today = datetime.date.today() + self.add_tomorrow(today, self.current_msg_month) if today.month == self.current_month: # No change in month, we're fine. returnValue(False) @@ -241,7 +261,6 @@ def update_rotating_tables(self): self.current_msg_month = message_table.table_name self.message_tables[self.current_msg_month] = \ Message(message_table, self.metrics) - returnValue(True) def update(self, **kwargs): diff --git a/autopush/tests/test_db.py b/autopush/tests/test_db.py index f6e36c4b..45228d67 100644 --- a/autopush/tests/test_db.py +++ b/autopush/tests/test_db.py @@ -70,6 +70,14 @@ def test_hasher(self): 'd8f614d06cdd592cb8470f31177c8331a') db.key_hash = "" + def test_normalize_id(self): + import autopush.db as db + abnormal = "deadbeef00000000decafbad00000000" + normal = "deadbeef-0000-0000-deca-fbad00000000" + eq_(db.normalize_id(abnormal), normal) + self.assertRaises(ValueError, db.normalize_id, "invalid") + eq_(db.normalize_id(abnormal.upper()), normal) + class StorageTestCase(unittest.TestCase): def setUp(self): diff --git a/autopush/tests/test_websocket.py b/autopush/tests/test_websocket.py index 248e4839..dec0cc49 100644 --- a/autopush/tests/test_websocket.py +++ b/autopush/tests/test_websocket.py @@ -1,4 +1,5 @@ import json +import datetime import time import uuid from hashlib import sha256 @@ -18,7 +19,9 @@ from twisted.trial import unittest import autopush.db as db -from autopush.db import create_rotating_message_table +from autopush.db import ( + create_rotating_message_table, +) from autopush.settings import AutopushSettings from autopush.websocket import ( PushState, @@ -406,6 +409,131 @@ def wait_for_agent_call(): # pragma: nocover reactor.callLater(0.1, wait_for_agent_call) return d + def test_hello_old(self): + orig_uaid = "deadbeef12345678decafbad12345678" + # router.register_user returns (registered, previous + target_day = datetime.date(2016, 2, 29) + msg_day = datetime.date(2015, 12, 15) + msg_date = "{}_{}_{}".format( + self.proto.ap_settings._message_prefix, + msg_day.year, + msg_day.month) + msg_data = { + "router_type": "webpush", + "node_id": "http://localhost", + "last_connect": int(msg_day.strftime("%s")), + "current_month": msg_date, + } + + def fake_msg(data): + return (True, msg_data, data) + + mock_msg = Mock(wraps=db.Message) + mock_msg.fetch_messages.return_value = [] + self.proto.ap_settings.router.register_user = fake_msg + # massage message_tables to include our fake range + mt = self.proto.ps.settings.message_tables + for k in mt.keys(): + del(mt[k]) + mt['message_2016_1'] = mock_msg + mt['message_2016_2'] = mock_msg + mt['message_2016_3'] = mock_msg + with patch.object(datetime, 'date', + Mock(wraps=datetime.date)) as patched: + patched.today.return_value = target_day + self._connect() + self._send_message(dict(messageType="hello", + uaid=orig_uaid, + channelIDs=[], + use_webpush=True)) + + def check_result(msg): + eq_(self.proto.ps.rotate_message_table, False) + # it's fine you've not connected in a while, but + # you should recycle your endpoints since they're probably + # invalid by now anyway. + eq_(msg["status"], 200) + ok_(msg["uaid"] != orig_uaid) + + return self._check_response(check_result) + + def test_hello_tomorrow(self): + orig_uaid = "deadbeef12345678decafbad12345678" + # router.register_user returns (registered, previous + target_day = datetime.date(2016, 2, 29) + msg_day = datetime.date(2016, 3, 1) + msg_date = "{}_{}_{}".format( + self.proto.ap_settings._message_prefix, + msg_day.year, + msg_day.month) + msg_data = { + "router_type": "webpush", + "node_id": "http://localhost", + "last_connect": int(msg_day.strftime("%s")), + "current_month": msg_date, + } + + def fake_msg(data): + return (True, msg_data, data) + + mock_msg = Mock(wraps=db.Message) + mock_msg.fetch_messages.return_value = [] + self.proto.ap_settings.router.register_user = fake_msg + # massage message_tables to include our fake range + mt = self.proto.ps.settings.message_tables + for k in mt.keys(): + del(mt[k]) + mt['message_2016_1'] = mock_msg + mt['message_2016_2'] = mock_msg + mt['message_2016_3'] = mock_msg + with patch.object(datetime, 'date', + Mock(wraps=datetime.date)) as patched: + patched.today.return_value = target_day + self._connect() + self._send_message(dict(messageType="hello", + uaid=orig_uaid, + channelIDs=[], + use_webpush=True)) + + def check_result(msg): + eq_(self.proto.ps.rotate_message_table, False) + # it's fine you've not connected in a while, but + # you should recycle your endpoints since they're probably + # invalid by now anyway. + eq_(msg["status"], 200) + eq_(msg["uaid"], orig_uaid) + + return self._check_response(check_result) + + def test_add_tomorrow(self): + today = datetime.date(2016, 2, 29) + yester = datetime.date(2016, 1, 1) + tomorrow = datetime.date(2016, 3, 1) + today_table = "{}_{}_{}".format( + self.proto.ap_settings._message_prefix, + today.year, + today.month) + yester_table = "{}_{}_{}".format( + self.proto.ap_settings._message_prefix, + yester.year, + yester.month) + tomorrow_table = "{}_{}_{}".format( + self.proto.ap_settings._message_prefix, + tomorrow.year, + tomorrow.month) + + mock_msg = Mock(wraps=db.Message) + mock_msg.fetch_messages.return_value = [] + mt = self.proto.ps.settings.message_tables + for k in mt.keys(): + del(mt[k]) + mt[yester_table] = mock_msg + mt[today_table] = mock_msg + + self._connect() + self.proto.ps.settings.add_tomorrow(today, today_table) + ok_(tomorrow_table in self.proto.ps.settings.message_tables) + def test_hello(self): self._connect() self._send_message(dict(messageType="hello", channelIDs=[])) diff --git a/autopush/websocket.py b/autopush/websocket.py index 62016d72..f1db662f 100644 --- a/autopush/websocket.py +++ b/autopush/websocket.py @@ -776,9 +776,14 @@ def _check_message_table_rotation(self, previous): self.transport.pauseProducing() # Check for table rotation cur_month = previous.get("current_month") + # Previous month user or new user, flag for message rotation and + # set the message_month to the router month if cur_month != self.ps.message_month: - # Previous month user or new user, flag for message rotation and - # set the message_month to the router month + if cur_month not in self.ps.settings.message_tables.keys(): + # This UAID has expired. Force client to reregister. + self.ps.uaid = uuid.uuid4().hex + self._finish_webpush_hello() + return self.ps.message_month = cur_month self.ps.rotate_message_table = True