chia-blockchain/chia/wallet/wallet_retry_store.py
Matt Hauff e5f7232098
Search for launcher with correct spent_height (#13827)
* Search for launcher with correct spent_height

* Remove unecessary assert

* Add a PRIMARY KEY to the retry store
2022-11-08 00:04:18 -06:00

58 lines
2.0 KiB
Python

from __future__ import annotations
from typing import List, Optional, Tuple
from chia_rs import CoinState
from chia.types.blockchain_format.sized_bytes import bytes32
from chia.util.db_wrapper import DBWrapper2
from chia.util.ints import uint32
class WalletRetryStore:
"""
Persistent coin states that we have received but failed to add
"""
db_wrapper: DBWrapper2
@classmethod
async def create(cls, db_wrapper: DBWrapper2) -> "WalletRetryStore":
self = cls()
self.db_wrapper = db_wrapper
async with self.db_wrapper.writer_maybe_transaction() as conn:
await conn.execute(
"CREATE TABLE IF NOT EXISTS retry_store("
" coin_state blob PRIMARY KEY,"
" peer blob,"
" fork_height int)"
)
return self
async def get_all_states_to_retry(self) -> List[Tuple[CoinState, bytes32, uint32]]:
"""
Return all states that were failed to sync
"""
async with self.db_wrapper.reader_no_transaction() as conn:
rows = await conn.execute_fetchall("SELECT * from retry_store")
return [(CoinState.from_bytes(row[0]), bytes32(row[1]), uint32(row[2])) for row in rows]
async def add_state(self, state: CoinState, peer_id: bytes32, fork_height: Optional[uint32]) -> None:
"""
Adds object to key val store. Obj MUST support __bytes__ and bytes() methods.
"""
async with self.db_wrapper.writer_maybe_transaction() as conn:
cursor = await conn.execute(
"INSERT OR IGNORE INTO retry_store VALUES(?, ?, ?)",
(bytes(state), peer_id, 0 if fork_height is None else fork_height),
)
await cursor.close()
async def remove_state(self, state: CoinState) -> None:
async with self.db_wrapper.writer_maybe_transaction() as conn:
cursor = await conn.execute("DELETE FROM retry_store where coin_state=?", (bytes(state),))
await cursor.close()