matrix-org · squahtx · Dec 14, 2021 · Dec 9, 2021 · Dec 9, 2021 · Dec 13, 2021
diff --git a/changelog.d/11548.misc b/changelog.d/11548.misc
@@ -0,0 +1 @@
+Add missing type hints to storage classes.
@@ -28,7 +28,6 @@ exclude = (?x)
    |synapse/storage/databases/main/account_data.py
    |synapse/storage/databases/main/cache.py
    |synapse/storage/databases/main/devices.py
-   |synapse/storage/databases/main/e2e_room_keys.py
    |synapse/storage/databases/main/end_to_end_keys.py
    |synapse/storage/databases/main/event_federation.py
    |synapse/storage/databases/main/event_push_actions.py
@@ -187,6 +186,9 @@ disallow_untyped_defs = True
 [mypy-synapse.storage.databases.main.directory]
 disallow_untyped_defs = True
 
+[mypy-synapse.storage.databases.main.e2e_room_keys]
+disallow_untyped_defs = True
+
 [mypy-synapse.storage.databases.main.events_worker]
 disallow_untyped_defs = True
 

@@ -14,7 +14,9 @@
 # limitations under the License.
 
 import logging
-from typing import TYPE_CHECKING, List, Optional
+from typing import TYPE_CHECKING, Dict, Optional
+
+from typing_extensions import Literal
 
 from synapse.api.errors import (
     Codes,
@@ -24,6 +26,7 @@
     SynapseError,
 )
 from synapse.logging.opentracing import log_kv, trace
+from synapse.storage.databases.main.e2e_room_keys import RoomKey
 from synapse.types import JsonDict
 from synapse.util.async_helpers import Linearizer
 
@@ -58,7 +61,9 @@ async def get_room_keys(
         version: str,
         room_id: Optional[str] = None,
         session_id: Optional[str] = None,
-    ) -> List[JsonDict]:
+    ) -> Dict[
+        Literal["rooms"], Dict[str, Dict[Literal["sessions"], Dict[str, RoomKey]]]
+    ]:
 room_keys = await self.e2e_room_keys_handler.get_room_keys( 
     user_id, version, room_id, session_id 
 ) 
 # Convert room_keys to the right format to return. 
 if session_id: 
     # If the client requests a specific session, but that session was 
     # not backed up, then return an M_NOT_FOUND. 
     if room_keys["rooms"] == {}: 
         raise NotFoundError("No room_keys found") 
     else: 
         room_keys = room_keys["rooms"][room_id]["sessions"][session_id] 
 elif room_id: 
     # If the client requests all sessions from a room, but no sessions 
     # are found, then return an empty result rather than an error, so 
     # that clients don't have to handle an error condition, and an 
     # empty result is valid.  (Similarly if the client requests all 
     # sessions from the backup, but in that case, room_keys is already 
     # in the right format, so we don't need to do anything about it.) 
     if room_keys["rooms"] == {}: 
         room_keys = {"sessions": {}} 
     else: 
         room_keys = room_keys["rooms"][room_id] 
 return 200, room_keys 
 res = self.get_success( 
     self.handler.get_room_keys( 
         self.local_user, version, room_id="!abc:matrix.org" 
     ) 
 ) 
 self.assertDictEqual(res, room_keys) 
 room_keys = await self.e2e_room_keys_handler.get_room_keys( 
     user_id, version, room_id, session_id 
 ) 
  
 # Convert room_keys to the right format to return. 
 if session_id: 
     # If the client requests a specific session, but that session was 
     # not backed up, then return an M_NOT_FOUND. 
     if room_keys["rooms"] == {}: 
         raise NotFoundError("No room_keys found") 
     else: 
         room_keys = room_keys["rooms"][room_id]["sessions"][session_id] 
 elif room_id: 
     # If the client requests all sessions from a room, but no sessions 
     # are found, then return an empty result rather than an error, so 
     # that clients don't have to handle an error condition, and an 
     # empty result is valid.  (Similarly if the client requests all 
     # sessions from the backup, but in that case, room_keys is already 
     # in the right format, so we don't need to do anything about it.) 
     if room_keys["rooms"] == {}: 
         room_keys = {"sessions": {}} 
     else: 
         room_keys = room_keys["rooms"][room_id] 
  
 return 200, room_keys 
 res = self.get_success( 
     self.handler.get_room_keys( 
         self.local_user, version, room_id="!abc:matrix.org" 
     ) 
 ) 
 self.assertDictEqual(res, room_keys) 
         """Bulk get the E2E room keys for a given backup, optionally filtered to a given
         room, or a given session.
         See EndToEndRoomKeyStore.get_e2e_room_keys for full details.
@@ -72,8 +77,8 @@ async def get_room_keys(
         Raises:
             NotFoundError: if the backup version does not exist
         Returns:
-            A list of dicts giving the session_data and message metadata for
-            these room keys.
+            A dict giving the session_data and message metadata for these room keys.
+            `{"rooms": {room_id: {"sessions": {session_id: room_key}}}}`
         """
 
         # we deliberately take the lock to get keys so that changing the version
@@ -273,7 +278,7 @@ async def upload_room_keys(
 
     @staticmethod
     def _should_replace_room_key(
-        current_room_key: Optional[JsonDict], room_key: JsonDict
+        current_room_key: Optional[RoomKey], room_key: RoomKey
     ) -> bool:
         """
         Determine whether to replace a given current_room_key (if any)