chia-blockchain 2.5.4rc2__py3-none-any.whl → 2.5.5__py3-none-any.whl

chia/full_node/coin_store.py

@@ -5,21 +5,19 @@ import logging
 import sqlite3
 import time
 from collections.abc import Collection
-from typing import Any, Optional
+from typing import Any, ClassVar, Optional
 
 import typing_extensions
 from aiosqlite import Cursor
+from chia_rs import CoinState
 from chia_rs.sized_bytes import bytes32
 from chia_rs.sized_ints import uint32, uint64
-from clvm.casts import int_from_bytes
 
-from chia.protocols.wallet_protocol import CoinState
 from chia.types.blockchain_format.coin import Coin
 from chia.types.coin_record import CoinRecord
-from chia.types.eligible_coin_spends import UnspentLineageInfo
+from chia.types.mempool_item import UnspentLineageInfo
 from chia.util.batches import to_batches
 from chia.util.db_wrapper import SQLITE_MAX_VARIABLE_NUMBER, DBWrapper2
-from chia.util.lru_cache import LRUCache
 
 log = logging.getLogger(__name__)
 
@@ -32,13 +30,15 @@ class CoinStore:
     """
 
     db_wrapper: DBWrapper2
-    coins_added_at_height_cache: LRUCache[uint32, list[CoinRecord]]
+    # Fall back to the `coin_puzzle_hash` index if the ff unspent index
+    # does not exist.
+    _unspent_lineage_for_ph_idx: str = "coin_puzzle_hash"
 
     @classmethod
     async def create(cls, db_wrapper: DBWrapper2) -> CoinStore:
         if db_wrapper.db_version != 2:
             raise RuntimeError(f"CoinStore does not support database schema v{db_wrapper.db_version}")
-        self = CoinStore(db_wrapper, LRUCache(100))
+        self = CoinStore(db_wrapper)
 
         async with self.db_wrapper.writer_maybe_transaction() as conn:
             log.info("DB: Creating coin store tables and indexes.")
@@ -69,11 +69,34 @@ class CoinStore:
             log.info("DB: Creating index coin_parent_index")
             await conn.execute("CREATE INDEX IF NOT EXISTS coin_parent_index on coin_record(coin_parent)")
 
+            async with conn.execute("SELECT 1 FROM coin_record LIMIT 1") as cursor:
+                is_new_db = await cursor.fetchone() is None
+            if is_new_db:
+                log.info("DB: Creating index coin_record_ph_ff_unspent_idx")
+                # This partial index optimizes fast forward singleton latest
+                # unspent queries. We're only adding it to new DBs to avoid
+                # complex migrations that affect the huge coin records table.
+                # The performance benefit outweighs the cost of this partial
+                # index as it only includes rows where spent_index is -1.
+                await conn.execute(
+                    """
+                    CREATE INDEX IF NOT EXISTS coin_record_ph_ff_unspent_idx
+                        ON coin_record(puzzle_hash, spent_index)
+                        WHERE spent_index = -1
+                    """
+                )
+            async with conn.execute(
+                "SELECT 1 FROM sqlite_master WHERE type = 'index' AND name = 'coin_record_ph_ff_unspent_idx'"
+            ) as cursor:
+                has_ff_unspent_idx = await cursor.fetchone() is not None
+            if has_ff_unspent_idx:
+                self._unspent_lineage_for_ph_idx = "coin_record_ph_ff_unspent_idx"
+
         return self
 
     async def num_unspent(self) -> int:
         async with self.db_wrapper.reader_no_transaction() as conn:
-            async with conn.execute("SELECT COUNT(*) FROM coin_record WHERE spent_index=0") as cursor:
+            async with conn.execute("SELECT COUNT(*) FROM coin_record WHERE spent_index <= 0") as cursor:
                 row = await cursor.fetchone()
         if row is not None:
             count: int = row[0]
@@ -85,27 +108,33 @@ class CoinStore:
         height: uint32,
         timestamp: uint64,
         included_reward_coins: Collection[Coin],
-        tx_additions: Collection[Coin],
+        tx_additions: Collection[tuple[bytes32, Coin, bool]],
         tx_removals: list[bytes32],
-    ) -> list[CoinRecord]:
+    ) -> None:
         """
         Only called for blocks which are blocks (and thus have rewards and transactions)
-        Returns a list of the CoinRecords that were added by this block
         """
 
         start = time.monotonic()
 
-        additions = []
+        db_values_to_insert = []
 
-        for coin in tx_additions:
-            record: CoinRecord = CoinRecord(
-                coin,
-                height,
-                uint32(0),
-                False,
-                timestamp,
+        for coin_id, coin, same_as_parent in tx_additions:
+            db_values_to_insert.append(
+                (
+                    coin_id,
+                    # confirmed_index
+                    height,
+                    # spent_index
+                    -1 if same_as_parent else 0,
+                    # coinbase
+                    0,
+                    coin.puzzle_hash,
+                    coin.parent_coin_info,
+                    coin.amount.stream_to_bytes(),
+                    timestamp,
+                )
             )
-            additions.append(record)
 
         if height == 0:
             assert len(included_reward_coins) == 0
@@ -113,16 +142,24 @@ class CoinStore:
             assert len(included_reward_coins) >= 2
 
         for coin in included_reward_coins:
-            reward_coin_r: CoinRecord = CoinRecord(
-                coin,
-                height,
-                uint32(0),
-                True,
-                timestamp,
+            db_values_to_insert.append(
+                (
+                    coin.name(),
+                    # confirmed_index
+                    height,
+                    # spent_index
+                    0,
+                    # coinbase
+                    1,
+                    coin.puzzle_hash,
+                    coin.parent_coin_info,
+                    coin.amount.stream_to_bytes(),
+                    timestamp,
+                )
             )
-            additions.append(reward_coin_r)
 
-        await self._add_coin_records(additions)
+        async with self.db_wrapper.writer_maybe_transaction() as conn:
+            await conn.executemany("INSERT INTO coin_record VALUES(?, ?, ?, ?, ?, ?, ?, ?)", db_values_to_insert)
         await self._set_spent(tx_removals, height)
 
         end = time.monotonic()
@@ -133,8 +170,6 @@ class CoinStore:
             + "blockchain database is on a fast drive",
         )
 
-        return additions
-
     # Checks DB and DiffStores for CoinRecord with coin_name and returns it
     async def get_coin_record(self, coin_name: bytes32) -> Optional[CoinRecord]:
         async with self.db_wrapper.reader_no_transaction() as conn:
@@ -146,7 +181,8 @@ class CoinStore:
                 row = await cursor.fetchone()
                 if row is not None:
                     coin = self.row_to_coin(row)
-                    return CoinRecord(coin, row[0], row[1], row[2], row[6])
+                    spent_index = uint32(0) if row[1] <= 0 else uint32(row[1])
+                    return CoinRecord(coin, row[0], spent_index, row[2], row[6])
         return None
 
     async def get_coin_records(self, names: Collection[bytes32]) -> list[CoinRecord]:
@@ -171,16 +207,13 @@ class CoinStore:
             for cursor in cursors:
                 for row in await cursor.fetchall():
                     coin = self.row_to_coin(row)
-                    record = CoinRecord(coin, row[0], row[1], row[2], row[6])
+                    spent_index = uint32(0) if row[1] <= 0 else uint32(row[1])
+                    record = CoinRecord(coin, row[0], spent_index, row[2], row[6])
                     coins.append(record)
 
         return coins
 
     async def get_coins_added_at_height(self, height: uint32) -> list[CoinRecord]:
-        coins_added: Optional[list[CoinRecord]] = self.coins_added_at_height_cache.get(height)
-        if coins_added is not None:
-            return coins_added
-
         async with self.db_wrapper.reader_no_transaction() as conn:
             async with conn.execute(
                 "SELECT confirmed_index, spent_index, coinbase, puzzle_hash, "
@@ -191,8 +224,8 @@ class CoinStore:
                 coins = []
                 for row in rows:
                     coin = self.row_to_coin(row)
-                    coins.append(CoinRecord(coin, row[0], row[1], row[2], row[6]))
-                self.coins_added_at_height_cache.put(height, coins)
+                    spent_index = uint32(0) if row[1] <= 0 else uint32(row[1])
+                    coins.append(CoinRecord(coin, row[0], spent_index, row[2], row[6]))
                 return coins
 
     async def get_coins_removed_at_height(self, height: uint32) -> list[CoinRecord]:
@@ -207,29 +240,12 @@ class CoinStore:
             ) as cursor:
                 coins = []
                 for row in await cursor.fetchall():
-                    if row[1] != 0:
+                    if row[1] > 0:
                         coin = self.row_to_coin(row)
                         coin_record = CoinRecord(coin, row[0], row[1], row[2], row[6])
                         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,
@@ -245,12 +261,13 @@ class CoinStore:
                 f"SELECT confirmed_index, spent_index, coinbase, puzzle_hash, "
                 f"coin_parent, amount, timestamp FROM coin_record INDEXED BY coin_puzzle_hash WHERE puzzle_hash=? "
                 f"AND confirmed_index>=? AND confirmed_index<? "
-                f"{'' if include_spent_coins else 'AND spent_index=0'}",
+                f"{'' if include_spent_coins else 'AND spent_index <= 0'}",
                 (puzzle_hash, start_height, end_height),
             ) 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]))
+                    spent_index = uint32(0) if row[1] <= 0 else uint32(row[1])
+                    coins.add(CoinRecord(coin, row[0], spent_index, row[2], row[6]))
                 return list(coins)
 
     async def get_coin_records_by_puzzle_hashes(
@@ -273,12 +290,13 @@ class CoinStore:
                 f"coin_parent, amount, timestamp FROM coin_record INDEXED BY coin_puzzle_hash "
                 f"WHERE puzzle_hash in ({'?,' * (len(puzzle_hashes) - 1)}?) "
                 f"AND confirmed_index>=? AND confirmed_index<? "
-                f"{'' if include_spent_coins else 'AND spent_index=0'}",
+                f"{'' if include_spent_coins else 'AND spent_index <= 0'}",
                 (*puzzle_hashes_db, start_height, end_height),
             ) 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]))
+                    spent_index = uint32(0) if row[1] <= 0 else uint32(row[1])
+                    coins.add(CoinRecord(coin, row[0], spent_index, row[2], row[6]))
                 return list(coins)
 
     async def get_coin_records_by_names(
@@ -299,12 +317,13 @@ class CoinStore:
                 f"coin_parent, amount, timestamp FROM coin_record INDEXED BY sqlite_autoindex_coin_record_1 "
                 f"WHERE coin_name in ({'?,' * (len(names) - 1)}?) "
                 f"AND confirmed_index>=? AND confirmed_index<? "
-                f"{'' if include_spent_coins else 'AND spent_index=0'}",
+                f"{'' if include_spent_coins else 'AND spent_index <= 0'}",
                 [*names, start_height, end_height],
             ) 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]))
+                    spent_index = uint32(0) if row[1] <= 0 else uint32(row[1])
+                    coins.add(CoinRecord(coin, row[0], spent_index, row[2], row[6]))
 
         return list(coins)
 
@@ -314,7 +333,7 @@ class CoinStore:
     def row_to_coin_state(self, row: sqlite3.Row) -> CoinState:
         coin = self.row_to_coin(row)
         spent_h = None
-        if row[1] != 0:
+        if row[1] > 0:
             spent_h = row[1]
         return CoinState(coin, spent_h, row[0])
 
@@ -338,7 +357,7 @@ class CoinStore:
                     f"coin_parent, amount, timestamp FROM coin_record INDEXED BY coin_puzzle_hash "
                     f"WHERE puzzle_hash in ({'?,' * (len(batch.entries) - 1)}?) "
                     f"AND (confirmed_index>=? OR spent_index>=?)"
-                    f"{'' if include_spent_coins else 'AND spent_index=0'}"
+                    f"{'' if include_spent_coins else ' AND spent_index <= 0'}"
                     " LIMIT ?",
                     (*puzzle_hashes_db, min_height, min_height, max_items - len(coins)),
                 ) as cursor:
@@ -369,12 +388,13 @@ class CoinStore:
                     f"SELECT confirmed_index, spent_index, coinbase, puzzle_hash, coin_parent, amount, timestamp "
                     f"FROM coin_record WHERE coin_parent in ({'?,' * (len(batch.entries) - 1)}?) "
                     f"AND confirmed_index>=? AND confirmed_index<? "
-                    f"{'' if include_spent_coins else 'AND spent_index=0'}",
+                    f"{'' if include_spent_coins else 'AND spent_index <= 0'}",
                     (*parent_ids_db, start_height, end_height),
                 ) as cursor:
                     async for row in cursor:
                         coin = self.row_to_coin(row)
-                        coins.add(CoinRecord(coin, row[0], row[1], row[2], row[6]))
+                        spent_index = uint32(0) if row[1] <= 0 else uint32(row[1])
+                        coins.add(CoinRecord(coin, row[0], spent_index, row[2], row[6]))
 
         return list(coins)
 
@@ -403,7 +423,7 @@ class CoinStore:
                     f"SELECT confirmed_index, spent_index, coinbase, puzzle_hash, coin_parent, amount, timestamp "
                     f"FROM coin_record WHERE coin_name in ({'?,' * (len(batch.entries) - 1)}?) "
                     f"AND (confirmed_index>=? OR spent_index>=?) {max_height_sql}"
-                    f"{'' if include_spent_coins else 'AND spent_index=0'}"
+                    f"{'' if include_spent_coins else 'AND spent_index <= 0'}"
                     " LIMIT ?",
                     (*coin_ids_db, min_height, min_height, max_items - len(coins)),
                 ) as cursor:
@@ -414,7 +434,7 @@ class CoinStore:
 
         return coins
 
-    MAX_PUZZLE_HASH_BATCH_SIZE = SQLITE_MAX_VARIABLE_NUMBER - 10
+    MAX_PUZZLE_HASH_BATCH_SIZE: ClassVar[int] = SQLITE_MAX_VARIABLE_NUMBER - 10
 
     async def batch_coin_states_by_puzzle_hashes(
         self,
@@ -448,7 +468,7 @@ class CoinStore:
             puzzle_hash_count = len(puzzle_hashes_db)
 
             require_spent = "spent_index>0"
-            require_unspent = "spent_index=0"
+            require_unspent = "spent_index <= 0"
             amount_filter = "AND amount>=? " if min_amount > 0 else ""
 
             if include_spent and include_unspent:
@@ -531,66 +551,70 @@ class CoinStore:
 
         return coin_states, next_height
 
-    async def rollback_to_block(self, block_index: int) -> list[CoinRecord]:
+    async def rollback_to_block(self, block_index: int) -> dict[bytes32, CoinRecord]:
         """
         Note that block_index can be negative, in which case everything is rolled back
-        Returns the list of coin records that have been modified
+        Returns a map of coin ID to coin record for modified items.
         """
 
         coin_changes: dict[bytes32, CoinRecord] = {}
         # Add coins that are confirmed in the reverted blocks to the list of updated coins.
         async with self.db_wrapper.writer_maybe_transaction() as conn:
-            async with conn.execute(
+            rows = await conn.execute_fetchall(
                 "SELECT confirmed_index, spent_index, coinbase, puzzle_hash, "
-                "coin_parent, amount, timestamp FROM coin_record WHERE confirmed_index>?",
+                "coin_parent, amount, timestamp, coin_name FROM coin_record WHERE confirmed_index>?",
                 (block_index,),
-            ) as cursor:
-                for row in await cursor.fetchall():
-                    coin = self.row_to_coin(row)
-                    record = CoinRecord(coin, uint32(0), row[1], row[2], uint64(0))
-                    coin_changes[record.name] = record
+            )
+            for row in rows:
+                coin = self.row_to_coin(row)
+                spent_index = uint32(0) if row[1] <= 0 else uint32(row[1])
+                record = CoinRecord(coin, uint32(0), spent_index, row[2], uint64(0))
+                coin_name = bytes32(row[7])
+                coin_changes[coin_name] = record
 
             # Delete reverted blocks from storage
             await conn.execute("DELETE FROM coin_record WHERE confirmed_index>?", (block_index,))
 
             # Add coins that are confirmed in the reverted blocks to the list of changed coins.
-            async with conn.execute(
+            rows = await conn.execute_fetchall(
                 "SELECT confirmed_index, spent_index, coinbase, puzzle_hash, "
-                "coin_parent, amount, timestamp FROM coin_record WHERE spent_index>?",
+                "coin_parent, amount, timestamp, coin_name FROM coin_record WHERE spent_index>?",
                 (block_index,),
-            ) as cursor:
-                for row in await cursor.fetchall():
-                    coin = self.row_to_coin(row)
-                    record = CoinRecord(coin, row[0], uint32(0), row[2], row[6])
-                    if record.name not in coin_changes:
-                        coin_changes[record.name] = record
-
-            await conn.execute("UPDATE coin_record SET spent_index=0 WHERE spent_index>?", (block_index,))
-        self.coins_added_at_height_cache = LRUCache(self.coins_added_at_height_cache.capacity)
-        return list(coin_changes.values())
-
-    # Store CoinRecord in DB
-    async def _add_coin_records(self, records: list[CoinRecord]) -> None:
-        values2 = []
-        for record in records:
-            values2.append(
-                (
-                    record.coin.name(),
-                    record.confirmed_block_index,
-                    record.spent_block_index,
-                    int(record.coinbase),
-                    record.coin.puzzle_hash,
-                    record.coin.parent_coin_info,
-                    uint64(record.coin.amount).stream_to_bytes(),
-                    record.timestamp,
-                )
             )
-        if len(values2) > 0:
-            async with self.db_wrapper.writer_maybe_transaction() as conn:
-                await conn.executemany(
-                    "INSERT INTO coin_record VALUES(?, ?, ?, ?, ?, ?, ?, ?)",
-                    values2,
-                )
+            for row in rows:
+                coin = self.row_to_coin(row)
+                record = CoinRecord(coin, row[0], uint32(0), row[2], row[6])
+                coin_name = bytes32(row[7])
+                if coin_name not in coin_changes:
+                    coin_changes[coin_name] = record
+
+            # If the coin to update is not a reward coin and its parent is
+            # spent and has the same puzzle hash and amount, we set its
+            # spent_index to -1 as a potential fast forward singleton unspent
+            # otherwise we set it to 0 as a normal unspent.
+            await conn.execute(
+                """
+                UPDATE coin_record INDEXED BY coin_spent_index
+                SET spent_index = CASE
+                    WHEN
+                        coinbase = 0 AND
+                        EXISTS (
+                            SELECT 1
+                            FROM coin_record AS parent INDEXED BY sqlite_autoindex_coin_record_1
+                            WHERE
+                                parent.coin_name = coin_record.coin_parent AND
+                                parent.puzzle_hash = coin_record.puzzle_hash AND
+                                parent.amount = coin_record.amount AND
+                                parent.spent_index > 0
+                        )
+                    THEN -1
+                    ELSE 0
+                END
+                WHERE spent_index > ?
+                """,
+                (block_index,),
+            )
+        return coin_changes
 
     # Update coin_record to be spent in DB
     async def _set_spent(self, coin_names: list[bytes32], index: uint32) -> None:
@@ -606,7 +630,7 @@ class CoinStore:
                 ret: Cursor = await conn.execute(
                     f"UPDATE coin_record INDEXED BY sqlite_autoindex_coin_record_1 "
                     f"SET spent_index={index} "
-                    f"WHERE spent_index=0 "
+                    f"WHERE spent_index <= 0 "
                     f"AND coin_name IN ({name_params})",
                     batch.entries,
                 )
@@ -621,27 +645,32 @@ class CoinStore:
         async with self.db_wrapper.reader_no_transaction() as conn:
             async with conn.execute(
                 "SELECT unspent.coin_name, "
-                "unspent.amount, "
                 "unspent.coin_parent, "
-                "parent.amount, "
                 "parent.coin_parent "
-                "FROM coin_record AS unspent INDEXED BY coin_puzzle_hash "
+                "FROM coin_record AS unspent "
+                f"INDEXED BY {self._unspent_lineage_for_ph_idx} "
                 "LEFT JOIN coin_record AS parent ON unspent.coin_parent = parent.coin_name "
-                "WHERE unspent.spent_index = 0 "
+                "WHERE unspent.spent_index = -1 "
                 "AND parent.spent_index > 0 "
                 "AND unspent.puzzle_hash = ? "
-                "AND parent.puzzle_hash = unspent.puzzle_hash",
+                "AND parent.puzzle_hash = unspent.puzzle_hash "
+                "AND parent.amount = unspent.amount",
                 (puzzle_hash,),
             ) as cursor:
                 rows = list(await cursor.fetchall())
                 if len(rows) != 1:
                     log.debug("Expected 1 unspent with puzzle hash %s, but found %s", puzzle_hash.hex(), len(rows))
                     return None
-                coin_id, coin_amount, parent_id, parent_amount, parent_parent_id = rows[0]
+                coin_id, parent_id, parent_parent_id = rows[0]
                 return UnspentLineageInfo(
-                    coin_id=bytes32(coin_id),
-                    coin_amount=uint64(int_from_bytes(coin_amount)),
-                    parent_id=bytes32(parent_id),
-                    parent_amount=uint64(int_from_bytes(parent_amount)),
-                    parent_parent_id=bytes32(parent_parent_id),
+                    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