Improve case handling.

chia/_tests/core/data_layer/test_data_store.py

@@ -1981,3 +1981,35 @@ async def test_insert_key_already_present(data_store: DataStore, tree_id: bytes3
     )
     with pytest.raises(Exception, match=f"Key already present: {key.hex()}"):
         await data_store.insert(key=key, value=value, tree_id=tree_id, reference_node_hash=None, side=None)
+
+
+@pytest.mark.anyio
+async def test_update_keys(data_store: DataStore, tree_id: bytes32) -> None:
+    num_keys = 10
+    missing_keys = 50
+    num_values = 10
+    new_keys = 10
+    for value in range(num_values):
+        changelist: List[Dict[str, Any]] = []
+        bytes_value = value.to_bytes(4, byteorder="big")
+        for key in range(num_keys + missing_keys):
+            bytes_key = key.to_bytes(4, byteorder="big")
+            changelist.append({"action": "delete", "key": bytes_key})
+        for key in range(num_keys):
+            bytes_key = key.to_bytes(4, byteorder="big")
+            changelist.append({"action": "insert", "key": bytes_key, "value": bytes_value})
+
+        await data_store.insert_batch(
+            tree_id=tree_id,
+            changelist=changelist,
+            status=Status.COMMITTED,
+        )
+        for key in range(num_keys):
+            bytes_key = key.to_bytes(4, byteorder="big")
+            node = await data_store.get_node_by_key(bytes_key, tree_id)
+            assert node.value == bytes_value
+        for key in range(num_keys, num_keys + missing_keys):
+            bytes_key = key.to_bytes(4, byteorder="big")
+            with pytest.raises(KeyNotFoundError, match=f"Key not found: {bytes_key.hex()}"):
+                await data_store.get_node_by_key(bytes_key, tree_id)
+        num_keys += new_keys

chia/data_layer/data_store.py

@@ -1353,10 +1353,13 @@ async def insert_batch(
             assert latest_local_root is not None
 
             key_hash_frequency: Dict[bytes32, int] = {}
+            first_action: Dict[bytes32, str] = {}
             for change in changelist:
                 key = change["key"]
                 hash = key_hash(key)
                 key_hash_frequency[hash] = key_hash_frequency.get(hash, 0) + 1
+                if hash not in first_action:
+                    first_action[hash] = change["action"]
 
             pending_autoinsert_hashes: List[bytes32] = []
             for change in changelist:
@@ -1370,10 +1373,15 @@ async def insert_batch(
                         # The key is not referenced in any other operation but this autoinsert, hence the order
                         # of performing these should not matter. We perform all these autoinserts as a batch
                         # at the end, to speed up the tree processing operations.
-                        if key_hash_frequency[hash] == 1 and enable_batch_autoinsert:
-                            terminal_node_hash = await self._insert_terminal_node(key, value)
-                            pending_autoinsert_hashes.append(terminal_node_hash)
-                            continue
+                        # Additionally, if the first action is a delete, we can still perform the autoinsert at the
+                        # end, since the order will be preserved.
+                        if enable_batch_autoinsert:
+                            if key_hash_frequency[hash] == 1 or (
+                                key_hash_frequency[hash] == 2 and first_action[hash] == "delete"
+                            ):
+                                terminal_node_hash = await self._insert_terminal_node(key, value)
+                                pending_autoinsert_hashes.append(terminal_node_hash)
+                                continue
                         insert_result = await self.autoinsert(
                             key, value, tree_id, True, Status.COMMITTED, root=latest_local_root
                         )