Add memory-cache wrapper for InstallationStore/AsyncInstallationStore

Author: seratch

slack_sdk/oauth/installation_store/async_cacheable_installation_store.py

@@ -0,0 +1,35 @@
+from logging import Logger
+from typing import Optional, Dict
+
+from slack_sdk.oauth.installation_store import Bot, Installation
+from slack_sdk.oauth.installation_store.async_installation_store import AsyncInstallationStore
+
+
+class AsyncCacheableInstallationStore(AsyncInstallationStore):
+    underlying: AsyncInstallationStore
+    cached_bots: Dict[str, Bot]
+
+    def __init__(self, installation_store: AsyncInstallationStore):
+        """A simple memory cache wrapper for any installation stores.
+
+        :param installation_store: the installation store to wrap
+        """
+        self.underlying = installation_store
+        self.cached_bots = {}
+
+    @property
+    def logger(self) -> Logger:
+        return self.underlying.logger
+
+    async def async_save(self, installation: Installation):
+        return await self.underlying.async_save(installation)
+
+    async def async_find_bot(self, *, enterprise_id: Optional[str], team_id: Optional[str]) -> Optional[Bot]:
+        pass
+        key = f"{enterprise_id}-{team_id}"
+        if key in self.cached_bots:
+            return self.cached_bots[key]
+        bot = await self.underlying.async_find_bot(enterprise_id=enterprise_id, team_id=team_id)
+        if bot:
+            self.cached_bots[key] = bot
+        return bot

slack_sdk/oauth/installation_store/cacheable_installation_store.py

@@ -0,0 +1,36 @@
+from logging import Logger
+from typing import Optional, Dict
+
+from slack_sdk.oauth import InstallationStore
+from slack_sdk.oauth.installation_store import Bot, Installation
+
+
+class CacheableInstallationStore(InstallationStore):
+    underlying: InstallationStore
+    cached_bots: Dict[str, Bot]
+
+    def __init__(self, installation_store: InstallationStore):
+        """A simple memory cache wrapper for any installation stores.
+
+        :param installation_store: the installation store to wrap
+        """
+        self.underlying = installation_store
+        self.cached_bots = {}
+
+    @property
+    def logger(self) -> Logger:
+        return self.underlying.logger
+
+    def save(self, installation: Installation):
+        return self.underlying.save(installation)
+
+    def find_bot(
+        self, *, enterprise_id: Optional[str], team_id: Optional[str]
+    ) -> Optional[Bot]:
+        key = f"{enterprise_id}-{team_id}"
+        if key in self.cached_bots:
+            return self.cached_bots[key]
+        bot = self.underlying.find_bot(enterprise_id=enterprise_id, team_id=team_id)
+        if bot:
+            self.cached_bots[key] = bot
+        return bot

tests/slack_sdk/oauth/installation_store/test_simple_cache.py

@@ -0,0 +1,39 @@
+import os
+import unittest
+
+from slack_sdk.oauth.installation_store import Installation
+from slack_sdk.oauth.installation_store.cacheable_installation_store import CacheableInstallationStore
+from slack_sdk.oauth.installation_store.sqlite3 import SQLite3InstallationStore
+
+
+class TestCacheable(unittest.TestCase):
+
+    def test_save_and_find(self):
+        sqlite3_store = SQLite3InstallationStore(
+            database="logs/cacheable.db",
+            client_id="111.222"
+        )
+        sqlite3_store.init()
+        store = CacheableInstallationStore(sqlite3_store)
+
+        installation = Installation(
+            app_id="A111",
+            enterprise_id="E111",
+            team_id="T111",
+            user_id="U111",
+            bot_id="B111",
+            bot_token="xoxb-111",
+            bot_scopes=["chat:write"],
+            bot_user_id="U222"
+        )
+        store.save(installation)
+
+        bot = store.find_bot(enterprise_id="E111", team_id="T111")
+        self.assertIsNotNone(bot)
+
+        os.remove("logs/cacheable.db")
+
+        bot = sqlite3_store.find_bot(enterprise_id="E111", team_id="T111")
+        self.assertIsNone(bot)
+        bot = store.find_bot(enterprise_id="E111", team_id="T111")
+        self.assertIsNotNone(bot)

tests/slack_sdk_async/oauth/__init__.py

@@ -0,0 +1,41 @@
+import os
+import unittest
+
+from slack_sdk.oauth.installation_store import Installation
+from slack_sdk.oauth.installation_store.async_cacheable_installation_store import AsyncCacheableInstallationStore
+from slack_sdk.oauth.installation_store.sqlite3 import SQLite3InstallationStore
+from tests.helpers import async_test
+
+
+class TestCacheable(unittest.TestCase):
+
+    @async_test
+    async def test_save_and_find(self):
+        sqlite3_store = SQLite3InstallationStore(
+            database="logs/cacheable.db",
+            client_id="111.222"
+        )
+        sqlite3_store.init()
+        store = AsyncCacheableInstallationStore(sqlite3_store)
+
+        installation = Installation(
+            app_id="A111",
+            enterprise_id="E111",
+            team_id="T111",
+            user_id="U111",
+            bot_id="B111",
+            bot_token="xoxb-111",
+            bot_scopes=["chat:write"],
+            bot_user_id="U222"
+        )
+        await store.async_save(installation)
+
+        bot = await store.async_find_bot(enterprise_id="E111", team_id="T111")
+        self.assertIsNotNone(bot)
+
+        os.remove("logs/cacheable.db")
+
+        bot = await sqlite3_store.async_find_bot(enterprise_id="E111", team_id="T111")
+        self.assertIsNone(bot)
+        bot = await store.async_find_bot(enterprise_id="E111", team_id="T111")
+        self.assertIsNotNone(bot)