From ec8145c7f6c492eed1d1ecc8f1b2bebb7fa0257d Mon Sep 17 00:00:00 2001
From: David Stephens <git@dwrs.org>
Date: Sat, 8 Jun 2024 21:32:01 +0100
Subject: [PATCH] Fix retrieve_config, so it actually retrieves all entries if
 not key is provided

---
 botto/storage/config_storage.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

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