Repository: Chia-Network/chia-blockchain
File: chia/data_layer/data_store.py (Python)

Code before:
                    raise Exception("Internal error") from e
                return kv_id

            raise

        if row is None:
            raise Exception("Internal error")
        if use_file:
            path = self.get_key_value_path(store_id=store_id, blob_hash=blob_hash)
            path.parent.mkdir(parents=True, exist_ok=True)
            # TODO: consider file-system based locking of either the file or the store directory
            path.write_bytes(zstd.compress(blob))
        return KeyOrValueId(row[0])

    def delete_unconfirmed_kvids(self, store_id: bytes32) -> None:
        for blob_hash in self.unconfirmed_keys_values[store_id]:
            with log_exceptions(log=log, consume=True):
                path = self.get_key_value_path(store_id=store_id, blob_hash=blob_hash)
                try:
                    path.unlink()
                except FileNotFoundError:
                    log.error(f"Cannot find key/value path {path} for hash {blob_hash}")
        self.unconfirmed_keys_values[store_id].clear()

    def confirm_all_kvids(self, store_id: bytes32) -> None:
        self.unconfirmed_keys_values[store_id].clear()

    @contextmanager
    def manage_kv_files(self, store_id: bytes32) -> Iterator[None]:
        if self.unconfirmed_keys_values[store_id]:
            raise Exception("Internal error: unconfirmed keys values cache not cleaned")

        try:
            yield
        except Exception:
            self.delete_unconfirmed_kvids(store_id)
            raise
        else:
            self.confirm_all_kvids(store_id)

    async def add_key_value(
        self, key: bytes, value: bytes, store_id: bytes32, writer: aiosqlite.Connection
    ) -> tuple[KeyId, ValueId]:
        kid = KeyId(await self.add_kvid(key, store_id, writer=writer))
        vid = ValueId(await self.add_kvid(value, store_id, writer=writer))

        return (kid, vid)

    async def get_terminal_node_by_hash(
        self,
        node_hash: bytes32,

Review comment (suggestion):
```suggestion
        del self.unconfirmed_keys_values[store_id]
```

Code after:

        if row is None:
            raise Exception("Internal error")
        if use_file:
            path = self.get_key_value_path(store_id=store_id, blob_hash=blob_hash)
            path.parent.mkdir(parents=True, exist_ok=True)
            self.unconfirmed_keys_values[store_id].append(blob_hash)
            # TODO: consider file-system based locking of either the file or the store directory
            path.write_bytes(zstd.compress(blob))
        return KeyOrValueId(row[0])

    def delete_unconfirmed_kvids(self, store_id: bytes32) -> None:
        for blob_hash in self.unconfirmed_keys_values[store_id]:
            with log_exceptions(log=log, consume=True):
                path = self.get_key_value_path(store_id=store_id, blob_hash=blob_hash)
                try:
                    path.unlink()
                except FileNotFoundError:
                    log.error(f"Cannot find key/value path {path} for hash {blob_hash}")
        del self.unconfirmed_keys_values[store_id]

    def confirm_all_kvids(self, store_id: bytes32) -> None:
        del self.unconfirmed_keys_values[store_id]

    @contextmanager
    def manage_kv_files(self, store_id: bytes32) -> Iterator[None]:
        if store_id not in self.unconfirmed_keys_values:
            self.unconfirmed_keys_values[store_id] = []
        else:
            raise Exception("Internal error: unconfirmed keys values cache not cleaned")

        try:
            yield
        except:
            self.delete_unconfirmed_kvids(store_id)
            raise
        else:
            self.confirm_all_kvids(store_id)

    async def add_key_value(
        self, key: bytes, value: bytes, store_id: bytes32, writer: aiosqlite.Connection
    ) -> tuple[KeyId, ValueId]:
        kid = KeyId(await self.add_kvid(key, store_id, writer=writer))
        vid = ValueId(await self.add_kvid(value, store_id, writer=writer))

        return (kid, vid)

    async def get_terminal_node_by_hash(
        self,
        node_hash: bytes32,
        store_id: bytes32,