Create CoinStoreProtocol to remove a dependencies on chia.full_node from chia.consensus (#19741)

* `CoinStoreProtocol`

* Remove `.get_all_coins`

* Apply suggestions from code review

Uses `...` instead of values.

Co-authored-by: Kyle Altendorf <sda@fstab.net>

* Add `.is_empty` to protocol

* Improve test fixture stuff

---------

Co-authored-by: Kyle Altendorf <sda@fstab.net>

chia/_tests/core/consensus/stores/test_coin_store_protocol.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+import pytest
+from chia_rs.sized_bytes import bytes32
+from chia_rs.sized_ints import uint32, uint64
+
+from chia._tests.util.db_connection import DBConnection
+from chia.consensus.coinbase import create_farmer_coin, create_pool_coin
+from chia.full_node.coin_store import CoinStore
+
+# black box tests from `chia/_tests/core/full_node/stores/test_coin_store.py`
+# should be moved here
+
+
+@pytest.mark.anyio
+async def test_is_empty_when_empty(db_version: int) -> None:
+    async with DBConnection(db_version) as db_wrapper:
+        coin_store = await CoinStore.create(db_wrapper)
+        assert await coin_store.is_empty() is True
+
+
+@pytest.mark.anyio
+async def test_is_empty_when_not_empty(db_version: int) -> None:
+    async with DBConnection(db_version) as db_wrapper:
+        coin_store = await CoinStore.create(db_wrapper)
+        assert await coin_store.is_empty() is True
+        height = uint32(1)
+        genesis_challenge = bytes32(b"\0" * 32)
+        pool_puzzle_hash = bytes32(b"\x01" * 32)
+        farmer_puzzle_hash = bytes32(b"\x02" * 32)
+        pool_coin = create_pool_coin(height, pool_puzzle_hash, uint64(1_750_000_000_000), genesis_challenge)
+        farmer_coin = create_farmer_coin(height, farmer_puzzle_hash, uint64(1_750_000_000_000), genesis_challenge)
+        await coin_store.new_block(
+            height=height,
+            timestamp=uint64(1234567890),
+            included_reward_coins=[pool_coin, farmer_coin],
+            tx_additions=[],
+            tx_removals=[],
+        )
+        assert await coin_store.is_empty() is False

chia/_tests/core/full_node/stores/test_coin_store.py

@@ -423,7 +423,7 @@ async def test_get_coin_states(db_version: int) -> None:
             for i in range(1, 301)
         ]
         coin_store = await CoinStore.create(db_wrapper)
-        await add_coin_records_to_db(coin_store.db_wrapper, crs)
+        await add_coin_records_to_db(coin_store, crs)
 
         assert len(await coin_store.get_coin_states_by_puzzle_hashes(True, {std_hash(b"2")}, uint32(0))) == 300
         assert len(await coin_store.get_coin_states_by_puzzle_hashes(False, {std_hash(b"2")}, uint32(0))) == 0
@@ -551,7 +551,7 @@ async def test_coin_state_batches(
         coin_store = await CoinStore.create(db_wrapper)
         hint_store = await HintStore.create(db_wrapper)
 
-        await add_coin_records_to_db(coin_store.db_wrapper, random_coin_records.items)
+        await add_coin_records_to_db(coin_store, random_coin_records.items)
         await hint_store.add_hints(random_coin_records.hints)
 
         # Make sure all of the coin states are found when batching.
@@ -645,7 +645,7 @@ async def test_batch_many_coin_states(db_version: int, cut_off_middle: bool) ->
         coin_store = await CoinStore.create(db_wrapper)
         await HintStore.create(db_wrapper)
 
-        await add_coin_records_to_db(coin_store.db_wrapper, coin_records)
+        await add_coin_records_to_db(coin_store, coin_records)
 
         # Make sure all of the coin states are found.
         (all_coin_states, next_height) = await coin_store.batch_coin_states_by_puzzle_hashes([ph])
@@ -659,7 +659,7 @@ async def test_batch_many_coin_states(db_version: int, cut_off_middle: bool) ->
 
         # For the middle case, insert a coin record between the two heights 10 and 12.
         await add_coin_records_to_db(
-            coin_store.db_wrapper,
+            coin_store,
             [
                 CoinRecord(
                     coin=Coin(std_hash(b"extra coin"), ph, uint64(0)),
@@ -707,7 +707,7 @@ async def test_duplicate_by_hint(db_version: int) -> None:
             uint64(12321312),
         )
 
-        await add_coin_records_to_db(coin_store.db_wrapper, [cr])
+        await add_coin_records_to_db(coin_store, [cr])
         await hint_store.add_hints([(cr.coin.name(), cr.coin.puzzle_hash)])
 
         coin_states, height = await coin_store.batch_coin_states_by_puzzle_hashes([cr.coin.puzzle_hash])
@@ -884,7 +884,7 @@ async def test_add_coin_records_to_db() -> None:
             )
             for i in range(5)
         ]
-        await add_coin_records_to_db(db_wrapper, test_records)
+        await add_coin_records_to_db(coin_store, test_records)
         # Verify all records got inserted correctly
         for record in test_records:
             resulting_record = await coin_store.get_coin_record(record.coin.name())

chia/_tests/core/full_node/test_full_node.py

@@ -47,11 +47,11 @@ from chia._tests.util.time_out_assert import time_out_assert, time_out_assert_cu
 from chia.consensus.augmented_chain import AugmentedBlockchain
 from chia.consensus.block_body_validation import ForkInfo
 from chia.consensus.blockchain import Blockchain
+from chia.consensus.coin_store_protocol import CoinStoreProtocol
 from chia.consensus.get_block_challenge import get_block_challenge
 from chia.consensus.multiprocess_validation import PreValidationResult, pre_validate_block
 from chia.consensus.pot_iterations import is_overflow_block
 from chia.consensus.signage_point import SignagePoint
-from chia.full_node.coin_store import CoinStore
 from chia.full_node.full_node import WalletUpdate
 from chia.full_node.full_node_api import FullNodeAPI
 from chia.full_node.sync_store import Peak
@@ -2505,7 +2505,7 @@ def print_coin_records(records: dict[bytes32, CoinRecord]) -> None:  # pragma: n
         print(f"{rec}")
 
 
-async def validate_coin_set(coin_store: CoinStore, blocks: list[FullBlock]) -> None:
+async def validate_coin_set(coin_store: CoinStoreProtocol, blocks: list[FullBlock]) -> None:
     prev_height = blocks[0].height - 1
     prev_hash = blocks[0].prev_header_hash
     for block in blocks:

chia/_tests/util/coin_store.py

@@ -1,12 +1,15 @@
 from __future__ import annotations
 
+from chia.consensus.coin_store_protocol import CoinStoreProtocol
 from chia.types.coin_record import CoinRecord
 from chia.util.db_wrapper import DBWrapper2
 
 
-async def add_coin_records_to_db(db_wrapper: DBWrapper2, records: list[CoinRecord]) -> None:
+async def add_coin_records_to_db(coin_store: CoinStoreProtocol, records: list[CoinRecord]) -> None:
     if len(records) == 0:
         return
+    db_wrapper = getattr(coin_store, "db_wrapper", None)
+    assert isinstance(db_wrapper, DBWrapper2), "CoinStore must use DBWrapper2"
     async with db_wrapper.writer_maybe_transaction() as conn:
         await conn.executemany(
             "INSERT INTO coin_record VALUES(?, ?, ?, ?, ?, ?, ?, ?)",

chia/_tests/wallet/test_new_wallet_protocol.py

@@ -278,7 +278,7 @@ async def test_request_coin_state(one_node: OneNode, self_hostname: str) -> None
         coinbase=False,
         timestamp=uint64(1),
     )
-    await add_coin_records_to_db(simulator.full_node.coin_store.db_wrapper, [*coin_records, ignored_coin])
+    await add_coin_records_to_db(simulator.full_node.coin_store, [*coin_records, ignored_coin])
 
     # Request no coin states
     resp = await simulator.request_coin_state(wallet_protocol.RequestCoinState([], None, genesis, False), peer)
@@ -378,7 +378,7 @@ async def test_request_coin_state_limit(one_node: OneNode, self_hostname: str) -
             )
             coin_records[coin_record.coin.name()] = coin_record
 
-    await add_coin_records_to_db(simulator.full_node.coin_store.db_wrapper, list(coin_records.values()))
+    await add_coin_records_to_db(simulator.full_node.coin_store, list(coin_records.values()))
 
     # Fetch the coin records using the wallet protocol,
     # with more coin ids than the limit of 100,000, but only after height 10000.
@@ -441,7 +441,7 @@ async def test_request_puzzle_state(one_node: OneNode, self_hostname: str) -> No
         timestamp=uint64(1),
     )
 
-    await add_coin_records_to_db(simulator.full_node.coin_store.db_wrapper, [*coin_records, ignored_coin])
+    await add_coin_records_to_db(simulator.full_node.coin_store, [*coin_records, ignored_coin])
 
     # We already test permutations of CoinStateFilters in the CoinStore tests
     # So it's redundant to do so here
@@ -570,7 +570,7 @@ async def test_request_puzzle_state_limit(one_node: OneNode, self_hostname: str)
             )
             coin_records[coin_record.coin.name()] = coin_record
 
-    await add_coin_records_to_db(simulator.full_node.coin_store.db_wrapper, list(coin_records.values()))
+    await add_coin_records_to_db(simulator.full_node.coin_store, list(coin_records.values()))
 
     # Fetch the coin records using the wallet protocol,
     # only after height 10000, so that the limit of 100000 isn't exceeded
@@ -724,7 +724,7 @@ async def test_sync_puzzle_state(
             if coin_ph != puzzle_hash:
                 hints.append((coin.name(), puzzle_hash))
 
-    await add_coin_records_to_db(simulator.full_node.coin_store.db_wrapper, list(coin_records.values()))
+    await add_coin_records_to_db(simulator.full_node.coin_store, list(coin_records.values()))
     await simulator.full_node.hint_store.add_hints(hints)
 
     # Farm peak

chia/consensus/blockchain.py

@@ -26,6 +26,7 @@ from chia_rs.sized_ints import uint16, uint32, uint64, uint128
 
 from chia.consensus.block_body_validation import ForkInfo, validate_block_body
 from chia.consensus.block_header_validation import validate_unfinished_header_block
+from chia.consensus.coin_store_protocol import CoinStoreProtocol
 from chia.consensus.cost_calculator import NPCResult
 from chia.consensus.difficulty_adjustment import get_next_sub_slot_iters_and_difficulty
 from chia.consensus.find_fork_point import lookup_fork_chain
@@ -35,7 +36,6 @@ from chia.consensus.get_block_generator import get_block_generator
 from chia.consensus.multiprocess_validation import PreValidationResult
 from chia.full_node.block_height_map import BlockHeightMap
 from chia.full_node.block_store import BlockStore
-from chia.full_node.coin_store import CoinStore
 from chia.types.blockchain_format.coin import Coin
 from chia.types.blockchain_format.vdf import VDFInfo
 from chia.types.coin_record import CoinRecord
@@ -102,7 +102,7 @@ class Blockchain:
     # epoch summaries
     __height_map: BlockHeightMap
     # Unspent Store
-    coin_store: CoinStore
+    coin_store: CoinStoreProtocol
     # Store
     block_store: BlockStore
     # Used to verify blocks in parallel
@@ -121,7 +121,7 @@ class Blockchain:
 
     @staticmethod
     async def create(
-        coin_store: CoinStore,
+        coin_store: CoinStoreProtocol,
         block_store: BlockStore,
         height_map: BlockHeightMap,
         consensus_constants: ConsensusConstants,

chia/consensus/coin_store_protocol.py

@@ -0,0 +1,151 @@
+from __future__ import annotations
+
+from collections.abc import Collection
+from typing import Optional, Protocol
+
+from chia_rs import CoinState
+from chia_rs.sized_bytes import bytes32
+from chia_rs.sized_ints import uint32, uint64
+
+from chia.types.blockchain_format.coin import Coin
+from chia.types.coin_record import CoinRecord
+from chia.types.mempool_item import UnspentLineageInfo
+
+
+class CoinStoreProtocol(Protocol):
+    """
+    Protocol defining the interface for CoinStore.
+    This is a substitute for importing from chia.full_node.coin_store directly.
+    """
+
+    async def new_block(
+        self,
+        height: uint32,
+        timestamp: uint64,
+        included_reward_coins: Collection[Coin],
+        tx_additions: Collection[tuple[bytes32, Coin]],
+        tx_removals: list[bytes32],
+    ) -> None:
+        """
+        Add a new block to the coin store
+        """
+
+    async def get_coin_record(self, coin_id: bytes32) -> Optional[CoinRecord]:
+        """
+        Returns the coin record for the specified coin id
+        """
+
+    async def get_coin_records(self, coin_ids: Collection[bytes32]) -> list[CoinRecord]:
+        """
+        Returns the coin records for the specified coin ids
+        """
+
+    async def get_coins_added_at_height(self, height: uint32) -> list[CoinRecord]:
+        """
+        Returns the coins added at a specific height
+        """
+
+    async def get_coins_removed_at_height(self, height: uint32) -> list[CoinRecord]:
+        """
+        Returns the coins removed at a specific height
+        """
+
+    async def get_coin_records_by_puzzle_hash(
+        self,
+        include_spent_coins: bool,
+        puzzle_hash: bytes32,
+        start_height: uint32 = ...,
+        end_height: uint32 = ...,
+    ) -> list[CoinRecord]:
+        """
+        Returns the coin records for a specific puzzle hash
+        """
+
+    async def get_coin_records_by_puzzle_hashes(
+        self,
+        coins: bool,
+        puzzle_hashes: list[bytes32],
+        start_height: uint32 = ...,
+        end_height: uint32 = ...,
+    ) -> list[CoinRecord]:
+        """
+        Returns the coin records for a list of puzzle hashes
+        """
+
+    async def get_coin_records_by_names(
+        self,
+        include_spent_coins: bool,
+        names: list[bytes32],
+        start_height: uint32 = ...,
+        end_height: uint32 = ...,
+    ) -> list[CoinRecord]:
+        """
+        Returns the coin records for a list of coin names
+        """
+
+    async def get_coin_states_by_puzzle_hashes(
+        self,
+        include_spent_coins: bool,
+        puzzle_hashes: set[bytes32],
+        min_height: uint32 = uint32(0),
+        *,
+        max_items: int = ...,
+    ) -> set[CoinState]:
+        """
+        Returns the coin states for a set of puzzle hashes
+        """
+
+    async def get_coin_records_by_parent_ids(
+        self,
+        include_spent_coins: bool,
+        parent_ids: list[bytes32],
+        start_height: uint32 = ...,
+        end_height: uint32 = ...,
+    ) -> list[CoinRecord]:
+        """
+        Returns the coin records for a list of parent ids
+        """
+
+    async def get_coin_states_by_ids(
+        self,
+        include_spent_coins: bool,
+        coin_ids: Collection[bytes32],
+        min_height: uint32 = uint32(0),
+        *,
+        max_height: uint32 = ...,
+        max_items: int = ...,
+    ) -> list[CoinState]:
+        """
+        Returns the coin states for a collection of coin ids
+        """
+
+    async def batch_coin_states_by_puzzle_hashes(
+        self,
+        puzzle_hashes: list[bytes32],
+        *,
+        min_height: uint32 = ...,
+        include_spent: bool = ...,
+        include_unspent: bool = ...,
+        include_hinted: bool = ...,
+        min_amount: uint64 = ...,
+        max_items: int = ...,
+    ) -> tuple[list[CoinState], Optional[uint32]]:
+        """
+        Returns the coin states, as well as the next block height (or `None` if finished).
+        """
+
+    async def get_unspent_lineage_info_for_puzzle_hash(self, puzzle_hash: bytes32) -> Optional[UnspentLineageInfo]:
+        """
+        Lookup the most recent unspent lineage that matches a puzzle hash
+        """
+
+    async def rollback_to_block(self, block_index: int) -> dict[bytes32, CoinRecord]:
+        """
+        Rolls back the blockchain to the specified block index
+        """
+
+    # DEPRECATED: do not use in new code
+    async def is_empty(self) -> bool:
+        """
+        Returns True if the coin store is empty
+        """

chia/full_node/coin_store.py

@@ -217,23 +217,6 @@ class CoinStore:
                         coins.append(coin_record)
                 return coins
 
-    async def get_all_coins(self, include_spent_coins: bool) -> list[CoinRecord]:
-        # WARNING: this should only be used for testing or in a simulation,
-        # running it on a synced testnet or mainnet node will most likely result in an OOM error.
-        coins = set()
-
-        async with self.db_wrapper.reader_no_transaction() as conn:
-            async with conn.execute(
-                f"SELECT confirmed_index, spent_index, coinbase, puzzle_hash, "
-                f"coin_parent, amount, timestamp FROM coin_record "
-                f"{'' if include_spent_coins else 'INDEXED BY coin_spent_index WHERE spent_index=0'}"
-                f" ORDER BY confirmed_index"
-            ) as cursor:
-                for row in await cursor.fetchall():
-                    coin = self.row_to_coin(row)
-                    coins.add(CoinRecord(coin, row[0], row[1], row[2], row[6]))
-                return list(coins)
-
     # Checks DB and DiffStores for CoinRecords with puzzle_hash and returns them
     async def get_coin_records_by_puzzle_hash(
         self,
@@ -622,3 +605,12 @@ class CoinStore:
                 return UnspentLineageInfo(
                     coin_id=bytes32(coin_id), parent_id=bytes32(parent_id), parent_parent_id=bytes32(parent_parent_id)
                 )
+
+    async def is_empty(self) -> bool:
+        """
+        Returns True if the coin store is empty, False otherwise.
+        """
+        async with self.db_wrapper.reader_no_transaction() as conn:
+            async with conn.execute("SELECT coin_name FROM coin_record LIMIT 1") as cursor:
+                row = await cursor.fetchone()
+                return row is None or len(row) == 0

chia/full_node/full_node.py

@@ -41,6 +41,7 @@ from chia.consensus.block_body_validation import ForkInfo
 from chia.consensus.block_creation import unfinished_block_to_full_block
 from chia.consensus.blockchain import AddBlockResult, Blockchain, BlockchainMutexPriority, StateChangeSummary
 from chia.consensus.blockchain_interface import BlockchainInterface
+from chia.consensus.coin_store_protocol import CoinStoreProtocol
 from chia.consensus.condition_tools import pkm_pairs
 from chia.consensus.cost_calculator import NPCResult
 from chia.consensus.difficulty_adjustment import get_next_sub_slot_iters_and_difficulty
@@ -156,7 +157,7 @@ class FullNode:
     _db_wrapper: Optional[DBWrapper2] = None
     _hint_store: Optional[HintStore] = None
     _block_store: Optional[BlockStore] = None
-    _coin_store: Optional[CoinStore] = None
+    _coin_store: Optional[CoinStoreProtocol] = None
     _mempool_manager: Optional[MempoolManager] = None
     _init_weight_proof: Optional[asyncio.Task[None]] = None
     _blockchain: Optional[Blockchain] = None
@@ -306,10 +307,9 @@ class FullNode:
             peak: Optional[BlockRecord] = self.blockchain.get_peak()
             if peak is None:
                 self.log.info(f"Initialized with empty blockchain time taken: {int(time_taken)}s")
-                num_unspent = await self.coin_store.num_unspent()
-                if num_unspent > 0:
+                if not await self.coin_store.is_empty():
                     self.log.error(
-                        f"Inconsistent blockchain DB file! Could not find peak block but found {num_unspent} coins! "
+                        "Inconsistent blockchain DB file! Could not find peak block but found some coins! "
                         "This is a fatal error. The blockchain database may be corrupt"
                     )
                     raise RuntimeError("corrupt blockchain DB")
@@ -350,9 +350,9 @@ class FullNode:
             self.initialized = True
 
             try:
-                async with contextlib.AsyncExitStack() as exit_stack:
+                async with contextlib.AsyncExitStack() as aexit_stack:
                     if self.full_node_peers is not None:
-                        await exit_stack.enter_async_context(self.full_node_peers.manage())
+                        await aexit_stack.enter_async_context(self.full_node_peers.manage())
                     yield
             finally:
                 self._shut_down = True
@@ -419,7 +419,7 @@ class FullNode:
         return self._blockchain
 
     @property
-    def coin_store(self) -> CoinStore:
+    def coin_store(self) -> CoinStoreProtocol:
         assert self._coin_store is not None
         return self._coin_store
 

chia/simulator/full_node_simulator.py

@@ -123,7 +123,33 @@ class FullNodeSimulator(FullNodeAPI):
             return self.auto_farm
 
     async def get_all_coins(self, request: GetAllCoinsProtocol) -> list[CoinRecord]:
-        return await self.full_node.coin_store.get_all_coins(request.include_spent_coins)
+        """
+        Simulates fetching all coins by querying coins added at each block height.
+
+        Args:
+            request: An object containing the `include_spent_coins` flag.
+
+        Returns:
+            A combined list of CoinRecords (including spent coins if requested).
+        """
+        coin_records: list[CoinRecord] = []
+        current_height = 0
+
+        # `.get_peak_height` can return `None`. We use -1 in that case to exit early
+        max_block_height = self.full_node.blockchain.get_peak_height() or -1
+
+        while current_height <= max_block_height:
+            # Fetch coins added at the current block height
+            records_at_height = await self.full_node.coin_store.get_coins_added_at_height(uint32(current_height))
+
+            if not request.include_spent_coins:
+                # Filter out spent coins if not requested
+                records_at_height = [record for record in records_at_height if not record.spent]
+
+            coin_records.extend(records_at_height)
+            current_height += 1
+
+        return coin_records
 
     async def revert_block_height(self, new_height: uint32) -> None:
         """