make it work and fix pep8

1 files changed, 43 insertions, 26 deletions

diff --git a/synapse/handlers/e2e_room_keys.py b/synapse/handlers/e2e_room_keys.py
index 15e3beb5ed..93f4ad5194 100644
--- a/synapse/handlers/e2e_room_keys.py
+++ b/synapse/handlers/e2e_room_keys.py
@@ -13,15 +13,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import ujson as json
 import logging
 
-from canonicaljson import encode_canonical_json
 from twisted.internet import defer
 
-from synapse.api.errors import SynapseError, CodeMessageException
+from synapse.api.errors import StoreError
 from synapse.util.async import Linearizer
-from synapse.util.retryutils import NotRetryingDestination
 
 logger = logging.getLogger(__name__)
 
@@ -29,11 +26,13 @@ logger = logging.getLogger(__name__)
 class E2eRoomKeysHandler(object):
     def __init__(self, hs):
         self.store = hs.get_datastore()
-        self._upload_linearizer = async.Linearizer("upload_room_keys_lock")
+        self._upload_linearizer = Linearizer("upload_room_keys_lock")
 
     @defer.inlineCallbacks
     def get_room_keys(self, user_id, version, room_id, session_id):
-        results = yield self.store.get_e2e_room_keys(user_id, version, room_id, session_id)
+        results = yield self.store.get_e2e_room_keys(
+            user_id, version, room_id, session_id
+        )
         defer.returnValue(results)
 
     @defer.inlineCallbacks
@@ -46,31 +45,49 @@ class E2eRoomKeysHandler(object):
         # TODO: Validate the JSON to make sure it has the right keys.
 
         # XXX: perhaps we should use a finer grained lock here?
-        with (yield self._upload_linearizer.queue(user_id):
+        with (yield self._upload_linearizer.queue(user_id)):
 
             # go through the room_keys
             for room_id in room_keys['rooms']:
                 for session_id in room_keys['rooms'][room_id]['sessions']:
                     room_key = room_keys['rooms'][room_id]['sessions'][session_id]
 
-                    # get the room_key for this particular row
-                    current_room_key = yield self.store.get_e2e_room_key(
-                        user_id, version, room_id, session_id
+                    yield self._upload_room_key(
+                        user_id, version, room_id, session_id, room_key
                     )
 
-                    # check whether we merge or not. spelling it out with if/elifs rather than
-                    # lots of booleans for legibility.
-                    replace = False
-                    if current_room_key:
-                        if room_key['is_verified'] and not current_room_key['is_verified']:
-                            replace = True
-                        elif room_key['first_message_index'] < current_room_key['first_message_index']:
-                            replace = True
-                        elif room_key['forwarded_count'] < room_key['forwarded_count']:
-                            replace = True
-
-                    # if so, we set the new room_key
-                    if replace:
-                        yield self.store.set_e2e_room_key(
-                            user_id, version, room_id, session_id, room_key
-                        )
+    @defer.inlineCallbacks
+    def _upload_room_key(self, user_id, version, room_id, session_id, room_key):
+        # get the room_key for this particular row
+        current_room_key = None
+        try:
+            current_room_key = yield self.store.get_e2e_room_key(
+                user_id, version, room_id, session_id
+            )
+        except StoreError as e:
+            if e.code == 404:
+                pass
+            else:
+                raise
+
+        # check whether we merge or not. spelling it out with if/elifs rather
+        # than lots of booleans for legibility.
+        upsert = True
+        if current_room_key:
+            if room_key['is_verified'] and not current_room_key['is_verified']:
+                pass
+            elif (
+                room_key['first_message_index'] <
+                current_room_key['first_message_index']
+            ):
+                pass
+            elif room_key['forwarded_count'] < room_key['forwarded_count']:
+                pass
+            else:
+                upsert = False
+
+        # if so, we set the new room_key
+        if upsert:
+            yield self.store.set_e2e_room_key(
+                user_id, version, room_id, session_id, room_key
+            )