Change get_pending_root method.

This commit is contained in:
Florin Chirica 2022-07-08 01:18:05 +02:00
parent e591f654eb
commit f232e7c85c
No known key found for this signature in database
GPG Key ID: 1805593F7B529698
3 changed files with 30 additions and 18 deletions

View File

@ -189,11 +189,10 @@ class DataLayer:
if singleton_record is None:
return
if root is None:
pending_roots = await self.data_store.get_pending_roots(tree_id=tree_id)
if len(pending_roots) > 0:
root = pending_roots[0]
if root.generation == 0 and root.node_hash is None:
await self.data_store.change_root_status(root, Status.COMMITTED)
pending_root = await self.data_store.get_pending_root(tree_id=tree_id)
if pending_root is not None:
if pending_root.generation == 0 and pending_root.node_hash is None:
await self.data_store.change_root_status(pending_root, Status.COMMITTED)
await self.data_store.clear_pending_roots(tree_id=tree_id)
return
else:
@ -224,10 +223,13 @@ class DataLayer:
await self.data_store.shift_root_generations(tree_id=tree_id, shift_size=generation_shift)
else:
expected_root_hash = None if new_hashes[0] == self.none_bytes else new_hashes[0]
pending_roots = await self.data_store.get_pending_roots(tree_id=tree_id)
expected_root = next((root for root in pending_roots if root.node_hash == expected_root_hash), None)
if expected_root is not None and expected_root.generation == root.generation + 1:
await self.data_store.change_root_status(expected_root, Status.COMMITTED)
pending_root = await self.data_store.get_pending_root(tree_id=tree_id)
if (
pending_root is not None
and pending_root.generation == root.generation + 1
and pending_root.node_hash == expected_root_hash
):
await self.data_store.change_root_status(pending_root, Status.COMMITTED)
await self.data_store.build_ancestor_table_for_latest_root(tree_id=tree_id)
await self.data_store.clear_pending_roots(tree_id=tree_id)

View File

@ -283,14 +283,23 @@ class DataStore:
return node_hash
async def get_pending_roots(self, tree_id: bytes32, *, lock: bool = True) -> List[Root]:
async def get_pending_root(self, tree_id: bytes32, *, lock: bool = True) -> Optional[Root]:
async with self.db_wrapper.locked_transaction(lock=lock):
cursor = await self.db.execute(
"SELECT * FROM root WHERE tree_id == :tree_id AND status == :status",
{"tree_id": tree_id.hex(), "status": Status.PENDING.value},
)
return [Root.from_row(row=row) async for row in cursor]
row = await cursor.fetchone()
if row is None:
return None
maybe_extra_result = await cursor.fetchone()
if maybe_extra_result is not None:
raise Exception(f"multiple pending roots found for id: {tree_id.hex()}")
return Root.from_row(row=row)
async def clear_pending_roots(self, tree_id: bytes32, *, lock: bool = True) -> None:
async with self.db_wrapper.locked_transaction(lock=lock):

View File

@ -1046,8 +1046,9 @@ async def test_change_root_state(data_store: DataStore, tree_id: bytes32) -> Non
reference_node_hash=None,
side=None,
)
roots = await data_store.get_pending_roots(tree_id)
await data_store.change_root_status(roots[0], Status.COMMITTED)
root = await data_store.get_pending_root(tree_id)
assert root is not None
await data_store.change_root_status(root, Status.COMMITTED)
root = await data_store.get_tree_root(tree_id)
is_empty = await data_store.table_is_empty(tree_id=tree_id)
assert not is_empty
@ -1226,10 +1227,10 @@ async def test_pending_roots(data_store: DataStore, tree_id: bytes32) -> None:
tree_id=tree_id,
status=Status.PENDING,
)
pending_roots = await data_store.get_pending_roots(tree_id=tree_id)
assert len(pending_roots) == 1
assert pending_roots[0].generation == 2 and pending_roots[0].status == Status.PENDING
pending_root = await data_store.get_pending_root(tree_id=tree_id)
assert pending_root is not None
assert pending_root.generation == 2 and pending_root.status == Status.PENDING
await data_store.clear_pending_roots(tree_id=tree_id)
pending_roots = await data_store.get_pending_roots(tree_id=tree_id)
assert len(pending_roots) == 0
pending_root = await data_store.get_pending_root(tree_id=tree_id)
assert pending_root is None