Support indefinitely large calls to WalletCoinStore.get_multiple_coin_records()

chia/wallet/wallet_coin_store.py

@@ -3,9 +3,12 @@
 import sqlite3
 from typing import Dict, List, Optional, Set
 
+from aiosqlite import Cursor
+
 from chia.types.blockchain_format.coin import Coin
 from chia.types.blockchain_format.sized_bytes import bytes32
-from chia.util.db_wrapper import DBWrapper2, execute_fetchone
+from chia.util.chunks import chunks
+from chia.util.db_wrapper import SQLITE_MAX_VARIABLE_NUMBER, DBWrapper2, execute_fetchone
 from chia.util.ints import uint32, uint64
 from chia.wallet.util.wallet_types import WalletType
 from chia.wallet.wallet_coin_record import WalletCoinRecord
@@ -64,18 +67,30 @@ async def count_small_unspent(self, cutoff: int) -> int:
             )
             return int(0 if row is None else row[0])
 
-    async def get_multiple_coin_records(self, coin_names: List[bytes32]) -> List[WalletCoinRecord]:
+    async def get_multiple_coin_records(
+        self, coin_names: List[bytes32], maximum_query_size: int = SQLITE_MAX_VARIABLE_NUMBER
+    ) -> List[WalletCoinRecord]:
         """Return WalletCoinRecord(s) that have a coin name in the specified list"""
         if len(coin_names) == 0:
             return []
 
-        as_hexes = [cn.hex() for cn in coin_names]
+        coins: List[WalletCoinRecord] = []
+
         async with self.db_wrapper.reader_no_transaction() as conn:
-            rows = await conn.execute_fetchall(
-                f'SELECT * from coin_record WHERE coin_name in ({"?," * (len(as_hexes) - 1)}?)', tuple(as_hexes)
-            )
+            cursors: List[Cursor] = []
+            for names_chunk in chunks(coin_names, maximum_query_size):
+                as_hexes = tuple(cn.hex() for cn in names_chunk)
+                cursors.append(
+                    await conn.execute(
+                        f'SELECT * from coin_record WHERE coin_name in ({"?," * (len(as_hexes) - 1)}?)', tuple(as_hexes)
+                    )
+                )
 
-        return [self.coin_record_from_row(row) for row in rows]
+            for cursor in cursors:
+                for row in await cursor.fetchall():
+                    coins.append(self.coin_record_from_row(row))
+
+        return coins
 
     # Store CoinRecord in DB and ram cache
     async def add_coin_record(self, record: WalletCoinRecord, name: Optional[bytes32] = None) -> None:

tests/wallet/test_wallet_coin_store.py

@@ -267,6 +267,37 @@ async def test_get_multiple_coin_records() -> None:
         ) == set([record_1, record_2, record_3, record_4, record_5, record_6, record_7])
 
 
+@pytest.mark.asyncio
+async def test_get_multiple_coin_records_chunking() -> None:
+    # faster testing of chunking than using an actual sql limit that we can't control
+    maximum_query_size = 50
+    extra_count = 10
+    unrequested_count = 5
+    # three chunks worth
+    requested_coin_count = (2 * maximum_query_size) + extra_count
+    total_coin_count = requested_coin_count + unrequested_count
+    coins = [Coin(token_bytes(32), token_bytes(32), uint64(12312)) for _ in range(total_coin_count)]
+    records = [
+        WalletCoinRecord(coin, uint32(4), uint32(0), False, True, WalletType.STANDARD_WALLET, 0) for coin in coins
+    ]
+
+    async with DBConnection(1) as db_wrapper:
+        store = await WalletCoinStore.create(db_wrapper)
+
+        for record in records:
+            await store.add_coin_record(record)
+
+        coins_to_query = coins[:requested_coin_count]
+        names_to_query = [coin.name() for coin in coins_to_query]
+        expected_records = set(records[:requested_coin_count])
+        returned_records = await store.get_multiple_coin_records(
+            coin_names=names_to_query,
+            maximum_query_size=maximum_query_size,
+        )
+
+    assert set(returned_records) == expected_records
+
+
 @pytest.mark.asyncio
 async def test_delete_coin_record() -> None:
     async with DBConnection(1) as db_wrapper: