add DBWrapper2.managed() (#16880)

Author: altendky

benchmarks/block_store.py

@@ -30,7 +30,6 @@
 from chia.types.blockchain_format.sized_bytes import bytes32
 from chia.types.blockchain_format.sub_epoch_summary import SubEpochSummary
 from chia.types.full_block import FullBlock
-from chia.util.db_wrapper import DBWrapper2
 from chia.util.ints import uint8, uint32, uint64, uint128
 
 # to run this benchmark:
@@ -44,7 +43,6 @@
 
 async def run_add_block_benchmark(version: int) -> None:
     verbose: bool = "--verbose" in sys.argv
-    db_wrapper: DBWrapper2 = await setup_db("block-store-benchmark.db", version)
 
     # keep track of benchmark total time
     all_test_time = 0.0
@@ -54,7 +52,7 @@ async def run_add_block_benchmark(version: int) -> None:
 
     header_hashes = []
 
-    try:
+    async with setup_db("block-store-benchmark.db", version) as db_wrapper:
         block_store = await BlockStore.create(db_wrapper)
 
         block_height = 1
@@ -495,9 +493,6 @@ async def run_add_block_benchmark(version: int) -> None:
         db_size = os.path.getsize(Path("block-store-benchmark.db"))
         print(f"database size: {db_size/1000000:.3f} MB")
 
-    finally:
-        await db_wrapper.close()
-
 
 if __name__ == "__main__":
     print("version 2")

benchmarks/coin_store.py

@@ -12,7 +12,6 @@
 from chia.full_node.coin_store import CoinStore
 from chia.types.blockchain_format.coin import Coin
 from chia.types.blockchain_format.sized_bytes import bytes32
-from chia.util.db_wrapper import DBWrapper2
 from chia.util.ints import uint32, uint64
 
 # to run this benchmark:
@@ -41,12 +40,11 @@ def make_coins(num: int) -> Tuple[List[Coin], List[bytes32]]:
 
 async def run_new_block_benchmark(version: int) -> None:
     verbose: bool = "--verbose" in sys.argv
-    db_wrapper: DBWrapper2 = await setup_db("coin-store-benchmark.db", version)
 
     # keep track of benchmark total time
     all_test_time = 0.0
 
-    try:
+    async with setup_db("coin-store-benchmark.db", version) as db_wrapper:
         coin_store = await CoinStore.create(db_wrapper)
 
         all_unspent: List[bytes32] = []
@@ -301,9 +299,6 @@ async def run_new_block_benchmark(version: int) -> None:
         all_test_time += total_time
         print(f"all tests completed in {all_test_time:0.4f}s")
 
-    finally:
-        await db_wrapper.close()
-
     db_size = os.path.getsize(Path("coin-store-benchmark.db"))
     print(f"database size: {db_size/1000000:.3f} MB")
 

benchmarks/utils.py

@@ -1,15 +1,14 @@
 from __future__ import annotations
 
+import contextlib
 import enum
 import os
 import random
 import subprocess
 import sys
-from datetime import datetime
 from pathlib import Path
-from typing import Any, Generic, Optional, Tuple, Type, TypeVar, Union
+from typing import Any, AsyncIterator, Generic, Optional, Tuple, Type, TypeVar, Union
 
-import aiosqlite
 import click
 from chia_rs import AugSchemeMPL, G1Element, G2Element
 
@@ -179,30 +178,29 @@ def rand_full_block() -> FullBlock:
     return full_block
 
 
-async def setup_db(name: Union[str, os.PathLike[str]], db_version: int) -> DBWrapper2:
+@contextlib.asynccontextmanager
+async def setup_db(name: Union[str, os.PathLike[str]], db_version: int) -> AsyncIterator[DBWrapper2]:
     db_filename = Path(name)
     try:
         os.unlink(db_filename)
     except FileNotFoundError:
         pass
-    connection = await aiosqlite.connect(db_filename)
-
-    def sql_trace_callback(req: str) -> None:
-        sql_log_path = "sql.log"
-        timestamp = datetime.now().strftime("%H:%M:%S.%f")
-        log = open(sql_log_path, "a")
-        log.write(timestamp + " " + req + "\n")
-        log.close()
 
+    log_path: Optional[Path]
     if "--sql-logging" in sys.argv:
-        await connection.set_trace_callback(sql_trace_callback)
-
-    await connection.execute("pragma journal_mode=wal")
-    await connection.execute("pragma synchronous=full")
-
-    ret = DBWrapper2(connection, db_version)
-    await ret.add_connection(await aiosqlite.connect(db_filename))
-    return ret
+        log_path = Path("sql.log")
+    else:
+        log_path = None
+
+    async with DBWrapper2.managed(
+        database=db_filename,
+        log_path=log_path,
+        db_version=db_version,
+        reader_count=1,
+        journal_mode="wal",
+        synchronous="full",
+    ) as db_wrapper:
+        yield db_wrapper
 
 
 def get_commit_hash() -> str:

chia/cmds/check_wallet_db.py

@@ -358,18 +358,18 @@ def check_unexpected_derivation_entries(
 
     async def scan(self, db_path: Path) -> int:
         """Returns number of lines of error output (not warnings)"""
-        self.db_wrapper = await DBWrapper2.create(
+        async with DBWrapper2.managed(
             database=db_path,
             reader_count=self.config.get("db_readers", 4),
             log_path=self.sql_log_path,
             synchronous=db_synchronous_on("auto"),
-        )
-        # TODO: Pass down db_wrapper
-        wallets = await self.get_all_wallets()
-        derivation_paths = await self.get_derivation_paths()
-        errors = []
-        warnings = []
-        try:
+        ) as self.db_wrapper:
+            # TODO: Pass down db_wrapper
+            wallets = await self.get_all_wallets()
+            derivation_paths = await self.get_derivation_paths()
+            errors = []
+            warnings = []
+
             if self.verbose:
                 await self.show_tables()
                 print_min_max_derivation_for_wallets(derivation_paths)
@@ -387,8 +387,7 @@ async def scan(self, db_path: Path) -> int:
             if len(errors) > 0:
                 print(f"    ---- Errors Found for {db_path.name}----")
                 print("\n".join(errors))
-        finally:
-            await self.db_wrapper.close()
+
         return len(errors)