diff --git a/litecord/gateway/websocket.py b/litecord/gateway/websocket.py index 6cc65a2..7526c57 100644 --- a/litecord/gateway/websocket.py +++ b/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) diff --git a/litecord/storage.py b/litecord/storage.py index 27912ec..3324d4b 100644 --- a/litecord/storage.py +++ b/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) diff --git a/litecord/user_storage.py b/litecord/user_storage.py index f02c15c..df0d8d3 100644 --- a/litecord/user_storage.py +++ b/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