gateway.websocket: fix more Storage -> UserStorage references

litecord/gateway/websocket.py

@@ -91,6 +91,7 @@ class GatewayWebsocket:
         )
 
         self.storage = self.ext.storage
+        self.user_storage = self.ext.user_storage
         self.presence = self.ext.presence
         self.ws = ws
 
@@ -358,7 +359,7 @@ class GatewayWebsocket:
         await self.ext.dispatcher.sub_many('guild', user_id, guild_ids)
 
         # subscribe the user to all dms they have OPENED.
-        dms = await self.storage.get_dms(user_id)
+        dms = await self.user_storage.get_dms(user_id)
         dm_ids = [int(dm['id']) for dm in dms]
 
         log.info('subscribing to {} dms', len(dm_ids))
@@ -367,7 +368,7 @@ class GatewayWebsocket:
         # subscribe to all friends
         # (their friends will also subscribe back
         #  when they come online)
-        friend_ids = await self.storage.get_friend_ids(user_id)
+        friend_ids = await self.user_storage.get_friend_ids(user_id)
         log.info('subscribing to {} friends', len(friend_ids))
         await self.ext.dispatcher.sub_many('friend', user_id, friend_ids)
 

litecord/storage.py

@@ -67,7 +67,7 @@ class Storage:
         self.db = db
         self.presence = None
 
-    async def _fetchrow_with_json(self, query: str, *args):
+    async def fetchrow_with_json(self, query: str, *args):
         """Fetch a single row with JSON/JSONB support."""
         # the pool by itself doesn't have
         # set_type_codec, so we must set it manually
@@ -76,7 +76,7 @@ class Storage:
             await _set_json(con)
             return await con.fetchrow(query, *args)
 
-    async def _fetch_with_json(self, query: str, *args):
+    async def fetch_with_json(self, query: str, *args):
         """Fetch many rows with JSON/JSONB support."""
         async with self.db.acquire() as con:
             await _set_json(con)

litecord/user_storage.py

@@ -25,7 +25,7 @@ class UserStorage:
 
     async def get_user_settings(self, user_id: int) -> Dict[str, Any]:
         """Get current user settings."""
-        row = await self._fetchrow_with_json("""
+        row = await self.storage.fetchrow_with_json("""
         SELECT *
         FROM user_settings
         WHERE id = $1