from typing import List, Optional import aiosqlite from chia.util.db_wrapper import DBWrapper from chia.util.ints import uint32 from chia.wallet.util.wallet_types import WalletType from chia.wallet.wallet_info import WalletInfo class WalletUserStore: """ WalletUserStore keeps track of all user created wallets and necessary smart-contract data """ db_connection: aiosqlite.Connection cache_size: uint32 db_wrapper: DBWrapper @classmethod async def create(cls, db_wrapper: DBWrapper): self = cls() self.db_wrapper = db_wrapper self.db_connection = db_wrapper.db await self.db_connection.execute("pragma journal_mode=wal") await self.db_connection.execute("pragma synchronous=2") await self.db_connection.execute( ( "CREATE TABLE IF NOT EXISTS users_wallets(" "id INTEGER PRIMARY KEY AUTOINCREMENT," " name text," " wallet_type int," " data text)" ) ) await self.db_connection.execute("CREATE INDEX IF NOT EXISTS name on users_wallets(name)") await self.db_connection.execute("CREATE INDEX IF NOT EXISTS type on users_wallets(wallet_type)") await self.db_connection.execute("CREATE INDEX IF NOT EXISTS data on users_wallets(data)") await self.db_connection.commit() await self.init_wallet() return self async def init_wallet(self): all_wallets = await self.get_all_wallet_info_entries() if len(all_wallets) == 0: await self.create_wallet("Chia Wallet", WalletType.STANDARD_WALLET, "") async def _clear_database(self): cursor = await self.db_connection.execute("DELETE FROM users_wallets") await cursor.close() await self.db_connection.commit() async def create_wallet( self, name: str, wallet_type: int, data: str, id: Optional[int] = None ) -> Optional[WalletInfo]: async with self.db_wrapper.lock: cursor = await self.db_connection.execute( "INSERT INTO users_wallets VALUES(?, ?, ?, ?)", (id, name, wallet_type, data), ) await cursor.close() await self.db_connection.commit() return await self.get_last_wallet() async def delete_wallet(self, id: int): async with self.db_wrapper.lock: cursor = await self.db_connection.execute(f"DELETE FROM users_wallets where id={id}") await cursor.close() await self.db_connection.commit() async def update_wallet(self, wallet_info: WalletInfo, in_transaction): if not in_transaction: await self.db_wrapper.lock.acquire() try: cursor = await self.db_connection.execute( "INSERT or REPLACE INTO users_wallets VALUES(?, ?, ?, ?)", ( wallet_info.id, wallet_info.name, wallet_info.type, wallet_info.data, ), ) await cursor.close() finally: if not in_transaction: await self.db_connection.commit() self.db_wrapper.lock.release() async def get_last_wallet(self) -> Optional[WalletInfo]: cursor = await self.db_connection.execute("SELECT MAX(id) FROM users_wallets;") row = await cursor.fetchone() await cursor.close() if row is None: return None return await self.get_wallet_by_id(row[0]) async def get_all_wallet_info_entries(self) -> List[WalletInfo]: """ Return a set containing all wallets """ cursor = await self.db_connection.execute("SELECT * from users_wallets") rows = await cursor.fetchall() await cursor.close() result = [] for row in rows: result.append(WalletInfo(row[0], row[1], row[2], row[3])) return result async def get_wallet_by_id(self, id: int) -> Optional[WalletInfo]: """ Return a wallet by id """ cursor = await self.db_connection.execute("SELECT * from users_wallets WHERE id=?", (id,)) row = await cursor.fetchone() await cursor.close() if row is None: return None return WalletInfo(row[0], row[1], row[2], row[3])