diff --git a/botto/storage/config_storage.py b/botto/storage/config_storage.py index d1a8597..d838e31 100644 --- a/botto/storage/config_storage.py +++ b/botto/storage/config_storage.py @@ -43,7 +43,7 @@ async def list_config_by_server(self) -> ConfigCache: async def retrieve_config( self, server_id: str | int, key: Optional[Union[str, int]] - ) -> Optional[ConfigEntry]: + ) -> Optional[ConfigEntry | dict[str, ConfigEntry]]: log.debug(f"Fetching {key or 'config'} for {server_id}") filter_by_formula = f"AND({{Server ID}}='{server_id}'" if key := key: @@ -56,13 +56,13 @@ async def retrieve_config( ) config_iterator = (ConfigEntry.from_airtable(x) async for x in result_iterator) try: - config = await config_iterator.__anext__() async with self.config_lock: server_config = self.config_cache.get(str(server_id), {}) - server_config[config.config_key] = config - self.config_cache[config.server_id] = server_config - return config - except (StopIteration, StopAsyncIteration): + async for config in config_iterator: + server_config[config.config_key] = config + self.config_cache[str(server_id)] = server_config + return server_config if not key else server_config[key] + except (StopIteration, StopAsyncIteration, KeyError): log.info(f"No config found for Key {key} with Server ID {server_id}") return None