fixup

chia/_tests/core/data_layer/test_merkle_blob.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import hashlib
 import itertools
+import re
 from dataclasses import dataclass
 from random import Random
 from typing import Generic, Protocol, TypeVar, final
@@ -30,6 +31,7 @@ from chia.data_layer.util.merkle_blob import (
     raw_node_classes,
     raw_node_type_to_class,
     spacing,
+    undefined_index,
     unpack_raw_node,
 )
 from chia.types.blockchain_format.sized_bytes import bytes32
@@ -360,7 +362,7 @@ def test_proof_of_inclusion_merkle_blob() -> None:
             del keys_values[kv_id]
 
         for kv_id in delete_ordering:
-            with pytest.raises(Exception, match=f"Key {kv_id} not present in the store"):
+            with pytest.raises(Exception, match=f"Key {re.escape(str(kv_id))} not present in the store"):
                 merkle_blob.get_proof_of_inclusion(kv_id)
 
         new_keys_values: dict[KeyId, ValueId] = {}
@@ -379,7 +381,7 @@ def test_proof_of_inclusion_merkle_blob() -> None:
             assert proof_of_inclusion.valid()
 
 
-@pytest.mark.parametrize(argnames="index", argvalues=[-1, 1, None])
+@pytest.mark.parametrize(argnames="index", argvalues=[TreeIndex(uint32(1)), undefined_index])
 def test_get_raw_node_raises_for_invalid_indexes(index: TreeIndex) -> None:
     merkle_blob = MerkleBlob(blob=bytearray())
     merkle_blob.insert(

chia/data_layer/data_store.py

@@ -51,6 +51,7 @@ from chia.data_layer.util.merkle_blob import (
 from chia.types.blockchain_format.sized_bytes import bytes32
 from chia.util.batches import to_batches
 from chia.util.db_wrapper import SQLITE_MAX_VARIABLE_NUMBER, DBWrapper2
+from chia.util.ints import uint32
 from chia.util.lru_cache import LRUCache
 
 log = logging.getLogger(__name__)
@@ -227,7 +228,7 @@ class DataStore:
                     terminal_nodes[node_hash] = (kid, vid)
 
         missing_hashes: list[bytes32] = []
-        merkle_blob_queries: dict[bytes32, list[int]] = defaultdict(list)
+        merkle_blob_queries: dict[bytes32, list[TreeIndex]] = defaultdict(list)
 
         for _, (left, right) in internal_nodes.items():
             for node_hash in (left, right):
@@ -252,7 +253,7 @@ class DataStore:
                     for row in rows:
                         node_hash = bytes32(row["hash"])
                         root_hash_blob = bytes32(row["root_hash"])
-                        index = row["idx"]
+                        index = TreeIndex(uint32(row["idx"]))
                         if node_hash in found_hashes:
                             raise Exception("Internal error: duplicate node_hash found in nodes table")
                         merkle_blob_queries[root_hash_blob].append(index)
@@ -477,8 +478,8 @@ class DataStore:
                 "INSERT OR REPLACE INTO hashes (hash, kid, vid, store_id) VALUES (?, ?, ?, ?)",
                 (
                     hash,
-                    kid,
-                    vid,
+                    kid.raw,
+                    vid.raw,
                     store_id,
                 ),
             )