WIP check-in

Author: thewhaleking

async_substrate_interface/utils/cache.py

@@ -9,6 +9,9 @@
 from pathlib import Path
 from typing import Callable, Any, Awaitable, Hashable, Optional
 
+import aiosqlite
+
+
 USE_CACHE = True if os.getenv("NO_CACHE") != "1" else False
 CACHE_LOCATION = (
     os.path.expanduser(
@@ -21,6 +24,70 @@
 logger = logging.getLogger("async_substrate_interface")
 
 
+class AsyncSqliteDB:
+    _instances: dict[str, "AsyncSqliteDB"] = {}
+    _db: Optional[aiosqlite.Connection] = None
+
+    def __new__(cls, chain_endpoint: str):
+        try:
+            return cls._instances[chain_endpoint]
+        except KeyError:
+            instance = super().__new__(cls)
+            cls._instances[chain_endpoint] = instance
+            return instance
+
+    async def __call__(self, chain, func, args, kwargs) -> Optional[Any]:
+        if not self._db:
+            _ensure_dir()
+            self._db = await aiosqlite.connect(CACHE_LOCATION)
+        table_name = _get_table_name(func)
+        key = None
+        if not (local_chain := _check_if_local(chain)) or not USE_CACHE:
+            await self._db.execute(
+                f"""CREATE TABLE IF NOT EXISTS {table_name} 
+                    (
+                       rowid INTEGER PRIMARY KEY AUTOINCREMENT,
+                       key BLOB,
+                       value BLOB,
+                       chain TEXT,
+                       created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+                    );
+                    """
+            )
+            await self._db.execute(
+                f"""CREATE TRIGGER IF NOT EXISTS prune_rows_trigger AFTER INSERT ON {table_name}
+                        BEGIN
+                          DELETE FROM {table_name}
+                          WHERE rowid IN (
+                            SELECT rowid FROM {table_name}
+                            ORDER BY created_at DESC
+                            LIMIT -1 OFFSET 500
+                          );
+                        END;"""
+            )
+            key = pickle.dumps((args, kwargs))
+            try:
+                cursor: aiosqlite.Cursor = await self._db.execute(
+                    f"SELECT value FROM {table_name} WHERE key=? AND chain=?",
+                    (key, chain),
+                )
+                result = await cursor.fetchone()
+                if result is not None:
+                    return pickle.loads(result[0])
+            except (pickle.PickleError, sqlite3.Error) as e:
+                logger.exception("Cache error", exc_info=e)
+                pass
+
+        result = await func(*args, **kwargs)
+        if not local_chain or not USE_CACHE:
+            # TODO use a task here
+            await self._db.execute(
+                f"INSERT OR REPLACE INTO {table_name} (key, value, chain) VALUES (?,?,?)",
+                (key, pickle.dumps(result), chain),
+            )
+        return result
+
+
 def _ensure_dir():
     path = Path(CACHE_LOCATION).parent
     if not path.exists():
@@ -115,7 +182,8 @@ def inner(self, *args, **kwargs):
             )
 
             # If not in DB, call func and store in DB
-            result = func(self, *args, **kwargs)
+            if result is None:
+                result = func(self, *args, **kwargs)
 
             if not local_chain or not USE_CACHE:
                 _insert_into_cache(c, conn, table_name, key, result, chain)
@@ -131,15 +199,8 @@ def async_sql_lru_cache(maxsize: Optional[int] = None):
     def decorator(func):
         @cached_fetcher(max_size=maxsize)
         async def inner(self, *args, **kwargs):
-            c, conn, table_name, key, result, chain, local_chain = (
-                _shared_inner_fn_logic(func, self, args, kwargs)
-            )
-
-            # If not in DB, call func and store in DB
-            result = await func(self, *args, **kwargs)
-            if not local_chain or not USE_CACHE:
-                _insert_into_cache(c, conn, table_name, key, result, chain)
-
+            async_sql_db = AsyncSqliteDB(self.url)
+            result = await async_sql_db(self.url, func, args, kwargs)
             return result
 
         return inner

pyproject.toml

@@ -11,7 +11,8 @@ dependencies = [
   "bt-decode==v0.6.0",
   "scalecodec~=1.2.11",
   "websockets>=14.1",
-  "xxhash"
+  "xxhash",
+  "aiosqlite>=0.21.0,<1.0.0"
 ]
 
 requires-python = ">=3.9,<3.14"