make it work and fix pep8

3 files changed, 64 insertions, 43 deletions

diff --git a/synapse/storage/__init__.py b/synapse/storage/__init__.py
index 134e4a80f1..69cb28268a 100644
--- a/synapse/storage/__init__.py
+++ b/synapse/storage/__init__.py
@@ -30,6 +30,7 @@ from .appservice import ApplicationServiceStore, ApplicationServiceTransactionSt
 from .client_ips import ClientIpStore
 from .deviceinbox import DeviceInboxStore
 from .directory import DirectoryStore
+from .e2e_room_keys import EndToEndRoomKeyStore
 from .end_to_end_keys import EndToEndKeyStore
 from .engines import PostgresEngine
 from .event_federation import EventFederationStore
@@ -76,6 +77,7 @@ class DataStore(RoomMemberStore, RoomStore,
                 ApplicationServiceTransactionStore,
                 ReceiptsStore,
                 EndToEndKeyStore,
+                EndToEndRoomKeyStore,
                 SearchStore,
                 TagsStore,
                 AccountDataStore,
diff --git a/synapse/storage/e2e_room_keys.py b/synapse/storage/e2e_room_keys.py
index 903dc083f8..5982710bd5 100644
--- a/synapse/storage/e2e_room_keys.py
+++ b/synapse/storage/e2e_room_keys.py
@@ -15,11 +15,6 @@
 
 from twisted.internet import defer
 
-from synapse.util.caches.descriptors import cached
-
-from canonicaljson import encode_canonical_json
-import ujson as json
-
 from ._base import SQLBaseStore
 
 
@@ -45,29 +40,27 @@ class EndToEndRoomKeyStore(SQLBaseStore):
             desc="get_e2e_room_key",
         )
 
-        defer.returnValue(row);
+        defer.returnValue(row)
 
     def set_e2e_room_key(self, user_id, version, room_id, session_id, room_key):
 
         def _set_e2e_room_key_txn(txn):
 
-            self._simple_upsert(
+            self._simple_upsert_txn(
                 txn,
                 table="e2e_room_keys",
                 keyvalues={
                     "user_id": user_id,
                     "room_id": room_id,
-                    "session_id": session_id,   
-                }
-                values=[
-                    {
-                        "version": version,
-                        "first_message_index": room_key['first_message_index'],
-                        "forwarded_count": room_key['forwarded_count'],
-                        "is_verified": room_key['is_verified'],
-                        "session_data": room_key['session_data'],
-                    }
-                ],
+                    "session_id": session_id,
+                },
+                values={
+                    "version": version,
+                    "first_message_index": room_key['first_message_index'],
+                    "forwarded_count": room_key['forwarded_count'],
+                    "is_verified": room_key['is_verified'],
+                    "session_data": room_key['session_data'],
+                },
                 lock=False,
             )
 
@@ -77,7 +70,6 @@ class EndToEndRoomKeyStore(SQLBaseStore):
             "set_e2e_room_key", _set_e2e_room_key_txn
         )
 
-
     # XXX: this isn't currently used and isn't tested anywhere
     # it could be used in future for bulk-uploading new versions of room_keys
     # for a user or something though.
@@ -85,23 +77,27 @@ class EndToEndRoomKeyStore(SQLBaseStore):
 
         def _set_e2e_room_keys_txn(txn):
 
+            values = []
+            for room_id in room_keys['rooms']:
+                for session_id in room_keys['rooms'][room_id]['sessions']:
+                    session = room_keys['rooms'][room_id]['sessions'][session_id]
+                    values.append(
+                        {
+                            "user_id": user_id,
+                            "room_id": room_id,
+                            "session_id": session_id,
+                            "version": version,
+                            "first_message_index": session['first_message_index'],
+                            "forwarded_count": session['forwarded_count'],
+                            "is_verified": session['is_verified'],
+                            "session_data": session['session_data'],
+                        }
+                    )
+
             self._simple_insert_many_txn(
                 txn,
                 table="e2e_room_keys",
-                values=[
-                    {
-                        "user_id": user_id,
-                        "room_id": room_id,
-                        "session_id": session_id,
-                        "version": version,
-                        "first_message_index": room_keys['rooms'][room_id]['sessions'][session_id]['first_message_index'],
-                        "forwarded_count": room_keys['rooms'][room_id]['sessions'][session_id]['forwarded_count'],
-                        "is_verified": room_keys['rooms'][room_id]['sessions'][session_id]['is_verified'],
-                        "session_data": room_keys['rooms'][room_id]['sessions'][session_id]['session_data'],
-                    }
-                    for session_id in room_keys['rooms'][room_id]['sessions']
-                    for room_id in room_keys['rooms']
-                ]
+                values=values
             )
 
             return True
@@ -113,17 +109,22 @@ class EndToEndRoomKeyStore(SQLBaseStore):
     @defer.inlineCallbacks
     def get_e2e_room_keys(self, user_id, version, room_id, session_id):
 
-        keyvalues={
+        keyvalues = {
             "user_id": user_id,
             "version": version,
         }
-        if room_id: keyvalues['room_id'] = room_id
-        if session_id: keyvalues['session_id'] = session_id
+        if room_id:
+            keyvalues['room_id'] = room_id
+        if session_id:
+            keyvalues['session_id'] = session_id
 
         rows = yield self._simple_select_list(
             table="e2e_room_keys",
             keyvalues=keyvalues,
             retcols=(
+                "user_id",
+                "room_id",
+                "session_id",
                 "first_message_index",
                 "forwarded_count",
                 "is_verified",
@@ -132,19 +133,37 @@ class EndToEndRoomKeyStore(SQLBaseStore):
             desc="get_e2e_room_keys",
         )
 
-        sessions = {}
-        sessions['rooms'][roomId]['sessions'][session_id] = row for row in rows;
-        defer.returnValue(sessions);
+        # perlesque autovivification from https://stackoverflow.com/a/19829714/6764493
+        class AutoVivification(dict):
+            def __getitem__(self, item):
+                try:
+                    return dict.__getitem__(self, item)
+                except KeyError:
+                    value = self[item] = type(self)()
+                    return value
+
+        sessions = AutoVivification()
+        for row in rows:
+            sessions['rooms'][row['room_id']]['sessions'][row['session_id']] = {
+                "first_message_index": row["first_message_index"],
+                "forwarded_count": row["forwarded_count"],
+                "is_verified": row["is_verified"],
+                "session_data": row["session_data"],
+            }
+
+        defer.returnValue(sessions)
 
     @defer.inlineCallbacks
     def delete_e2e_room_keys(self, user_id, version, room_id, session_id):
 
-        keyvalues={
+        keyvalues = {
             "user_id": user_id,
             "version": version,
         }
-        if room_id: keyvalues['room_id'] = room_id
-        if session_id: keyvalues['session_id'] = session_id
+        if room_id:
+            keyvalues['room_id'] = room_id
+        if session_id:
+            keyvalues['session_id'] = session_id
 
         yield self._simple_delete(
             table="e2e_room_keys",
diff --git a/synapse/storage/schema/delta/46/e2e_room_keys.sql b/synapse/storage/schema/delta/46/e2e_room_keys.sql
index 51b826e8b3..6b344c5ad7 100644
--- a/synapse/storage/schema/delta/46/e2e_room_keys.sql
+++ b/synapse/storage/schema/delta/46/e2e_room_keys.sql
@@ -29,7 +29,7 @@ CREATE UNIQUE INDEX e2e_room_keys_user_idx ON e2e_room_keys(user_id);
 CREATE UNIQUE INDEX e2e_room_keys_room_idx ON e2e_room_keys(room_id);
 CREATE UNIQUE INDEX e2e_room_keys_session_idx ON e2e_room_keys(session_id);
 
--- the versioning metadata about versions of users' encrypted e2e session backups
+-- the metadata for each generation of encrypted e2e session backups
 CREATE TABLE e2e_room_key_versions (
     user_id TEXT NOT NULL,
     version INT NOT NULL,