Defer to DBWrapper2

Author: richardkiss

chia/_tests/blockchain/test_blockchain.py

@@ -3879,7 +3879,7 @@ async def test_chain_failed_rollback(empty_blockchain: Blockchain, bt: BlockTool
     print(f"{await b.consensus_store.get_coin_record(spend_bundle.coin_spends[0].coin.name())}")
     print(spend_bundle.coin_spends[0].coin.name())
     # await b.consensus_store._set_spent([spend_bundle.coin_spends[0].coin.name()], 8)
-    async with b.consensus_store as writer:
+    async with b.consensus_store.writer() as writer:
         await writer.rollback_to_block(2)
     print(f"{await b.consensus_store.get_coin_record(spend_bundle.coin_spends[0].coin.name())}")
 

chia/consensus/blockchain.py

@@ -413,7 +413,7 @@ async def add_block(
 
         try:
             # Always add the block to the database
-            async with self.consensus_store as writer:
+            async with self.consensus_store.writer() as writer:
                 # Perform the DB operations to update the state, and rollback if something goes wrong
                 await writer.add_full_block(header_hash, block, block_record)
                 records, state_change_summary = await self._reconsider_peak(block_record, genesis, fork_info)
@@ -486,7 +486,7 @@ async def _reconsider_peak(
         if genesis and peak is not None:
             return [], None
 
-        async with self.consensus_store as writer:
+        async with self.consensus_store.writer() as writer:
             if peak is not None:
                 if block_record.weight < peak.weight:
                     # This is not a heavier block than the heaviest we have seen, so we don't change the coin set
@@ -981,7 +981,7 @@ def add_block_record(self, block_record: BlockRecord) -> None:
     async def persist_sub_epoch_challenge_segments(
         self, ses_block_hash: bytes32, segments: list[SubEpochChallengeSegment]
     ) -> None:
-        async with self.consensus_store as writer:
+        async with self.consensus_store.writer() as writer:
             await writer.persist_sub_epoch_challenge_segments(ses_block_hash, segments)
 
     async def get_sub_epoch_challenge_segments(

chia/consensus/consensus_store_protocol.py

@@ -1,12 +1,12 @@
 from __future__ import annotations
 
 from collections.abc import Collection
-from types import TracebackType
 from typing import Optional, Protocol
 
 from chia_rs import BlockRecord, FullBlock, SubEpochChallengeSegment, SubEpochSummary
 from chia_rs.sized_bytes import bytes32
 from chia_rs.sized_ints import uint32, uint64
+from typing_extensions import AsyncContextManager
 
 from chia.types.blockchain_format.coin import Coin
 from chia.types.coin_record import CoinRecord
@@ -46,28 +46,22 @@ class ConsensusStoreProtocol(Protocol):
     """
     Read-only protocol for the consensus store.
 
-    This protocol also acts as an async context manager. Entering the context
+    This protocol is callable and returns an async context manager. Entering the context
     yields a ConsensusStoreWriteProtocol instance, which must be used for
     performing write (mutating) operations. This ensures atomic writes and
     makes it harder to accidentally perform writes outside a transaction.
 
     Example usage:
-        async with store as writer:
+        async with store.writer() as writer:
             await writer.add_full_block(...)
             await writer.set_peak(...)
 
         # Outside the context, only read methods are available
         br = await store.get_block_record(header_hash)
     """
 
-    # Async context manager methods
-    async def __aenter__(self) -> ConsensusStoreWriteProtocol: ...
-    async def __aexit__(
-        self,
-        exc_type: Optional[type[BaseException]],
-        exc: Optional[BaseException],
-        tb: Optional[TracebackType],
-    ) -> Optional[bool]: ...
+    # Writer method that returns async context manager
+    def writer(self) -> AsyncContextManager[ConsensusStoreWriteProtocol]: ...
 
     # Block store reads
     async def get_block_records_close_to_peak(

chia/full_node/consensus_store_sqlite3.py

@@ -1,11 +1,10 @@
 from __future__ import annotations
 
-from collections.abc import Collection
-from contextlib import AbstractAsyncContextManager
+from collections.abc import AsyncIterator, Collection
+from contextlib import asynccontextmanager
 from dataclasses import dataclass
 from pathlib import Path
-from types import TracebackType
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Optional
 
 from chia_rs import BlockRecord, FullBlock, SubEpochChallengeSegment, SubEpochSummary
 from chia_rs.sized_bytes import bytes32
@@ -55,6 +54,12 @@ async def new_block(
     ) -> None:
         await self.coin_store.new_block(height, timestamp, included_reward_coins, tx_additions, tx_removals)
 
+    @asynccontextmanager
+    async def writer(self) -> AsyncIterator[ConsensusStoreSQLite3Writer]:
+        # Return self as the writer facade
+        async with self.block_store.transaction():
+            yield self
+
 
 @dataclass
 class ConsensusStoreSQLite3:
@@ -66,11 +71,6 @@ class ConsensusStoreSQLite3:
     coin_store: CoinStoreProtocol
     height_map: BlockHeightMap
 
-    # Writer context and writer facade for transactional writes (re-entrant via depth counter)
-    _writer_ctx: Optional[AbstractAsyncContextManager[Any]] = None
-    _writer: Optional[ConsensusStoreSQLite3Writer] = None
-    _txn_depth: int = 0
-
     @classmethod
     async def create(
         cls,
@@ -99,41 +99,12 @@ async def create(
             height_map=height_map,
         )
 
-    # Async context manager yielding a writer for atomic writes
-    async def __aenter__(self) -> ConsensusStoreSQLite3Writer:
-        # Re-entrant async context manager:
-        # Begin a transaction only at the outermost level. CoinStore shares the same DB.
-        if self._writer is None:
-            self._writer_ctx = self.block_store.transaction()
-            await self._writer_ctx.__aenter__()
-            # Create writer facade bound to this transaction
-            self._writer = ConsensusStoreSQLite3Writer(self.block_store, self.coin_store)
-        self._txn_depth += 1
-        return self._writer  # Return the same writer for nested contexts
-
-    async def __aexit__(
-        self,
-        exc_type: Optional[type[BaseException]],
-        exc: Optional[BaseException],
-        tb: Optional[TracebackType],
-    ) -> Optional[bool]:
-        try:
-            # Check if we're at the outermost level before decrementing
-            if self._txn_depth == 1:
-                # This is the outermost context, handle transaction exit
-                if self._writer_ctx is not None:
-                    return await self._writer_ctx.__aexit__(exc_type, exc, tb)
-                return None
-            else:
-                # This is a nested context, just return None (don't suppress exceptions)
-                return None
-        finally:
-            # Always decrement depth and clean up if we're at the outermost level
-            if self._txn_depth > 0:
-                self._txn_depth -= 1
-            if self._txn_depth == 0:
-                self._writer_ctx = None
-                self._writer = None
+    @asynccontextmanager
+    async def writer(self) -> AsyncIterator[ConsensusStoreSQLite3Writer]:
+        """Async context manager that yields a writer facade for performing transactional writes."""
+        csw = ConsensusStoreSQLite3Writer(self.block_store, self.coin_store)
+        async with csw.writer() as writer:
+            yield writer
 
     # Block store methods