mypy

Author: richardkiss

chia/consensus/block_store_protocol.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from contextlib import AbstractAsyncContextManager
 from typing import Optional, Protocol
 
 from chia_rs import BlockRecord, FullBlock, SubEpochChallengeSegment
@@ -11,62 +12,37 @@
 
 class BlockStoreProtocol(Protocol):
     async def add_full_block(self, header_hash: bytes32, block: FullBlock, block_record: BlockRecord) -> None: ...
-
+    def get_block_from_cache(self, header_hash: bytes32) -> Optional[FullBlock]: ...
     async def get_full_block(self, header_hash: bytes32) -> Optional[FullBlock]: ...
-
     async def get_full_block_bytes(self, header_hash: bytes32) -> Optional[bytes]: ...
-
     async def get_full_blocks_at(self, heights: list[uint32]) -> list[FullBlock]: ...
-
     async def get_block_info(self, header_hash: bytes32) -> Optional[GeneratorBlockInfo]: ...
-
     async def get_generator(self, header_hash: bytes32) -> Optional[bytes]: ...
-
     async def get_generators_at(self, heights: set[uint32]) -> dict[uint32, bytes]: ...
-
     async def get_block_record(self, header_hash: bytes32) -> Optional[BlockRecord]: ...
-
-    async def get_block_records_in_range(
-        self,
-        start: int,
-        stop: int,
-    ) -> dict[bytes32, BlockRecord]: ...
-
+    async def get_block_records_in_range(self, start: int, stop: int) -> dict[bytes32, BlockRecord]: ...
     async def get_block_records_by_hash(self, header_hashes: list[bytes32]) -> list[BlockRecord]: ...
-
     async def get_block_bytes_by_hash(self, header_hashes: list[bytes32]) -> list[bytes]: ...
-
     async def get_blocks_by_hash(self, header_hashes: list[bytes32]) -> list[FullBlock]: ...
-
     async def get_peak(self) -> Optional[tuple[bytes32, uint32]]: ...
-
-    async def get_block_bytes_in_range(
-        self,
-        start: int,
-        stop: int,
-    ) -> list[bytes]: ...
-
+    async def get_block_bytes_in_range(self, start: int, stop: int) -> list[bytes]: ...
     async def get_random_not_compactified(self, number: int) -> list[int]: ...
-
     async def persist_sub_epoch_challenge_segments(
         self, ses_block_hash: bytes32, segments: list[SubEpochChallengeSegment]
     ) -> None: ...
-
     async def get_sub_epoch_challenge_segments(
         self,
         ses_block_hash: bytes32,
     ) -> Optional[list[SubEpochChallengeSegment]]: ...
-
     async def rollback(self, height: int) -> None: ...
-
+    def rollback_cache_block(self, header_hash: bytes32) -> None: ...
     async def set_in_chain(self, header_hashes: list[tuple[bytes32]]) -> None: ...
-
     async def set_peak(self, header_hash: bytes32) -> None: ...
-
     async def is_fully_compactified(self, header_hash: bytes32) -> Optional[bool]: ...
-
     async def replace_proof(self, header_hash: bytes32, block: FullBlock) -> None: ...
-
     async def count_compactified_blocks(self) -> int: ...
-
     async def count_uncompactified_blocks(self) -> int: ...
+    async def get_block_records_close_to_peak(self, blocks_n: int) -> tuple[dict[bytes32, BlockRecord], Optional[bytes32]]: ...
+    def get_host_parameter_limit(self) -> int: ...
+    async def get_prev_hash(self, header_hash: bytes32) -> bytes32: ...
+    def transaction(self) -> AbstractAsyncContextManager[None]: ...

chia/consensus/blockchain.py

@@ -421,7 +421,7 @@ async def add_block(
 
         try:
             # Always add the block to the database
-            async with self.block_store.db_wrapper.writer():
+            async with self.block_store.transaction():
                 # Perform the DB operations to update the state, and rollback if something goes wrong
                 await self.block_store.add_full_block(header_hash, block, block_record)
                 records, state_change_summary = await self._reconsider_peak(block_record, genesis, fork_info)
@@ -883,7 +883,7 @@ async def get_header_blocks_in_range(
 
         blocks: list[FullBlock] = []
         for hash in hashes.copy():
-            block = self.block_store.block_cache.get(hash)
+            block = self.block_store.get_block_from_cache(hash)
             if block is not None:
                 blocks.append(block)
                 hashes.remove(hash)
@@ -932,7 +932,7 @@ async def get_block_records_at(self, heights: list[uint32], batch_size: int = 90
         """
         records: list[BlockRecord] = []
         hashes: list[bytes32] = []
-        assert batch_size < self.block_store.db_wrapper.host_parameter_limit
+        assert batch_size < self.block_store.get_host_parameter_limit()
         for height in heights:
             header_hash: Optional[bytes32] = self.height_to_hash(height)
             if header_hash is None:

chia/full_node/block_store.py

@@ -189,6 +189,39 @@ async def get_sub_epoch_challenge_segments(
             return challenge_segments
         return None
 
+    def get_host_parameter_limit(self) -> int:
+        return self.db_wrapper.host_parameter_limit
+
+    def transaction(self) -> AbstractAsyncContextManager[None]:
+        return self.db_wrapper.writer()
+
+    def get_block_from_cache(self, header_hash: bytes32) -> Optional[FullBlock]:
+        return self.block_cache.get(header_hash)
+
+    async def get_block_records_close_to_peak(
+        self, blocks_n: int
+    ) -> tuple[dict[bytes32, BlockRecord], Optional[bytes32]]:
+        """
+        Returns a dictionary with all blocks that have height >= peak height - blocks_n, as well as the
+        peak header hash. Only blocks that are part of the main chain/current peak are included.
+        """
+
+        peak = await self.get_peak()
+        if peak is None:
+            return {}, None
+
+        ret: dict[bytes32, BlockRecord] = {}
+        async with self.db_wrapper.reader_no_transaction() as conn:
+            async with conn.execute(
+                "SELECT header_hash, block_record FROM full_blocks WHERE height >= ? AND in_main_chain=1",
+                (peak[1] - blocks_n,),
+            ) as cursor:
+                for row in await cursor.fetchall():
+                    header_hash = bytes32(row[0])
+                    ret[header_hash] = BlockRecord.from_bytes(row[1])
+
+        return ret, peak[0]
+
     def rollback_cache_block(self, header_hash: bytes32) -> None:
         try:
             self.block_cache.remove(header_hash)
@@ -197,6 +230,25 @@ def rollback_cache_block(self, header_hash: bytes32) -> None:
             # block to the cache yet
             pass
 
+    async def get_prev_hash(self, header_hash: bytes32) -> bytes32:
+        """
+        Returns the header hash preceeding the input header hash.
+        Throws an exception if the block is not present
+        """
+        cached = self.block_cache.get(header_hash)
+        if cached is not None:
+            return cached.prev_header_hash
+
+        async with self.db_wrapper.reader_no_transaction() as conn:
+            async with conn.execute(
+                "SELECT prev_hash FROM full_blocks WHERE header_hash=?",
+                (header_hash,),
+            ) as cursor:
+                row = await cursor.fetchone()
+                if row is None:
+                    raise KeyError("missing block in chain")
+                return bytes32(row[0])
+
     async def get_full_block(self, header_hash: bytes32) -> Optional[FullBlock]:
         cached: Optional[FullBlock] = self.block_cache.get(header_hash)
         if cached is not None: