chia-blockchain 2.5.6rc2__py3-none-any.whl → 2.5.7__py3-none-any.whl

chia/_tests/core/data_layer/test_data_store.py

@@ -1,54 +1,54 @@
 from __future__ import annotations
 
+import contextlib
+import importlib.resources as importlib_resources
 import itertools
 import logging
 import os
 import random
 import re
+import shutil
 import statistics
 import time
 from collections.abc import Awaitable
 from dataclasses import dataclass
+from hashlib import sha256
 from pathlib import Path
 from random import Random
-from typing import Any, Callable, Optional, cast
+from typing import Any, BinaryIO, Callable, Optional
 
 import aiohttp
-import aiosqlite
+import chia_rs.datalayer
 import pytest
+from chia_rs.datalayer import KeyAlreadyPresentError, MerkleBlob, TreeIndex
 from chia_rs.sized_bytes import bytes32
 
 from chia._tests.core.data_layer.util import Example, add_0123_example, add_01234567_example
 from chia._tests.util.misc import BenchmarkRunner, Marks, boolean_datacases, datacases
-from chia.data_layer.data_layer_errors import KeyNotFoundError, NodeHashError, TreeGenerationIncrementingError
+from chia.data_layer.data_layer_errors import KeyNotFoundError, TreeGenerationIncrementingError
 from chia.data_layer.data_layer_util import (
     DiffData,
     InternalNode,
-    Node,
-    NodeType,
     OperationType,
-    ProofOfInclusion,
-    ProofOfInclusionLayer,
     Root,
+    SerializedNode,
     ServerInfo,
     Side,
     Status,
     Subscription,
     TerminalNode,
     _debug_dump,
-    leaf_hash,
-)
-from chia.data_layer.data_store import DataStore
-from chia.data_layer.download_data import (
     get_delta_filename_path,
     get_full_tree_filename_path,
-    insert_from_delta_file,
-    insert_into_data_store_from_file,
-    write_files_for_root,
+    leaf_hash,
 )
+from chia.data_layer.data_store import DataStore
+from chia.data_layer.download_data import insert_from_delta_file, write_files_for_root
+from chia.data_layer.util.benchmark import generate_datastore
 from chia.types.blockchain_format.program import Program
 from chia.util.byte_types import hexstr_to_bytes
 from chia.util.db_wrapper import DBWrapper2, generate_in_memory_db_uri
+from chia.util.lru_cache import LRUCache
 
 log = logging.getLogger(__name__)
 
@@ -57,31 +57,48 @@ pytestmark = pytest.mark.data_layer
 
 
 table_columns: dict[str, list[str]] = {
-    "node": ["hash", "node_type", "left", "right", "key", "value"],
     "root": ["tree_id", "generation", "node_hash", "status"],
+    "subscriptions": ["tree_id", "url", "ignore_till", "num_consecutive_failures", "from_wallet"],
+    "schema": ["version_id", "applied_at"],
+    "ids": ["kv_id", "hash", "blob", "store_id"],
+    "nodes": ["store_id", "hash", "root_hash", "generation", "idx"],
 }
 
 
-# TODO: Someday add tests for malformed DB data to make sure we handle it gracefully
-#       and with good error messages.
-
-
 @pytest.mark.anyio
-async def test_valid_node_values_fixture_are_valid(data_store: DataStore, valid_node_values: dict[str, Any]) -> None:
-    async with data_store.db_wrapper.writer() as writer:
-        await writer.execute(
-            """
-            INSERT INTO node(hash, node_type, left, right, key, value)
-            VALUES(:hash, :node_type, :left, :right, :key, :value)
-            """,
-            valid_node_values,
-        )
+async def test_migrate_from_old_format(store_id: bytes32, tmp_path: Path) -> None:
+    old_format_resources = importlib_resources.files(__name__.rpartition(".")[0]).joinpath("old_format")
+    db_uri = tmp_path / "old_db.sqlite"
+    db_uri.write_bytes(old_format_resources.joinpath("old_db.sqlite").read_bytes())
+    files_resources = old_format_resources.joinpath("files")
+
+    with importlib_resources.as_file(files_resources) as files_path:
+        async with DataStore.managed(
+            database=db_uri,
+            uri=True,
+            merkle_blobs_path=tmp_path.joinpath("merkle-blobs"),
+            key_value_blobs_path=tmp_path.joinpath("key-value-blobs"),
+        ) as data_store:
+            await data_store.migrate_db(files_path)
+            root = await data_store.get_tree_root(store_id=store_id)
+            expected = Root(
+                store_id=bytes32.fromhex("2e2e2e2e2e2e2e2e2e2e2e2e2e2e2e2e2e2e2e2e2e2e612073746f7265206964"),
+                node_hash=bytes32.fromhex("aa77376d1ccd3664e5c6366e010c52a978fedbf40f5ce262fee71b2e7fe0c6a9"),
+                generation=50,
+                status=Status.COMMITTED,
+            )
+            assert root == expected
 
 
+# TODO: Someday add tests for malformed DB data to make sure we handle it gracefully
+#       and with good error messages.
 @pytest.mark.parametrize(argnames=["table_name", "expected_columns"], argvalues=table_columns.items())
 @pytest.mark.anyio
 async def test_create_creates_tables_and_columns(
-    database_uri: str, table_name: str, expected_columns: list[str]
+    database_uri: str,
+    table_name: str,
+    expected_columns: list[str],
+    tmp_path: Path,
 ) -> None:
     # Never string-interpolate sql queries...  Except maybe in tests when it does not
     # allow you to parametrize the query.
@@ -93,7 +110,12 @@ async def test_create_creates_tables_and_columns(
             columns = await cursor.fetchall()
             assert columns == []
 
-        async with DataStore.managed(database=database_uri, uri=True):
+        async with DataStore.managed(
+            database=database_uri,
+            uri=True,
+            merkle_blobs_path=tmp_path.joinpath("merkle-blobs"),
+            key_value_blobs_path=tmp_path.joinpath("key-value-blobs"),
+        ):
             async with db_wrapper.reader() as reader:
                 cursor = await reader.execute(query)
                 columns = await cursor.fetchall()
@@ -211,46 +233,6 @@ async def test_get_tree_generation_returns_none_when_none_available(
         await raw_data_store.get_tree_generation(store_id=store_id)
 
 
-@pytest.mark.anyio
-async def test_insert_internal_node_does_nothing_if_matching(data_store: DataStore, store_id: bytes32) -> None:
-    await add_01234567_example(data_store=data_store, store_id=store_id)
-
-    kv_node = await data_store.get_node_by_key(key=b"\x04", store_id=store_id)
-    ancestors = await data_store.get_ancestors(node_hash=kv_node.hash, store_id=store_id)
-    parent = ancestors[0]
-
-    async with data_store.db_wrapper.reader() as reader:
-        cursor = await reader.execute("SELECT * FROM node")
-        before = await cursor.fetchall()
-
-    await data_store._insert_internal_node(left_hash=parent.left_hash, right_hash=parent.right_hash)
-
-    async with data_store.db_wrapper.reader() as reader:
-        cursor = await reader.execute("SELECT * FROM node")
-        after = await cursor.fetchall()
-
-    assert after == before
-
-
-@pytest.mark.anyio
-async def test_insert_terminal_node_does_nothing_if_matching(data_store: DataStore, store_id: bytes32) -> None:
-    await add_01234567_example(data_store=data_store, store_id=store_id)
-
-    kv_node = await data_store.get_node_by_key(key=b"\x04", store_id=store_id)
-
-    async with data_store.db_wrapper.reader() as reader:
-        cursor = await reader.execute("SELECT * FROM node")
-        before = await cursor.fetchall()
-
-    await data_store._insert_terminal_node(key=kv_node.key, value=kv_node.value)
-
-    async with data_store.db_wrapper.reader() as reader:
-        cursor = await reader.execute("SELECT * FROM node")
-        after = await cursor.fetchall()
-
-    assert after == before
-
-
 @pytest.mark.anyio
 async def test_build_a_tree(
     data_store: DataStore,
@@ -293,7 +275,7 @@ async def test_get_ancestors(data_store: DataStore, store_id: bytes32) -> None:
         "c852ecd8fb61549a0a42f9eb9dde65e6c94a01934dbd9c1d35ab94e2a0ae58e2",
     ]
 
-    ancestors_2 = await data_store.get_ancestors_optimized(node_hash=reference_node_hash, store_id=store_id)
+    ancestors_2 = await data_store.get_ancestors(node_hash=reference_node_hash, store_id=store_id)
     assert ancestors == ancestors_2
 
 
@@ -306,6 +288,10 @@ async def test_get_ancestors_optimized(data_store: DataStore, store_id: bytes32)
     first_insertions = [True, False, True, False, True, True, False, True, False, True, True, False, False, True, False]
     deleted_all = False
     node_count = 0
+    node_hashes: list[bytes32] = []
+    hash_to_key: dict[bytes32, bytes] = {}
+    node_hash: Optional[bytes32]
+
     for i in range(1000):
         is_insert = False
         if i <= 14:
@@ -318,12 +304,10 @@ async def test_get_ancestors_optimized(data_store: DataStore, store_id: bytes32)
             if not deleted_all:
                 while node_count > 0:
                     node_count -= 1
-                    seed = bytes32(b"0" * 32)
-                    node_hash = await data_store.get_terminal_node_for_seed(store_id, seed)
+                    node_hash = random.choice(node_hashes)
                     assert node_hash is not None
-                    node = await data_store.get_node(node_hash)
-                    assert isinstance(node, TerminalNode)
-                    await data_store.delete(key=node.key, store_id=store_id, status=Status.COMMITTED)
+                    await data_store.delete(key=hash_to_key[node_hash], store_id=store_id, status=Status.COMMITTED)
+                    node_hashes.remove(node_hash)
                 deleted_all = True
                 is_insert = True
             else:
@@ -335,10 +319,10 @@ async def test_get_ancestors_optimized(data_store: DataStore, store_id: bytes32)
         key = (i % 200).to_bytes(4, byteorder="big")
         value = (i % 200).to_bytes(4, byteorder="big")
         seed = Program.to((key, value)).get_tree_hash()
-        node_hash = await data_store.get_terminal_node_for_seed(store_id, seed)
+        node_hash = None if len(node_hashes) == 0 else random.choice(node_hashes)
         if is_insert:
             node_count += 1
-            side = None if node_hash is None else data_store.get_side_for_seed(seed)
+            side = None if node_hash is None else (Side.LEFT if seed[0] < 128 else Side.RIGHT)
 
             insert_result = await data_store.insert(
                 key=key,
@@ -346,10 +330,11 @@ async def test_get_ancestors_optimized(data_store: DataStore, store_id: bytes32)
                 store_id=store_id,
                 reference_node_hash=node_hash,
                 side=side,
-                use_optimized=False,
                 status=Status.COMMITTED,
             )
             node_hash = insert_result.node_hash
+            hash_to_key[node_hash] = key
+            node_hashes.append(node_hash)
             if node_hash is not None:
                 generation = await data_store.get_tree_generation(store_id=store_id)
                 current_ancestors = await data_store.get_ancestors(node_hash=node_hash, store_id=store_id)
@@ -357,39 +342,38 @@ async def test_get_ancestors_optimized(data_store: DataStore, store_id: bytes32)
         else:
             node_count -= 1
             assert node_hash is not None
-            node = await data_store.get_node(node_hash)
-            assert isinstance(node, TerminalNode)
-            await data_store.delete(key=node.key, store_id=store_id, use_optimized=False, status=Status.COMMITTED)
+            node_hashes.remove(node_hash)
+            await data_store.delete(key=hash_to_key[node_hash], store_id=store_id, status=Status.COMMITTED)
 
     for generation, node_hash, expected_ancestors in ancestors:
-        current_ancestors = await data_store.get_ancestors_optimized(
+        current_ancestors = await data_store.get_ancestors(
             node_hash=node_hash, store_id=store_id, generation=generation
         )
         assert current_ancestors == expected_ancestors
 
 
 @pytest.mark.anyio
-@pytest.mark.parametrize(
-    "use_optimized",
-    [True, False],
-)
 @pytest.mark.parametrize(
     "num_batches",
     [1, 5, 10, 25],
 )
-async def test_batch_update(
+async def test_batch_update_against_single_operations(
     data_store: DataStore,
     store_id: bytes32,
-    use_optimized: bool,
     tmp_path: Path,
     num_batches: int,
 ) -> None:
-    total_operations = 1000 if use_optimized else 100
+    total_operations = 1000
     num_ops_per_batch = total_operations // num_batches
     saved_batches: list[list[dict[str, Any]]] = []
     saved_kv: list[list[TerminalNode]] = []
     db_uri = generate_in_memory_db_uri()
-    async with DataStore.managed(database=db_uri, uri=True) as single_op_data_store:
+    async with DataStore.managed(
+        database=db_uri,
+        uri=True,
+        merkle_blobs_path=tmp_path.joinpath("merkle-blobs"),
+        key_value_blobs_path=tmp_path.joinpath("key-value-blobs"),
+    ) as single_op_data_store:
         await single_op_data_store.create_tree(store_id, status=Status.COMMITTED)
         random = Random()
         random.seed(100, version=2)
@@ -412,7 +396,6 @@ async def test_batch_update(
                         key=key,
                         value=value,
                         store_id=store_id,
-                        use_optimized=use_optimized,
                         status=Status.COMMITTED,
                     )
                 else:
@@ -420,7 +403,6 @@ async def test_batch_update(
                         key=key,
                         new_value=value,
                         store_id=store_id,
-                        use_optimized=use_optimized,
                         status=Status.COMMITTED,
                     )
                 action = "insert" if op_type == "insert" else "upsert"
@@ -432,7 +414,6 @@ async def test_batch_update(
                 await single_op_data_store.delete(
                     key=key,
                     store_id=store_id,
-                    use_optimized=use_optimized,
                     status=Status.COMMITTED,
                 )
                 batch.append({"action": "delete", "key": key})
@@ -446,7 +427,6 @@ async def test_batch_update(
                     key=key,
                     new_value=new_value,
                     store_id=store_id,
-                    use_optimized=use_optimized,
                     status=Status.COMMITTED,
                 )
                 keys_values[key] = new_value
@@ -469,38 +449,13 @@ async def test_batch_update(
         assert {node.key: node.value for node in current_kv} == {
             node.key: node.value for node in saved_kv[batch_number]
         }
-        queue: list[bytes32] = [root.node_hash]
-        ancestors: dict[bytes32, bytes32] = {}
-        while len(queue) > 0:
-            node_hash = queue.pop(0)
-            expected_ancestors = []
-            ancestor = node_hash
-            while ancestor in ancestors:
-                ancestor = ancestors[ancestor]
-                expected_ancestors.append(ancestor)
-            result_ancestors = await data_store.get_ancestors_optimized(node_hash, store_id)
-            assert [node.hash for node in result_ancestors] == expected_ancestors
-            node = await data_store.get_node(node_hash)
-            if isinstance(node, InternalNode):
-                queue.append(node.left_hash)
-                queue.append(node.right_hash)
-                ancestors[node.left_hash] = node_hash
-                ancestors[node.right_hash] = node_hash
 
     all_kv = await data_store.get_keys_values(store_id)
     assert {node.key: node.value for node in all_kv} == keys_values
 
 
 @pytest.mark.anyio
-@pytest.mark.parametrize(
-    "use_optimized",
-    [True, False],
-)
-async def test_upsert_ignores_existing_arguments(
-    data_store: DataStore,
-    store_id: bytes32,
-    use_optimized: bool,
-) -> None:
+async def test_upsert_ignores_existing_arguments(data_store: DataStore, store_id: bytes32) -> None:
     key = b"key"
     value = b"value1"
 
@@ -508,7 +463,6 @@ async def test_upsert_ignores_existing_arguments(
         key=key,
         value=value,
         store_id=store_id,
-        use_optimized=use_optimized,
         status=Status.COMMITTED,
     )
     node = await data_store.get_node_by_key(key, store_id)
@@ -519,7 +473,6 @@ async def test_upsert_ignores_existing_arguments(
         key=key,
         new_value=new_value,
         store_id=store_id,
-        use_optimized=use_optimized,
         status=Status.COMMITTED,
     )
     node = await data_store.get_node_by_key(key, store_id)
@@ -529,7 +482,6 @@ async def test_upsert_ignores_existing_arguments(
         key=key,
         new_value=new_value,
         store_id=store_id,
-        use_optimized=use_optimized,
         status=Status.COMMITTED,
     )
     node = await data_store.get_node_by_key(key, store_id)
@@ -540,7 +492,6 @@ async def test_upsert_ignores_existing_arguments(
         key=key2,
         new_value=value,
         store_id=store_id,
-        use_optimized=use_optimized,
         status=Status.COMMITTED,
     )
     node = await data_store.get_node_by_key(key2, store_id)
@@ -575,30 +526,24 @@ async def test_insert_batch_reference_and_side(
     )
     assert new_root_hash is not None, "batch insert failed or failed to update root"
 
-    parent = await data_store.get_node(new_root_hash)
-    assert isinstance(parent, InternalNode)
+    merkle_blob = await data_store.get_merkle_blob(store_id=store_id, root_hash=new_root_hash)
+    nodes_with_indexes = merkle_blob.get_nodes_with_indexes()
+    nodes = [pair[1] for pair in nodes_with_indexes]
+    assert len(nodes) == 3
+    assert isinstance(nodes[1], chia_rs.datalayer.LeafNode)
+    assert isinstance(nodes[2], chia_rs.datalayer.LeafNode)
+    left_terminal_node = await data_store.get_terminal_node(nodes[1].key, nodes[1].value, store_id)
+    right_terminal_node = await data_store.get_terminal_node(nodes[2].key, nodes[2].value, store_id)
     if side == Side.LEFT:
-        child = await data_store.get_node(parent.left_hash)
-        assert parent.left_hash == child.hash
+        assert left_terminal_node.key == b"key2"
+        assert right_terminal_node.key == b"key1"
     elif side == Side.RIGHT:
-        child = await data_store.get_node(parent.right_hash)
-        assert parent.right_hash == child.hash
+        assert left_terminal_node.key == b"key1"
+        assert right_terminal_node.key == b"key2"
     else:  # pragma: no cover
         raise Exception("invalid side for test")
 
 
-@pytest.mark.anyio
-async def test_ancestor_table_unique_inserts(data_store: DataStore, store_id: bytes32) -> None:
-    await add_0123_example(data_store=data_store, store_id=store_id)
-    hash_1 = bytes32.from_hexstr("0763561814685fbf92f6ca71fbb1cb11821951450d996375c239979bd63e9535")
-    hash_2 = bytes32.from_hexstr("924be8ff27e84cba17f5bc918097f8410fab9824713a4668a21c8e060a8cab40")
-    await data_store._insert_ancestor_table(hash_1, hash_2, store_id, 2)
-    await data_store._insert_ancestor_table(hash_1, hash_2, store_id, 2)
-    with pytest.raises(Exception, match="^Requested insertion of ancestor"):
-        await data_store._insert_ancestor_table(hash_1, hash_1, store_id, 2)
-    await data_store._insert_ancestor_table(hash_1, hash_2, store_id, 2)
-
-
 @pytest.mark.anyio
 async def test_get_pairs(
     data_store: DataStore,
@@ -609,7 +554,7 @@ async def test_get_pairs(
 
     pairs = await data_store.get_keys_values(store_id=store_id)
 
-    assert [node.hash for node in pairs] == example.terminal_nodes
+    assert {node.hash for node in pairs} == set(example.terminal_nodes)
 
 
 @pytest.mark.anyio
@@ -662,37 +607,6 @@ async def test_inserting_duplicate_key_fails(
         )
 
 
-@pytest.mark.anyio()
-async def test_inserting_invalid_length_hash_raises_original_exception(
-    data_store: DataStore,
-) -> None:
-    with pytest.raises(aiosqlite.IntegrityError):
-        # casting since we are testing an invalid case
-        await data_store._insert_node(
-            node_hash=cast(bytes32, b"\x05"),
-            node_type=NodeType.TERMINAL,
-            left_hash=None,
-            right_hash=None,
-            key=b"\x06",
-            value=b"\x07",
-        )
-
-
-@pytest.mark.anyio()
-async def test_inserting_invalid_length_ancestor_hash_raises_original_exception(
-    data_store: DataStore,
-    store_id: bytes32,
-) -> None:
-    with pytest.raises(aiosqlite.IntegrityError):
-        # casting since we are testing an invalid case
-        await data_store._insert_ancestor_table(
-            left_hash=bytes32(b"\x01" * 32),
-            right_hash=bytes32(b"\x02" * 32),
-            store_id=store_id,
-            generation=0,
-        )
-
-
 @pytest.mark.anyio()
 async def test_autoinsert_balances_from_scratch(data_store: DataStore, store_id: bytes32) -> None:
     random = Random()
@@ -705,7 +619,7 @@ async def test_autoinsert_balances_from_scratch(data_store: DataStore, store_id:
         insert_result = await data_store.autoinsert(key, value, store_id, status=Status.COMMITTED)
         hashes.append(insert_result.node_hash)
 
-    heights = {node_hash: len(await data_store.get_ancestors_optimized(node_hash, store_id)) for node_hash in hashes}
+    heights = {node_hash: len(await data_store.get_ancestors(node_hash, store_id)) for node_hash in hashes}
     too_tall = {hash: height for hash, height in heights.items() if height > 14}
     assert too_tall == {}
     assert 11 <= statistics.mean(heights.values()) <= 12
@@ -715,7 +629,7 @@ async def test_autoinsert_balances_from_scratch(data_store: DataStore, store_id:
 async def test_autoinsert_balances_gaps(data_store: DataStore, store_id: bytes32) -> None:
     random = Random()
     random.seed(101, version=2)
-    hashes = []
+    hashes: list[bytes32] = []
 
     for i in range(2000):
         key = (i + 100).to_bytes(4, byteorder="big")
@@ -723,7 +637,7 @@ async def test_autoinsert_balances_gaps(data_store: DataStore, store_id: bytes32
         if i == 0 or i > 10:
             insert_result = await data_store.autoinsert(key, value, store_id, status=Status.COMMITTED)
         else:
-            reference_node_hash = await data_store.get_terminal_node_for_seed(store_id, bytes32.zeros)
+            reference_node_hash = hashes[-1]
             insert_result = await data_store.insert(
                 key=key,
                 value=value,
@@ -732,11 +646,11 @@ async def test_autoinsert_balances_gaps(data_store: DataStore, store_id: bytes32
                 side=Side.LEFT,
                 status=Status.COMMITTED,
             )
-            ancestors = await data_store.get_ancestors_optimized(insert_result.node_hash, store_id)
+            ancestors = await data_store.get_ancestors(insert_result.node_hash, store_id)
             assert len(ancestors) == i
         hashes.append(insert_result.node_hash)
 
-    heights = {node_hash: len(await data_store.get_ancestors_optimized(node_hash, store_id)) for node_hash in hashes}
+    heights = {node_hash: len(await data_store.get_ancestors(node_hash, store_id)) for node_hash in hashes}
     too_tall = {hash: height for hash, height in heights.items() if height > 14}
     assert too_tall == {}
     assert 11 <= statistics.mean(heights.values()) <= 12
@@ -874,24 +788,24 @@ async def test_proof_of_inclusion_by_hash(data_store: DataStore, store_id: bytes
     await _debug_dump(db=data_store.db_wrapper)
 
     expected_layers = [
-        ProofOfInclusionLayer(
+        chia_rs.datalayer.ProofOfInclusionLayer(
             other_hash_side=Side.RIGHT,
             other_hash=bytes32.fromhex("fb66fe539b3eb2020dfbfadfd601fa318521292b41f04c2057c16fca6b947ca1"),
             combined_hash=bytes32.fromhex("36cb1fc56017944213055da8cb0178fb0938c32df3ec4472f5edf0dff85ba4a3"),
         ),
-        ProofOfInclusionLayer(
+        chia_rs.datalayer.ProofOfInclusionLayer(
             other_hash_side=Side.RIGHT,
             other_hash=bytes32.fromhex("6d3af8d93db948e8b6aa4386958e137c6be8bab726db86789594b3588b35adcd"),
             combined_hash=bytes32.fromhex("5f67a0ab1976e090b834bf70e5ce2a0f0a9cd474e19a905348c44ae12274d30b"),
         ),
-        ProofOfInclusionLayer(
+        chia_rs.datalayer.ProofOfInclusionLayer(
             other_hash_side=Side.LEFT,
             other_hash=bytes32.fromhex("c852ecd8fb61549a0a42f9eb9dde65e6c94a01934dbd9c1d35ab94e2a0ae58e2"),
             combined_hash=bytes32.fromhex("7a5193a4e31a0a72f6623dfeb2876022ab74a48abb5966088a1c6f5451cc5d81"),
         ),
     ]
 
-    assert proof == ProofOfInclusion(node_hash=node.hash, layers=expected_layers)
+    assert proof == chia_rs.datalayer.ProofOfInclusion(node_hash=node.hash, layers=expected_layers)
 
 
 @pytest.mark.anyio
@@ -904,26 +818,7 @@ async def test_proof_of_inclusion_by_hash_no_ancestors(data_store: DataStore, st
 
     proof = await data_store.get_proof_of_inclusion_by_hash(node_hash=node.hash, store_id=store_id)
 
-    assert proof == ProofOfInclusion(node_hash=node.hash, layers=[])
-
-
-@pytest.mark.anyio
-async def test_proof_of_inclusion_by_hash_program(data_store: DataStore, store_id: bytes32) -> None:
-    """The proof of inclusion program has the expected Python equivalence."""
-
-    await add_01234567_example(data_store=data_store, store_id=store_id)
-    node = await data_store.get_node_by_key(key=b"\x04", store_id=store_id)
-
-    proof = await data_store.get_proof_of_inclusion_by_hash(node_hash=node.hash, store_id=store_id)
-
-    assert proof.as_program() == [
-        b"\x04",
-        [
-            bytes32.fromhex("fb66fe539b3eb2020dfbfadfd601fa318521292b41f04c2057c16fca6b947ca1"),
-            bytes32.fromhex("6d3af8d93db948e8b6aa4386958e137c6be8bab726db86789594b3588b35adcd"),
-            bytes32.fromhex("c852ecd8fb61549a0a42f9eb9dde65e6c94a01934dbd9c1d35ab94e2a0ae58e2"),
-        ],
-    ]
+    assert proof == chia_rs.datalayer.ProofOfInclusion(node_hash=node.hash, layers=[])
 
 
 @pytest.mark.anyio
@@ -939,27 +834,6 @@ async def test_proof_of_inclusion_by_hash_equals_by_key(data_store: DataStore, s
     assert proof_by_hash == proof_by_key
 
 
-@pytest.mark.anyio
-async def test_proof_of_inclusion_by_hash_bytes(data_store: DataStore, store_id: bytes32) -> None:
-    """The proof of inclusion provided by the data store is able to be converted to a
-    program and subsequently to bytes.
-    """
-    await add_01234567_example(data_store=data_store, store_id=store_id)
-    node = await data_store.get_node_by_key(key=b"\x04", store_id=store_id)
-
-    proof = await data_store.get_proof_of_inclusion_by_hash(node_hash=node.hash, store_id=store_id)
-
-    expected = (
-        b"\xff\x04\xff\xff\xa0\xfbf\xfeS\x9b>\xb2\x02\r\xfb\xfa\xdf\xd6\x01\xfa1\x85!)"
-        b"+A\xf0L W\xc1o\xcak\x94|\xa1\xff\xa0m:\xf8\xd9=\xb9H\xe8\xb6\xaaC\x86\x95"
-        b"\x8e\x13|k\xe8\xba\xb7&\xdb\x86x\x95\x94\xb3X\x8b5\xad\xcd\xff\xa0\xc8R\xec"
-        b"\xd8\xfbaT\x9a\nB\xf9\xeb\x9d\xdee\xe6\xc9J\x01\x93M\xbd\x9c\x1d5\xab\x94"
-        b"\xe2\xa0\xaeX\xe2\x80\x80"
-    )
-
-    assert bytes(proof.as_program()) == expected
-
-
 # @pytest.mark.anyio
 # async def test_create_first_pair(data_store: DataStore, store_id: bytes) -> None:
 #     key = SExp.to([1, 2])
@@ -1036,46 +910,6 @@ async def test_check_roots_are_incrementing_gap(raw_data_store: DataStore) -> No
         await raw_data_store._check_roots_are_incrementing()
 
 
-@pytest.mark.anyio
-async def test_check_hashes_internal(raw_data_store: DataStore) -> None:
-    async with raw_data_store.db_wrapper.writer() as writer:
-        await writer.execute(
-            "INSERT INTO node(hash, node_type, left, right) VALUES(:hash, :node_type, :left, :right)",
-            {
-                "hash": a_bytes_32,
-                "node_type": NodeType.INTERNAL,
-                "left": a_bytes_32,
-                "right": a_bytes_32,
-            },
-        )
-
-    with pytest.raises(
-        NodeHashError,
-        match=r"\n +000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f$",
-    ):
-        await raw_data_store._check_hashes()
-
-
-@pytest.mark.anyio
-async def test_check_hashes_terminal(raw_data_store: DataStore) -> None:
-    async with raw_data_store.db_wrapper.writer() as writer:
-        await writer.execute(
-            "INSERT INTO node(hash, node_type, key, value) VALUES(:hash, :node_type, :key, :value)",
-            {
-                "hash": a_bytes_32,
-                "node_type": NodeType.TERMINAL,
-                "key": Program.to((1, 2)).as_bin(),
-                "value": Program.to((1, 2)).as_bin(),
-            },
-        )
-
-    with pytest.raises(
-        NodeHashError,
-        match=r"\n +000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f$",
-    ):
-        await raw_data_store._check_hashes()
-
-
 @pytest.mark.anyio
 async def test_root_state(data_store: DataStore, store_id: bytes32) -> None:
     key = b"\x01\x02"
@@ -1127,28 +961,29 @@ async def test_kv_diff(data_store: DataStore, store_id: bytes32) -> None:
     insertions = 0
     expected_diff: set[DiffData] = set()
     root_start = None
+
     for i in range(500):
         key = (i + 100).to_bytes(4, byteorder="big")
         value = (i + 200).to_bytes(4, byteorder="big")
         seed = leaf_hash(key=key, value=value)
-        node_hash = await data_store.get_terminal_node_for_seed(store_id, seed)
+        node = await data_store.get_terminal_node_for_seed(seed, store_id)
+        side_seed = bytes(seed)[0]
+        side = None if node is None else (Side.LEFT if side_seed < 128 else Side.RIGHT)
+
         if random.randint(0, 4) > 0 or insertions < 10:
             insertions += 1
-            side = None if node_hash is None else data_store.get_side_for_seed(seed)
-
+            reference_node_hash = node.hash if node is not None else None
             await data_store.insert(
                 key=key,
                 value=value,
                 store_id=store_id,
-                reference_node_hash=node_hash,
-                side=side,
                 status=Status.COMMITTED,
+                reference_node_hash=reference_node_hash,
+                side=side,
             )
             if i > 200:
                 expected_diff.add(DiffData(OperationType.INSERT, key, value))
         else:
-            assert node_hash is not None
-            node = await data_store.get_node(node_hash)
             assert isinstance(node, TerminalNode)
             await data_store.delete(key=node.key, store_id=store_id, status=Status.COMMITTED)
             if i > 200:
@@ -1275,6 +1110,39 @@ async def test_subscribe_unsubscribe(data_store: DataStore, store_id: bytes32) -
     ]
 
 
+@pytest.mark.anyio
+async def test_unsubscribe_clears_databases(data_store: DataStore, store_id: bytes32) -> None:
+    num_inserts = 100
+    await data_store.subscribe(Subscription(store_id, []))
+    for value in range(num_inserts):
+        await data_store.insert(
+            key=value.to_bytes(4, byteorder="big"),
+            value=value.to_bytes(4, byteorder="big"),
+            store_id=store_id,
+            reference_node_hash=None,
+            side=None,
+            status=Status.COMMITTED,
+        )
+    await data_store.add_node_hashes(store_id)
+
+    tables = ["ids", "nodes"]
+    for table in tables:
+        async with data_store.db_wrapper.reader() as reader:
+            async with reader.execute(f"SELECT COUNT(*) FROM {table}") as cursor:
+                row_count = await cursor.fetchone()
+                assert row_count is not None
+                assert row_count[0] > 0
+
+    await data_store.unsubscribe(store_id)
+
+    for table in tables:
+        async with data_store.db_wrapper.reader() as reader:
+            async with reader.execute(f"SELECT COUNT(*) FROM {table}") as cursor:
+                row_count = await cursor.fetchone()
+                assert row_count is not None
+                assert row_count[0] == 0
+
+
 @pytest.mark.anyio
 async def test_server_selection(data_store: DataStore, store_id: bytes32) -> None:
     start_timestamp = 1000
@@ -1410,16 +1278,71 @@ async def test_server_http_ban(
     assert sinfo.ignore_till == start_timestamp  # we don't increase on second failure
 
 
-@pytest.mark.parametrize(
-    "test_delta",
-    [True, False],
-)
+async def get_first_generation(data_store: DataStore, node_hash: bytes32, store_id: bytes32) -> Optional[int]:
+    async with data_store.db_wrapper.reader() as reader:
+        cursor = await reader.execute(
+            "SELECT generation FROM nodes WHERE hash = ? AND store_id = ?",
+            (
+                node_hash,
+                store_id,
+            ),
+        )
+
+        row = await cursor.fetchone()
+        if row is None:
+            return None
+
+        return int(row[0])
+
+
+async def write_tree_to_file_old_format(
+    data_store: DataStore,
+    root: Root,
+    node_hash: bytes32,
+    store_id: bytes32,
+    writer: BinaryIO,
+    merkle_blob: Optional[MerkleBlob] = None,
+    hash_to_index: Optional[dict[bytes32, TreeIndex]] = None,
+) -> None:
+    if node_hash == bytes32.zeros:
+        return
+
+    if merkle_blob is None:
+        merkle_blob = await data_store.get_merkle_blob(store_id=store_id, root_hash=root.node_hash)
+    if hash_to_index is None:
+        hash_to_index = merkle_blob.get_hashes_indexes()
+
+    generation = await get_first_generation(data_store, node_hash, store_id)
+    # Root's generation is not the first time we see this hash, so it's not a new delta.
+    if root.generation != generation:
+        return
+
+    raw_index = hash_to_index[node_hash]
+    raw_node = merkle_blob.get_raw_node(raw_index)
+
+    if isinstance(raw_node, chia_rs.datalayer.InternalNode):
+        left_hash = merkle_blob.get_hash_at_index(raw_node.left)
+        right_hash = merkle_blob.get_hash_at_index(raw_node.right)
+        await write_tree_to_file_old_format(data_store, root, left_hash, store_id, writer, merkle_blob, hash_to_index)
+        await write_tree_to_file_old_format(data_store, root, right_hash, store_id, writer, merkle_blob, hash_to_index)
+        to_write = bytes(SerializedNode(False, bytes(left_hash), bytes(right_hash)))
+    elif isinstance(raw_node, chia_rs.datalayer.LeafNode):
+        node = await data_store.get_terminal_node(raw_node.key, raw_node.value, store_id)
+        to_write = bytes(SerializedNode(True, node.key, node.value))
+    else:
+        raise Exception(f"Node is neither InternalNode nor TerminalNode: {raw_node}")
+
+    writer.write(len(to_write).to_bytes(4, byteorder="big"))
+    writer.write(to_write)
+
+
+@pytest.mark.parametrize(argnames="test_delta", argvalues=["full", "delta", "old"])
 @boolean_datacases(name="group_files_by_store", false="group by singleton", true="don't group by singleton")
 @pytest.mark.anyio
 async def test_data_server_files(
     data_store: DataStore,
     store_id: bytes32,
-    test_delta: bool,
+    test_delta: str,
     group_files_by_store: bool,
     tmp_path: Path,
 ) -> None:
@@ -1428,45 +1351,70 @@ async def test_data_server_files(
     num_ops_per_batch = 100
 
     db_uri = generate_in_memory_db_uri()
-    async with DataStore.managed(database=db_uri, uri=True) as data_store_server:
+    async with DataStore.managed(
+        database=db_uri,
+        uri=True,
+        merkle_blobs_path=tmp_path.joinpath("merkle-blobs"),
+        key_value_blobs_path=tmp_path.joinpath("key-value-blobs"),
+    ) as data_store_server:
         await data_store_server.create_tree(store_id, status=Status.COMMITTED)
         random = Random()
         random.seed(100, version=2)
 
         keys: list[bytes] = []
         counter = 0
-
-        for batch in range(num_batches):
-            changelist: list[dict[str, Any]] = []
-            for operation in range(num_ops_per_batch):
-                if random.randint(0, 4) > 0 or len(keys) == 0:
-                    key = counter.to_bytes(4, byteorder="big")
-                    value = (2 * counter).to_bytes(4, byteorder="big")
-                    keys.append(key)
-                    changelist.append({"action": "insert", "key": key, "value": value})
+        num_repeats = 2
+
+        # Repeat twice to guarantee there will be hashes from the old file format
+        for _ in range(num_repeats):
+            for batch in range(num_batches):
+                changelist: list[dict[str, Any]] = []
+                if batch == num_batches - 1:
+                    for key in keys:
+                        changelist.append({"action": "delete", "key": key})
+                    keys = []
+                    counter = 0
                 else:
-                    key = random.choice(keys)
-                    keys.remove(key)
-                    changelist.append({"action": "delete", "key": key})
-                counter += 1
-            await data_store_server.insert_batch(store_id, changelist, status=Status.COMMITTED)
-            root = await data_store_server.get_tree_root(store_id)
-            await write_files_for_root(
-                data_store_server, store_id, root, tmp_path, 0, group_by_store=group_files_by_store
-            )
-            roots.append(root)
+                    for operation in range(num_ops_per_batch):
+                        if random.randint(0, 4) > 0 or len(keys) == 0:
+                            key = counter.to_bytes(4, byteorder="big")
+                            value = (2 * counter).to_bytes(4, byteorder="big")
+                            keys.append(key)
+                            changelist.append({"action": "insert", "key": key, "value": value})
+                        else:
+                            key = random.choice(keys)
+                            keys.remove(key)
+                            changelist.append({"action": "delete", "key": key})
+                        counter += 1
+
+                await data_store_server.insert_batch(store_id, changelist, status=Status.COMMITTED)
+                root = await data_store_server.get_tree_root(store_id)
+                await data_store_server.add_node_hashes(store_id)
+                if test_delta == "old":
+                    node_hash = root.node_hash if root.node_hash is not None else bytes32.zeros
+                    filename = get_delta_filename_path(
+                        tmp_path, store_id, node_hash, root.generation, group_files_by_store
+                    )
+                    filename.parent.mkdir(parents=True, exist_ok=True)
+                    with open(filename, "xb") as writer:
+                        await write_tree_to_file_old_format(data_store_server, root, node_hash, store_id, writer)
+                else:
+                    await write_files_for_root(
+                        data_store_server, store_id, root, tmp_path, 0, group_by_store=group_files_by_store
+                    )
+                roots.append(root)
 
     generation = 1
-    assert len(roots) == num_batches
+    assert len(roots) == num_batches * num_repeats
     for root in roots:
-        assert root.node_hash is not None
-        if not test_delta:
-            filename = get_full_tree_filename_path(tmp_path, store_id, root.node_hash, generation, group_files_by_store)
+        node_hash = root.node_hash if root.node_hash is not None else bytes32.zeros
+        if test_delta == "full":
+            filename = get_full_tree_filename_path(tmp_path, store_id, node_hash, generation, group_files_by_store)
             assert filename.exists()
         else:
-            filename = get_delta_filename_path(tmp_path, store_id, root.node_hash, generation, group_files_by_store)
+            filename = get_delta_filename_path(tmp_path, store_id, node_hash, generation, group_files_by_store)
             assert filename.exists()
-        await insert_into_data_store_from_file(data_store, store_id, root.node_hash, tmp_path.joinpath(filename))
+        await data_store.insert_into_data_store_from_file(store_id, root.node_hash, tmp_path.joinpath(filename))
         current_root = await data_store.get_tree_root(store_id=store_id)
         assert current_root.node_hash == root.node_hash
         generation += 1
@@ -1551,6 +1499,20 @@ class BatchesInsertBenchmarkCase:
         return f"count={self.count},batch_count={self.batch_count}"
 
 
+@dataclass
+class BatchUpdateBenchmarkCase:
+    pre: int
+    num_inserts: int
+    num_deletes: int
+    num_upserts: int
+    limit: float
+    marks: Marks = ()
+
+    @property
+    def id(self) -> str:
+        return f"pre={self.pre},inserts={self.num_inserts},deletes={self.num_deletes},upserts={self.num_upserts}"
+
+
 @datacases(
     BatchInsertBenchmarkCase(
         pre=0,
@@ -1589,16 +1551,55 @@ async def test_benchmark_batch_insert_speed(
     r.seed("shadowlands", version=2)
 
     changelist = [
+        {"action": "insert", "key": x.to_bytes(32, byteorder="big", signed=False), "value": r.randbytes(1200)}
+        for x in range(case.pre + case.count)
+    ]
+
+    pre = changelist[: case.pre]
+    batch = changelist[case.pre : case.pre + case.count]
+
+    if case.pre > 0:
+        await data_store.insert_batch(
+            store_id=store_id,
+            changelist=pre,
+            status=Status.COMMITTED,
+        )
+
+    with benchmark_runner.assert_runtime(seconds=case.limit):
+        await data_store.insert_batch(
+            store_id=store_id,
+            changelist=batch,
+        )
+
+
+@datacases(
+    BatchUpdateBenchmarkCase(
+        pre=1_000,
+        num_inserts=1_000,
+        num_deletes=500,
+        num_upserts=500,
+        limit=36,
+    ),
+)
+@pytest.mark.anyio
+async def test_benchmark_batch_update_speed(
+    data_store: DataStore,
+    store_id: bytes32,
+    benchmark_runner: BenchmarkRunner,
+    case: BatchUpdateBenchmarkCase,
+) -> None:
+    r = random.Random()
+    r.seed("shadowlands", version=2)
+
+    pre = [
         {
             "action": "insert",
             "key": x.to_bytes(32, byteorder="big", signed=False),
             "value": bytes(r.getrandbits(8) for _ in range(1200)),
         }
-        for x in range(case.pre + case.count)
+        for x in range(case.pre)
     ]
-
-    pre = changelist[: case.pre]
-    batch = changelist[case.pre : case.pre + case.count]
+    batch = []
 
     if case.pre > 0:
         await data_store.insert_batch(
@@ -1607,6 +1608,44 @@ async def test_benchmark_batch_insert_speed(
             status=Status.COMMITTED,
         )
 
+    keys = [x.to_bytes(32, byteorder="big", signed=False) for x in range(case.pre)]
+    for operation in range(case.num_inserts):
+        key = (operation + case.pre).to_bytes(32, byteorder="big", signed=False)
+        batch.append(
+            {
+                "action": "insert",
+                "key": key,
+                "value": bytes(r.getrandbits(8) for _ in range(1200)),
+            }
+        )
+        keys.append(key)
+
+    if case.num_deletes > 0:
+        r.shuffle(keys)
+        assert len(keys) >= case.num_deletes
+        batch.extend(
+            {
+                "action": "delete",
+                "key": key,
+            }
+            for key in keys[: case.num_deletes]
+        )
+        keys = keys[case.num_deletes :]
+
+    if case.num_upserts > 0:
+        assert len(keys) > 0
+        r.shuffle(keys)
+        batch.extend(
+            [
+                {
+                    "action": "upsert",
+                    "key": keys[operation % len(keys)],
+                    "value": bytes(r.getrandbits(8) for _ in range(1200)),
+                }
+                for operation in range(case.num_upserts)
+            ]
+        )
+
     with benchmark_runner.assert_runtime(seconds=case.limit):
         await data_store.insert_batch(
             store_id=store_id,
@@ -1614,6 +1653,22 @@ async def test_benchmark_batch_insert_speed(
         )
 
 
+@datacases(
+    BatchInsertBenchmarkCase(
+        pre=0,
+        count=50,
+        limit=2,
+    ),
+)
+@pytest.mark.anyio
+async def test_benchmark_tool(
+    benchmark_runner: BenchmarkRunner,
+    case: BatchInsertBenchmarkCase,
+) -> None:
+    with benchmark_runner.assert_runtime(seconds=case.limit):
+        await generate_datastore(case.count)
+
+
 @datacases(
     BatchesInsertBenchmarkCase(
         count=50,
@@ -1637,7 +1692,7 @@ async def test_benchmark_batch_insert_speed_multiple_batches(
                 {
                     "action": "insert",
                     "key": x.to_bytes(32, byteorder="big", signed=False),
-                    "value": bytes(r.getrandbits(8) for _ in range(10000)),
+                    "value": r.randbytes(10000),
                 }
                 for x in range(batch * case.count, (batch + 1) * case.count)
             ]
@@ -1648,189 +1703,6 @@ async def test_benchmark_batch_insert_speed_multiple_batches(
             )
 
 
-@pytest.mark.anyio
-async def test_delete_store_data(raw_data_store: DataStore) -> None:
-    store_id = bytes32.zeros
-    store_id_2 = bytes32(b"\0" * 31 + b"\1")
-    await raw_data_store.create_tree(store_id=store_id, status=Status.COMMITTED)
-    await raw_data_store.create_tree(store_id=store_id_2, status=Status.COMMITTED)
-    total_keys = 4
-    keys = [key.to_bytes(4, byteorder="big") for key in range(total_keys)]
-    batch1 = [
-        {"action": "insert", "key": keys[0], "value": keys[0]},
-        {"action": "insert", "key": keys[1], "value": keys[1]},
-    ]
-    batch2 = batch1.copy()
-    batch1.append({"action": "insert", "key": keys[2], "value": keys[2]})
-    batch2.append({"action": "insert", "key": keys[3], "value": keys[3]})
-    assert batch1 != batch2
-    await raw_data_store.insert_batch(store_id, batch1, status=Status.COMMITTED)
-    await raw_data_store.insert_batch(store_id_2, batch2, status=Status.COMMITTED)
-    keys_values_before = await raw_data_store.get_keys_values(store_id_2)
-    async with raw_data_store.db_wrapper.reader() as reader:
-        result = await reader.execute("SELECT * FROM node")
-        nodes = await result.fetchall()
-        kv_nodes_before = {}
-        for node in nodes:
-            if node["key"] is not None:
-                kv_nodes_before[node["key"]] = node["value"]
-    assert [kv_nodes_before[key] for key in keys] == keys
-    await raw_data_store.delete_store_data(store_id)
-    # Deleting from `node` table doesn't alter other stores.
-    keys_values_after = await raw_data_store.get_keys_values(store_id_2)
-    assert keys_values_before == keys_values_after
-    async with raw_data_store.db_wrapper.reader() as reader:
-        result = await reader.execute("SELECT * FROM node")
-        nodes = await result.fetchall()
-        kv_nodes_after = {}
-        for node in nodes:
-            if node["key"] is not None:
-                kv_nodes_after[node["key"]] = node["value"]
-    for i in range(total_keys):
-        if i != 2:
-            assert kv_nodes_after[keys[i]] == keys[i]
-        else:
-            # `keys[2]` was only present in the first store.
-            assert keys[i] not in kv_nodes_after
-    assert not await raw_data_store.store_id_exists(store_id)
-    await raw_data_store.delete_store_data(store_id_2)
-    async with raw_data_store.db_wrapper.reader() as reader:
-        async with reader.execute("SELECT COUNT(*) FROM node") as cursor:
-            row_count = await cursor.fetchone()
-            assert row_count is not None
-            assert row_count[0] == 0
-    assert not await raw_data_store.store_id_exists(store_id_2)
-
-
-@pytest.mark.anyio
-async def test_delete_store_data_multiple_stores(raw_data_store: DataStore) -> None:
-    # Make sure inserting and deleting the same data works
-    for repetition in range(2):
-        num_stores = 50
-        total_keys = 150
-        keys_deleted_per_store = 3
-        store_ids = [bytes32(i.to_bytes(32, byteorder="big")) for i in range(num_stores)]
-        for store_id in store_ids:
-            await raw_data_store.create_tree(store_id=store_id, status=Status.COMMITTED)
-        original_keys = [key.to_bytes(4, byteorder="big") for key in range(total_keys)]
-        batches = []
-        for i in range(num_stores):
-            batch = [
-                {"action": "insert", "key": key, "value": key} for key in original_keys[i * keys_deleted_per_store :]
-            ]
-            batches.append(batch)
-
-        for store_id, batch in zip(store_ids, batches):
-            await raw_data_store.insert_batch(store_id, batch, status=Status.COMMITTED)
-
-        for tree_index in range(num_stores):
-            async with raw_data_store.db_wrapper.reader() as reader:
-                result = await reader.execute("SELECT * FROM node")
-                nodes = await result.fetchall()
-
-            keys = {node["key"] for node in nodes if node["key"] is not None}
-            assert len(keys) == total_keys - tree_index * keys_deleted_per_store
-            keys_after_index = set(original_keys[tree_index * keys_deleted_per_store :])
-            keys_before_index = set(original_keys[: tree_index * keys_deleted_per_store])
-            assert keys_after_index.issubset(keys)
-            assert keys.isdisjoint(keys_before_index)
-            await raw_data_store.delete_store_data(store_ids[tree_index])
-
-        async with raw_data_store.db_wrapper.reader() as reader:
-            async with reader.execute("SELECT COUNT(*) FROM node") as cursor:
-                row_count = await cursor.fetchone()
-                assert row_count is not None
-                assert row_count[0] == 0
-
-
-@pytest.mark.parametrize("common_keys_count", [1, 250, 499])
-@pytest.mark.anyio
-async def test_delete_store_data_with_common_values(raw_data_store: DataStore, common_keys_count: int) -> None:
-    store_id_1 = bytes32(b"\x00" * 31 + b"\x01")
-    store_id_2 = bytes32(b"\x00" * 31 + b"\x02")
-
-    await raw_data_store.create_tree(store_id=store_id_1, status=Status.COMMITTED)
-    await raw_data_store.create_tree(store_id=store_id_2, status=Status.COMMITTED)
-
-    key_offset = 1000
-    total_keys_per_store = 500
-    assert common_keys_count < key_offset
-    common_keys = {key.to_bytes(4, byteorder="big") for key in range(common_keys_count)}
-    unique_keys_1 = {
-        (key + key_offset).to_bytes(4, byteorder="big") for key in range(total_keys_per_store - common_keys_count)
-    }
-    unique_keys_2 = {
-        (key + (2 * key_offset)).to_bytes(4, byteorder="big") for key in range(total_keys_per_store - common_keys_count)
-    }
-
-    batch1 = [{"action": "insert", "key": key, "value": key} for key in common_keys.union(unique_keys_1)]
-    batch2 = [{"action": "insert", "key": key, "value": key} for key in common_keys.union(unique_keys_2)]
-
-    await raw_data_store.insert_batch(store_id_1, batch1, status=Status.COMMITTED)
-    await raw_data_store.insert_batch(store_id_2, batch2, status=Status.COMMITTED)
-
-    await raw_data_store.delete_store_data(store_id_1)
-    async with raw_data_store.db_wrapper.reader() as reader:
-        result = await reader.execute("SELECT * FROM node")
-        nodes = await result.fetchall()
-
-    keys = {node["key"] for node in nodes if node["key"] is not None}
-    # Since one store got all its keys deleted, we're left only with the keys of the other store.
-    assert len(keys) == total_keys_per_store
-    assert keys.intersection(unique_keys_1) == set()
-    assert keys.symmetric_difference(common_keys.union(unique_keys_2)) == set()
-
-
-@pytest.mark.anyio
-@pytest.mark.parametrize("pending_status", [Status.PENDING, Status.PENDING_BATCH])
-async def test_delete_store_data_protects_pending_roots(raw_data_store: DataStore, pending_status: Status) -> None:
-    num_stores = 5
-    total_keys = 15
-    store_ids = [bytes32(i.to_bytes(32, byteorder="big")) for i in range(num_stores)]
-    for store_id in store_ids:
-        await raw_data_store.create_tree(store_id=store_id, status=Status.COMMITTED)
-    original_keys = [key.to_bytes(4, byteorder="big") for key in range(total_keys)]
-    batches = []
-    keys_per_pending_root = 2
-
-    for i in range(num_stores - 1):
-        start_index = i * keys_per_pending_root
-        end_index = (i + 1) * keys_per_pending_root
-        batch = [{"action": "insert", "key": key, "value": key} for key in original_keys[start_index:end_index]]
-        batches.append(batch)
-    for store_id, batch in zip(store_ids, batches):
-        await raw_data_store.insert_batch(store_id, batch, status=pending_status)
-
-    store_id = store_ids[-1]
-    batch = [{"action": "insert", "key": key, "value": key} for key in original_keys]
-    await raw_data_store.insert_batch(store_id, batch, status=Status.COMMITTED)
-
-    async with raw_data_store.db_wrapper.reader() as reader:
-        result = await reader.execute("SELECT * FROM node")
-        nodes = await result.fetchall()
-
-    keys = {node["key"] for node in nodes if node["key"] is not None}
-    assert keys == set(original_keys)
-
-    await raw_data_store.delete_store_data(store_id)
-    async with raw_data_store.db_wrapper.reader() as reader:
-        result = await reader.execute("SELECT * FROM node")
-        nodes = await result.fetchall()
-
-    keys = {node["key"] for node in nodes if node["key"] is not None}
-    assert keys == set(original_keys[: (num_stores - 1) * keys_per_pending_root])
-
-    for index in range(num_stores - 1):
-        store_id = store_ids[index]
-        root = await raw_data_store.get_pending_root(store_id)
-        assert root is not None
-        await raw_data_store.change_root_status(root, Status.COMMITTED)
-        kv = await raw_data_store.get_keys_values(store_id=store_id)
-        start_index = index * keys_per_pending_root
-        end_index = (index + 1) * keys_per_pending_root
-        assert {pair.key for pair in kv} == set(original_keys[start_index:end_index])
-
-
 @pytest.mark.anyio
 @boolean_datacases(name="group_files_by_store", true="group by singleton", false="don't group by singleton")
 @pytest.mark.parametrize("max_full_files", [1, 2, 5])
@@ -1843,7 +1715,6 @@ async def test_insert_from_delta_file(
     group_files_by_store: bool,
     max_full_files: int,
 ) -> None:
-    await data_store.create_tree(store_id=store_id, status=Status.COMMITTED)
     num_files = 5
     for generation in range(num_files):
         key = generation.to_bytes(4, byteorder="big")
@@ -1854,30 +1725,34 @@ async def test_insert_from_delta_file(
             store_id=store_id,
             status=Status.COMMITTED,
         )
+        await data_store.add_node_hashes(store_id)
 
     root = await data_store.get_tree_root(store_id=store_id)
-    assert root.generation == num_files + 1
+    assert root.generation == num_files
     root_hashes = []
 
     tmp_path_1 = tmp_path.joinpath("1")
     tmp_path_2 = tmp_path.joinpath("2")
 
-    for generation in range(1, num_files + 2):
+    for generation in range(1, num_files + 1):
         root = await data_store.get_tree_root(store_id=store_id, generation=generation)
         await write_files_for_root(data_store, store_id, root, tmp_path_1, 0, False, group_files_by_store)
         root_hashes.append(bytes32.zeros if root.node_hash is None else root.node_hash)
     store_path = tmp_path_1.joinpath(f"{store_id}") if group_files_by_store else tmp_path_1
     with os.scandir(store_path) as entries:
         filenames = {entry.name for entry in entries}
-        assert len(filenames) == 2 * (num_files + 1)
+        assert len(filenames) == 2 * num_files
     for filename in filenames:
         if "full" in filename:
             store_path.joinpath(filename).unlink()
     with os.scandir(store_path) as entries:
         filenames = {entry.name for entry in entries}
-        assert len(filenames) == num_files + 1
+        assert len(filenames) == num_files
     kv_before = await data_store.get_keys_values(store_id=store_id)
     await data_store.rollback_to_generation(store_id, 0)
+    with contextlib.suppress(FileNotFoundError):
+        shutil.rmtree(data_store.merkle_blobs_path)
+
     root = await data_store.get_tree_root(store_id=store_id)
     assert root.generation == 0
     os.rename(store_path, tmp_path_2)
@@ -1950,12 +1825,13 @@ async def test_insert_from_delta_file(
         assert success
 
     root = await data_store.get_tree_root(store_id=store_id)
-    assert root.generation == num_files + 1
+    assert root.generation == num_files
     with os.scandir(store_path) as entries:
         filenames = {entry.name for entry in entries}
-        assert len(filenames) == num_files + 1 + max_full_files  # 6 deltas and max_full_files full files
+        assert len(filenames) == num_files + max_full_files - 1
     kv = await data_store.get_keys_values(store_id=store_id)
-    assert kv == kv_before
+    # order agnostic comparison of the list
+    assert set(kv) == set(kv_before)
 
 
 @pytest.mark.anyio
@@ -1993,7 +1869,7 @@ async def test_get_node_by_key_with_overlapping_keys(raw_data_store: DataStore)
                 if random.randint(0, 4) == 0:
                     batch = [{"action": "delete", "key": key}]
                     await raw_data_store.insert_batch(store_id, batch, status=Status.COMMITTED)
-                    with pytest.raises(KeyNotFoundError, match=f"Key not found: {key.hex()}"):
+                    with pytest.raises(chia_rs.datalayer.UnknownKeyError):
                         await raw_data_store.get_node_by_key(store_id=store_id, key=key)
 
 
@@ -2013,6 +1889,7 @@ async def test_insert_from_delta_file_correct_file_exists(
             store_id=store_id,
             status=Status.COMMITTED,
         )
+        await data_store.add_node_hashes(store_id)
 
     root = await data_store.get_tree_root(store_id=store_id)
     assert root.generation == num_files + 1
@@ -2023,18 +1900,20 @@ async def test_insert_from_delta_file_correct_file_exists(
         root_hashes.append(bytes32.zeros if root.node_hash is None else root.node_hash)
     store_path = tmp_path.joinpath(f"{store_id}") if group_files_by_store else tmp_path
     with os.scandir(store_path) as entries:
-        filenames = {entry.name for entry in entries}
+        filenames = {entry.name for entry in entries if entry.name.endswith(".dat")}
         assert len(filenames) == 2 * (num_files + 1)
     for filename in filenames:
         if "full" in filename:
             store_path.joinpath(filename).unlink()
     with os.scandir(store_path) as entries:
-        filenames = {entry.name for entry in entries}
+        filenames = {entry.name for entry in entries if entry.name.endswith(".dat")}
         assert len(filenames) == num_files + 1
     kv_before = await data_store.get_keys_values(store_id=store_id)
     await data_store.rollback_to_generation(store_id, 0)
     root = await data_store.get_tree_root(store_id=store_id)
     assert root.generation == 0
+    with contextlib.suppress(FileNotFoundError):
+        shutil.rmtree(data_store.merkle_blobs_path)
 
     sinfo = ServerInfo("http://127.0.0.1/8003", 0, 0)
     success = await insert_from_delta_file(
@@ -2056,10 +1935,11 @@ async def test_insert_from_delta_file_correct_file_exists(
     root = await data_store.get_tree_root(store_id=store_id)
     assert root.generation == num_files + 1
     with os.scandir(store_path) as entries:
-        filenames = {entry.name for entry in entries}
+        filenames = {entry.name for entry in entries if entry.name.endswith(".dat")}
         assert len(filenames) == num_files + 2  # 1 full and 6 deltas
     kv = await data_store.get_keys_values(store_id=store_id)
-    assert kv == kv_before
+    # order agnostic comparison of the list
+    assert set(kv) == set(kv_before)
 
 
 @pytest.mark.anyio
@@ -2079,6 +1959,7 @@ async def test_insert_from_delta_file_incorrect_file_exists(
         store_id=store_id,
         status=Status.COMMITTED,
     )
+    await data_store.add_node_hashes(store_id)
 
     root = await data_store.get_tree_root(store_id=store_id)
     assert root.generation == 2
@@ -2087,7 +1968,7 @@ async def test_insert_from_delta_file_incorrect_file_exists(
     incorrect_root_hash = bytes32([0] * 31 + [1])
     store_path = tmp_path.joinpath(f"{store_id}") if group_files_by_store else tmp_path
     with os.scandir(store_path) as entries:
-        filenames = [entry.name for entry in entries]
+        filenames = [entry.name for entry in entries if entry.name.endswith(".dat")]
         assert len(filenames) == 2
         os.rename(
             store_path.joinpath(filenames[0]),
@@ -2119,7 +2000,7 @@ async def test_insert_from_delta_file_incorrect_file_exists(
     root = await data_store.get_tree_root(store_id=store_id)
     assert root.generation == 1
     with os.scandir(store_path) as entries:
-        filenames = [entry.name for entry in entries]
+        filenames = [entry.name for entry in entries if entry.name.endswith(".dat")]
         assert len(filenames) == 0
 
 
@@ -2130,7 +2011,7 @@ async def test_insert_key_already_present(data_store: DataStore, store_id: bytes
     await data_store.insert(
         key=key, value=value, store_id=store_id, reference_node_hash=None, side=None, status=Status.COMMITTED
     )
-    with pytest.raises(Exception, match=f"Key already present: {key.hex()}"):
+    with pytest.raises(KeyAlreadyPresentError):
         await data_store.insert(key=key, value=value, store_id=store_id, reference_node_hash=None, side=None)
 
 
@@ -2145,7 +2026,7 @@ async def test_batch_insert_key_already_present(
     value = b"bar"
     changelist = [{"action": "insert", "key": key, "value": value}]
     await data_store.insert_batch(store_id, changelist, Status.COMMITTED, use_batch_autoinsert)
-    with pytest.raises(Exception, match=f"Key already present: {key.hex()}"):
+    with pytest.raises(KeyAlreadyPresentError):
         await data_store.insert_batch(store_id, changelist, Status.COMMITTED, use_batch_autoinsert)
 
 
@@ -2188,7 +2069,7 @@ async def test_update_keys(data_store: DataStore, store_id: bytes32, use_upsert:
 
 
 @pytest.mark.anyio
-async def test_migration_unknown_version(data_store: DataStore) -> None:
+async def test_migration_unknown_version(data_store: DataStore, tmp_path: Path) -> None:
     async with data_store.db_wrapper.writer() as writer:
         await writer.execute(
             "INSERT INTO schema(version_id) VALUES(:version_id)",
@@ -2197,228 +2078,284 @@ async def test_migration_unknown_version(data_store: DataStore) -> None:
             },
         )
     with pytest.raises(Exception, match="Unknown version"):
-        await data_store.migrate_db()
-
-
-async def _check_ancestors(
-    data_store: DataStore, store_id: bytes32, root_hash: bytes32
-) -> dict[bytes32, Optional[bytes32]]:
-    ancestors: dict[bytes32, Optional[bytes32]] = {}
-    root_node: Node = await data_store.get_node(root_hash)
-    queue: list[Node] = [root_node]
-
-    while queue:
-        node = queue.pop(0)
-        if isinstance(node, InternalNode):
-            left_node = await data_store.get_node(node.left_hash)
-            right_node = await data_store.get_node(node.right_hash)
-            ancestors[left_node.hash] = node.hash
-            ancestors[right_node.hash] = node.hash
-            queue.append(left_node)
-            queue.append(right_node)
-
-    ancestors[root_hash] = None
-    for node_hash, ancestor_hash in ancestors.items():
-        ancestor_node = await data_store._get_one_ancestor(node_hash, store_id)
-        if ancestor_hash is None:
-            assert ancestor_node is None
-        else:
-            assert ancestor_node is not None
-            assert ancestor_node.hash == ancestor_hash
+        await data_store.migrate_db(tmp_path)
 
-    return ancestors
+
+@boolean_datacases(name="group_files_by_store", false="group by singleton", true="don't group by singleton")
+@pytest.mark.anyio
+async def test_migration(
+    data_store: DataStore,
+    store_id: bytes32,
+    group_files_by_store: bool,
+    tmp_path: Path,
+) -> None:
+    num_batches = 10
+    num_ops_per_batch = 100
+    keys: list[bytes] = []
+    counter = 0
+    random = Random()
+    random.seed(100, version=2)
+
+    for batch in range(num_batches):
+        changelist: list[dict[str, Any]] = []
+        for operation in range(num_ops_per_batch):
+            if random.randint(0, 4) > 0 or len(keys) == 0:
+                key = counter.to_bytes(4, byteorder="big")
+                value = (2 * counter).to_bytes(4, byteorder="big")
+                keys.append(key)
+                changelist.append({"action": "insert", "key": key, "value": value})
+            else:
+                key = random.choice(keys)
+                keys.remove(key)
+                changelist.append({"action": "delete", "key": key})
+            counter += 1
+        await data_store.insert_batch(store_id, changelist, status=Status.COMMITTED)
+        root = await data_store.get_tree_root(store_id)
+        await data_store.add_node_hashes(store_id)
+        await write_files_for_root(data_store, store_id, root, tmp_path, 0, group_by_store=group_files_by_store)
+
+    kv_before = await data_store.get_keys_values(store_id=store_id)
+    async with data_store.db_wrapper.writer(foreign_key_enforcement_enabled=False) as writer:
+        tables = [table for table in table_columns.keys() if table != "root"]
+        for table in tables:
+            await writer.execute(f"DELETE FROM {table}")
+
+    with contextlib.suppress(FileNotFoundError):
+        shutil.rmtree(data_store.merkle_blobs_path)
+    with contextlib.suppress(FileNotFoundError):
+        shutil.rmtree(data_store.key_value_blobs_path)
+
+    data_store.recent_merkle_blobs = LRUCache(capacity=128)
+    assert await data_store.get_keys_values(store_id=store_id) == []
+    await data_store.migrate_db(tmp_path)
+    # order agnostic comparison of the list
+    assert set(await data_store.get_keys_values(store_id=store_id)) == set(kv_before)
 
 
 @pytest.mark.anyio
-async def test_build_ancestor_table(data_store: DataStore, store_id: bytes32) -> None:
-    num_values = 1000
+@pytest.mark.parametrize("num_keys", [10, 1000])
+async def test_get_existing_hashes(
+    data_store: DataStore,
+    store_id: bytes32,
+    num_keys: int,
+) -> None:
     changelist: list[dict[str, Any]] = []
-    for value in range(num_values):
-        value_bytes = value.to_bytes(4, byteorder="big")
-        changelist.append({"action": "upsert", "key": value_bytes, "value": value_bytes})
-    await data_store.insert_batch(
-        store_id=store_id,
-        changelist=changelist,
-        status=Status.PENDING,
-    )
+    for i in range(num_keys):
+        key = i.to_bytes(4, byteorder="big")
+        value = (2 * i).to_bytes(4, byteorder="big")
+        changelist.append({"action": "insert", "key": key, "value": value})
+    await data_store.insert_batch(store_id, changelist, status=Status.COMMITTED)
+    await data_store.add_node_hashes(store_id)
 
-    pending_root = await data_store.get_pending_root(store_id=store_id)
-    assert pending_root is not None
-    assert pending_root.node_hash is not None
-    await data_store.change_root_status(pending_root, Status.COMMITTED)
-    await data_store.build_ancestor_table_for_latest_root(store_id=store_id)
-
-    assert pending_root.node_hash is not None
-    await _check_ancestors(data_store, store_id, pending_root.node_hash)
+    root = await data_store.get_tree_root(store_id=store_id)
+    merkle_blob = await data_store.get_merkle_blob(store_id=store_id, root_hash=root.node_hash)
+    hash_to_index = merkle_blob.get_hashes_indexes()
+    existing_hashes = list(hash_to_index.keys())
+    not_existing_hashes = [bytes32(i.to_bytes(32, byteorder="big")) for i in range(num_keys)]
+    result = await data_store.get_existing_hashes(existing_hashes + not_existing_hashes, store_id)
+    assert result == set(existing_hashes)
 
 
 @pytest.mark.anyio
-async def test_sparse_ancestor_table(data_store: DataStore, store_id: bytes32) -> None:
-    num_values = 100
-    for value in range(num_values):
-        value_bytes = value.to_bytes(4, byteorder="big")
-        await data_store.autoinsert(
-            key=value_bytes,
-            value=value_bytes,
-            store_id=store_id,
-            status=Status.COMMITTED,
-        )
-        root = await data_store.get_tree_root(store_id=store_id)
-        assert root.node_hash is not None
-        ancestors = await _check_ancestors(data_store, store_id, root.node_hash)
+@pytest.mark.parametrize(argnames="size_offset", argvalues=[-1, 0, 1])
+async def test_basic_key_value_db_vs_disk_cutoff(
+    data_store: DataStore,
+    store_id: bytes32,
+    seeded_random: random.Random,
+    size_offset: int,
+) -> None:
+    size = data_store.prefer_db_kv_blob_length + size_offset
 
-    # Check the ancestor table is sparse
-    root_generation = root.generation
-    current_generation_count = 0
-    previous_generation_count = 0
-    for node_hash, ancestor_hash in ancestors.items():
-        async with data_store.db_wrapper.reader() as reader:
-            if ancestor_hash is not None:
-                cursor = await reader.execute(
-                    "SELECT MAX(generation) AS generation FROM ancestors WHERE hash == :hash AND ancestor == :ancestor",
-                    {"hash": node_hash, "ancestor": ancestor_hash},
-                )
-            else:
-                cursor = await reader.execute(
-                    "SELECT MAX(generation) AS generation FROM ancestors WHERE hash == :hash AND ancestor IS NULL",
-                    {"hash": node_hash},
-                )
+    blob = bytes(seeded_random.getrandbits(8) for _ in range(size))
+    blob_hash = bytes32(sha256(blob).digest())
+    async with data_store.db_wrapper.writer() as writer:
+        with data_store.manage_kv_files(store_id):
+            await data_store.add_kvid(blob=blob, store_id=store_id, writer=writer)
+
+    file_exists = data_store.get_key_value_path(store_id=store_id, blob_hash=blob_hash).exists()
+    async with data_store.db_wrapper.writer() as writer:
+        async with writer.execute(
+            "SELECT blob FROM ids WHERE hash = :blob_hash",
+            {"blob_hash": blob_hash},
+        ) as cursor:
             row = await cursor.fetchone()
             assert row is not None
-            generation = row["generation"]
-            assert generation <= root_generation
-            if generation == root_generation:
-                current_generation_count += 1
-            else:
-                previous_generation_count += 1
+            db_blob: Optional[bytes] = row["blob"]
 
-    assert current_generation_count == 15
-    assert previous_generation_count == 184
+    if size_offset <= 0:
+        assert not file_exists
+        assert db_blob == blob
+    else:
+        assert file_exists
+        assert db_blob is None
 
 
-async def get_all_nodes(data_store: DataStore, store_id: bytes32) -> list[Node]:
-    root = await data_store.get_tree_root(store_id)
-    assert root.node_hash is not None
-    root_node = await data_store.get_node(root.node_hash)
-    nodes: list[Node] = []
-    queue: list[Node] = [root_node]
+@pytest.mark.anyio
+@pytest.mark.parametrize(argnames="size_offset", argvalues=[-1, 0, 1])
+@pytest.mark.parametrize(argnames="limit_change", argvalues=[-2, -1, 1, 2])
+async def test_changing_key_value_db_vs_disk_cutoff(
+    data_store: DataStore,
+    store_id: bytes32,
+    seeded_random: random.Random,
+    size_offset: int,
+    limit_change: int,
+) -> None:
+    size = data_store.prefer_db_kv_blob_length + size_offset
 
-    while len(queue) > 0:
-        node = queue.pop(0)
-        nodes.append(node)
-        if isinstance(node, InternalNode):
-            left_node = await data_store.get_node(node.left_hash)
-            right_node = await data_store.get_node(node.right_hash)
-            queue.append(left_node)
-            queue.append(right_node)
+    blob = bytes(seeded_random.getrandbits(8) for _ in range(size))
+    async with data_store.db_wrapper.writer() as writer:
+        with data_store.manage_kv_files(store_id):
+            kv_id = await data_store.add_kvid(blob=blob, store_id=store_id, writer=writer)
+
+    data_store.prefer_db_kv_blob_length += limit_change
+    retrieved_blob = await data_store.get_blob_from_kvid(kv_id=kv_id, store_id=store_id)
 
-    return nodes
+    assert blob == retrieved_blob
 
 
 @pytest.mark.anyio
-async def test_get_nodes(data_store: DataStore, store_id: bytes32) -> None:
-    num_values = 50
-    changelist: list[dict[str, Any]] = []
+async def test_get_keys_both_disk_and_db(
+    data_store: DataStore,
+    store_id: bytes32,
+    seeded_random: random.Random,
+) -> None:
+    inserted_keys: set[bytes] = set()
 
-    for value in range(num_values):
-        value_bytes = value.to_bytes(4, byteorder="big")
-        changelist.append({"action": "upsert", "key": value_bytes, "value": value_bytes})
-    await data_store.insert_batch(
-        store_id=store_id,
-        changelist=changelist,
-        status=Status.COMMITTED,
-    )
+    for size_offset in [-1, 0, 1]:
+        size = data_store.prefer_db_kv_blob_length + size_offset
+
+        blob = bytes(seeded_random.getrandbits(8) for _ in range(size))
+        await data_store.insert(key=blob, value=b"", store_id=store_id, status=Status.COMMITTED)
+        inserted_keys.add(blob)
 
-    expected_nodes = await get_all_nodes(data_store, store_id)
-    nodes = await data_store.get_nodes([node.hash for node in expected_nodes])
-    assert nodes == expected_nodes
+    retrieved_keys = set(await data_store.get_keys(store_id=store_id))
 
-    node_hash = bytes32.zeros
-    node_hash_2 = bytes32([0] * 31 + [1])
-    with pytest.raises(Exception, match=f"^Nodes not found for hashes: {node_hash.hex()}, {node_hash_2.hex()}"):
-        await data_store.get_nodes([node_hash, node_hash_2] + [node.hash for node in expected_nodes])
+    assert retrieved_keys == inserted_keys
 
 
 @pytest.mark.anyio
-@pytest.mark.parametrize("pre", [0, 2048])
-@pytest.mark.parametrize("batch_size", [25, 100, 500])
-async def test_get_leaf_at_minimum_height(
+async def test_get_keys_values_both_disk_and_db(
     data_store: DataStore,
     store_id: bytes32,
-    pre: int,
-    batch_size: int,
+    seeded_random: random.Random,
 ) -> None:
-    num_values = 1000
-    value_offset = 1000000
-    all_min_leafs: set[TerminalNode] = set()
+    inserted_keys_values: dict[bytes, bytes] = {}
+
+    for size_offset in [-1, 0, 1]:
+        size = data_store.prefer_db_kv_blob_length + size_offset
+
+        key = bytes(seeded_random.getrandbits(8) for _ in range(size))
+        value = bytes(seeded_random.getrandbits(8) for _ in range(size))
+        await data_store.insert(key=key, value=value, store_id=store_id, status=Status.COMMITTED)
+        inserted_keys_values[key] = value
+
+    terminal_nodes = await data_store.get_keys_values(store_id=store_id)
+    retrieved_keys_values = {node.key: node.value for node in terminal_nodes}
+
+    assert retrieved_keys_values == inserted_keys_values
+
 
-    if pre > 0:
-        # This builds a complete binary tree, in order to test more than one batch in the queue before finding the leaf
+@pytest.mark.anyio
+@boolean_datacases(name="success", false="invalid file", true="valid file")
+async def test_db_data_insert_from_file(
+    data_store: DataStore,
+    store_id: bytes32,
+    tmp_path: Path,
+    seeded_random: random.Random,
+    success: bool,
+) -> None:
+    num_keys = 1000
+    db_uri = generate_in_memory_db_uri()
+
+    async with DataStore.managed(
+        database=db_uri,
+        uri=True,
+        merkle_blobs_path=tmp_path.joinpath("merkle-blobs-tmp"),
+        key_value_blobs_path=tmp_path.joinpath("key-value-blobs-tmp"),
+    ) as tmp_data_store:
+        await tmp_data_store.create_tree(store_id, status=Status.COMMITTED)
         changelist: list[dict[str, Any]] = []
+        for _ in range(num_keys):
+            use_file = seeded_random.choice([True, False])
+            assert tmp_data_store.prefer_db_kv_blob_length > 7
+            size = tmp_data_store.prefer_db_kv_blob_length + 1 if use_file else 8
+            key = seeded_random.randbytes(size)
+            value = seeded_random.randbytes(size)
+            changelist.append({"action": "insert", "key": key, "value": value})
+
+        await tmp_data_store.insert_batch(store_id, changelist, status=Status.COMMITTED)
+        root = await tmp_data_store.get_tree_root(store_id)
+        files_path = tmp_path.joinpath("files")
+        await write_files_for_root(tmp_data_store, store_id, root, files_path, 1000)
+        assert root.node_hash is not None
+        filename = get_delta_filename_path(files_path, store_id, root.node_hash, 1)
+        assert filename.exists()
 
-        for value in range(pre):
-            value_bytes = (value * value).to_bytes(8, byteorder="big")
-            changelist.append({"action": "upsert", "key": value_bytes, "value": value_bytes})
-        await data_store.insert_batch(
-            store_id=store_id,
-            changelist=changelist,
-            status=Status.COMMITTED,
-        )
+    root_hash = bytes32([0] * 31 + [1]) if not success else root.node_hash
+    sinfo = ServerInfo("http://127.0.0.1/8003", 0, 0)
 
-    for value in range(num_values):
-        value_bytes = value.to_bytes(4, byteorder="big")
-        # Use autoinsert instead of `insert_batch` to get a more randomly shaped tree
-        await data_store.autoinsert(
-            key=value_bytes,
-            value=value_bytes,
-            store_id=store_id,
-            status=Status.COMMITTED,
-        )
+    if not success:
+        target_filename_path = get_delta_filename_path(files_path, store_id, root_hash, 1)
+        shutil.copyfile(filename, target_filename_path)
+        assert target_filename_path.exists()
 
-        if (value + 1) % batch_size == 0:
-            hash_to_parent: dict[bytes32, InternalNode] = {}
-            root = await data_store.get_tree_root(store_id)
-            assert root.node_hash is not None
-            min_leaf = await data_store.get_leaf_at_minimum_height(root.node_hash, hash_to_parent)
-            all_nodes = await get_all_nodes(data_store, store_id)
-            heights: dict[bytes32, int] = {}
-            heights[root.node_hash] = 0
-            min_leaf_height = None
-
-            for node in all_nodes:
-                if isinstance(node, InternalNode):
-                    heights[node.left_hash] = heights[node.hash] + 1
-                    heights[node.right_hash] = heights[node.hash] + 1
-                else:
-                    if min_leaf_height is not None:
-                        min_leaf_height = min(min_leaf_height, heights[node.hash])
-                    else:
-                        min_leaf_height = heights[node.hash]
-
-            assert min_leaf_height is not None
-            if pre > 0:
-                assert min_leaf_height >= 11
-            for node in all_nodes:
-                if isinstance(node, TerminalNode):
-                    assert node == min_leaf
-                    assert heights[min_leaf.hash] == min_leaf_height
-                    break
-                if node.left_hash in hash_to_parent:
-                    assert hash_to_parent[node.left_hash] == node
-                if node.right_hash in hash_to_parent:
-                    assert hash_to_parent[node.right_hash] == node
-
-            # Push down the min height leaf, so on the next iteration we get a different leaf
-            pushdown_height = 20
-            for repeat in range(pushdown_height):
-                value_bytes = (value + (repeat + 1) * value_offset).to_bytes(4, byteorder="big")
-                await data_store.insert(
-                    key=value_bytes,
-                    value=value_bytes,
-                    store_id=store_id,
-                    reference_node_hash=min_leaf.hash,
-                    side=Side.RIGHT,
-                    status=Status.COMMITTED,
-                )
-            assert min_leaf not in all_min_leafs
-            all_min_leafs.add(min_leaf)
+    keys_value_path = data_store.key_value_blobs_path.joinpath(store_id.hex())
+    assert sum(1 for path in keys_value_path.rglob("*") if path.is_file()) == 0
+
+    is_success = await insert_from_delta_file(
+        data_store=data_store,
+        store_id=store_id,
+        existing_generation=0,
+        target_generation=1,
+        root_hashes=[root_hash],
+        server_info=sinfo,
+        client_foldername=files_path,
+        timeout=aiohttp.ClientTimeout(total=15, sock_connect=5),
+        log=log,
+        proxy_url="",
+        downloader=None,
+    )
+    assert is_success == success
+
+    async with data_store.db_wrapper.reader() as reader:
+        async with reader.execute("SELECT COUNT(*) FROM ids") as cursor:
+            row_count = await cursor.fetchone()
+            assert row_count is not None
+            if success:
+                assert row_count[0] > 0
+            else:
+                assert row_count[0] == 0
+
+    if success:
+        assert sum(1 for path in keys_value_path.rglob("*") if path.is_file()) > 0
+    else:
+        assert sum(1 for path in keys_value_path.rglob("*") if path.is_file()) == 0
+
+
+@pytest.mark.anyio
+async def test_manage_kv_files(
+    data_store: DataStore,
+    store_id: bytes32,
+    seeded_random: random.Random,
+) -> None:
+    num_keys = 1000
+    num_files = 0
+    keys_value_path = data_store.key_value_blobs_path.joinpath(store_id.hex())
+
+    with pytest.raises(Exception, match="Test exception"):
+        async with data_store.db_wrapper.writer() as writer:
+            with data_store.manage_kv_files(store_id):
+                for _ in range(num_keys):
+                    use_file = seeded_random.choice([True, False])
+                    assert data_store.prefer_db_kv_blob_length > 7
+                    size = data_store.prefer_db_kv_blob_length + 1 if use_file else 8
+                    key = seeded_random.randbytes(size)
+                    value = seeded_random.randbytes(size)
+                    await data_store.add_key_value(key, value, store_id, writer)
+                    num_files += 2 * use_file
+
+                assert num_files > 0
+                assert sum(1 for path in keys_value_path.rglob("*") if path.is_file()) == num_files
+                raise Exception("Test exception")
+
+    assert sum(1 for path in keys_value_path.rglob("*") if path.is_file()) == 0