Ensure closing of sqlite DB connection when closing DiskCachedAsyncSubstrateInterface

Author: thewhaleking

async_substrate_interface/async_substrate.py

@@ -63,6 +63,7 @@
 from async_substrate_interface.utils.cache import (
     async_sql_lru_cache,
     cached_fetcher,
+    AsyncSqliteDB,
 )
 from async_substrate_interface.utils.decoding import (
     _determine_if_old_runtime_call,
@@ -4026,6 +4027,18 @@ class DiskCachedAsyncSubstrateInterface(AsyncSubstrateInterface):
     Experimental new class that uses disk-caching in addition to memory-caching for the cached methods
     """
 
+    async def close(self):
+        """
+        Closes the substrate connection, and the websocket connection.
+        """
+        try:
+            await self.ws.shutdown()
+        except AttributeError:
+            pass
+        db_conn = AsyncSqliteDB(self.url)
+        if db_conn._db is not None:
+            await db_conn._db.close()
+
     @async_sql_lru_cache(maxsize=SUBSTRATE_CACHE_METHOD_SIZE)
     async def get_parent_block_hash(self, block_hash):
         return await self._get_parent_block_hash(block_hash)

async_substrate_interface/utils/cache.py

@@ -36,7 +36,7 @@ def __new__(cls, chain_endpoint: str):
             cls._instances[chain_endpoint] = instance
             return instance
 
-    async def __call__(self, chain, func, args, kwargs) -> Optional[Any]:
+    async def __call__(self, chain, other_self, func, args, kwargs) -> Optional[Any]:
         if not self._db:
             _ensure_dir()
             self._db = await aiosqlite.connect(CACHE_LOCATION)
@@ -65,26 +65,28 @@ async def __call__(self, chain, func, args, kwargs) -> Optional[Any]:
                           );
                         END;"""
             )
-            key = pickle.dumps((args, kwargs))
+            await self._db.commit()
+            key = pickle.dumps((args, kwargs or None))
             try:
                 cursor: aiosqlite.Cursor = await self._db.execute(
                     f"SELECT value FROM {table_name} WHERE key=? AND chain=?",
                     (key, chain),
                 )
                 result = await cursor.fetchone()
+                await cursor.close()
                 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)
+        result = await func(other_self, *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),
             )
+            await self._db.commit()
         return result
 
 
@@ -200,7 +202,7 @@ def decorator(func):
         @cached_fetcher(max_size=maxsize)
         async def inner(self, *args, **kwargs):
             async_sql_db = AsyncSqliteDB(self.url)
-            result = await async_sql_db(self.url, func, args, kwargs)
+            result = await async_sql_db(self.url, self, func, args, kwargs)
             return result
 
         return inner