Prefetch storage (#764)

* wip: fetch-storage

* wip: fetch storages

* fix cli

* refactor

* refactor and fix

* worker

* fix

README.md

@@ -154,6 +154,7 @@ There are 2 types of plugins: `cli` and `rpc`. `cli` plugins are used to extend
 
 To create a new plugin, you could check out the [run-block plugin](packages/chopsticks/src/plugins/run-block/) as an example.
 
+
 ## RPC Methods
 
 Chopsticks allows you to load your extended rpc methods by adding the cli argument `--unsafe-rpc-methods=<file path>`or `-ur=<file path>`.
@@ -162,11 +163,11 @@ Chopsticks allows you to load your extended rpc methods by adding the cli argume
 
 It loads an **unverified** scripts, making it **unsafe**. Ensure you load a **trusted** script.
 
-**example**: 
+**example**:
 
 `npx @acala-network/chopsticks@latest --unsafe-rpc-methods=rpc-methods-scripts.js`
 
-**scripts example of rpc-methods-scripts:** 
+**scripts example of rpc-methods-scripts:**
 
 ```
 return {
@@ -181,3 +182,42 @@ return {
 }
 ```
 
+## Testing big migrations
+
+When testing migrations with lots of keys, you may want to fetch and cache some storages.
+
+There are two ways to fetch storages.
+
+The first way is to use a config file with a `prefetch-storages` section:
+
+```yml
+prefetch-storages:
+  - '0x123456' # fetch all storages with this prefix
+  - Balances # fetch all storages under Balances pallet
+  - Tokens.Accounts # fetch all storages under Tokens.Accounts stroage
+  - System: Account # fetch all storages under System.Account stroage
+  - Tokens:
+      Accounts: [5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY] # fetch all storages for Tokens.Accounts(Alice)
+  - Tokens.Accounts: [5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY, { token: DOT }] # fetch this particular storage
+```
+
+When you starts chopsticks, it will fetch these storages in background.
+
+Please note that only the formats mentioned above are supported for config files.
+
+The second way is use `fetch-storages` subcommand to only fetch and cache storages:
+
+```sh
+npx @acala-network/chopsticks@latest fetch-storages 0x123456 Balances Tokens.Accounts
+	--endlpint=wss://acala-rpc-0.aca-api.network
+	--block=<blockhash> # default to latest block
+	--db=acala.sqlite
+```
+
+The subcommand arguments could be:
+- hex: fetch all storages with this prefix
+- PalletName: fetch all storages for this pallet
+- PalletName.StorageName: fetch all storages for this storage
+
+Please note that for both ways, fetched storages will be saved in the sqlite file specified by `--db` option (`db: ./acala.sqlite` in a config file), if not provided, it will default to `./db-{network}-{block}.sqlite`.
+

packages/chopsticks/src/context.ts

@@ -4,6 +4,7 @@ import { Config } from './schema/index.js'
 import { HexString } from '@polkadot/util/types'
 import { SqliteDatabase } from '@acala-network/chopsticks-db'
 import { overrideStorage, overrideWasm } from './utils/override.js'
+import { startFetchStorageWorker } from './utils/fetch-storages.js'
 import axios from 'axios'
 
 const logger = defaultLogger.child({ name: 'setup-context' })
@@ -88,5 +89,12 @@ export const setupContext = async (argv: Config, overrideParent = false) => {
   await overrideWasm(chain, argv['wasm-override'], at)
   await overrideStorage(chain, argv['import-storage'], at)
 
-  return { chain }
+  const fetchStorageWorker = await startFetchStorageWorker({
+    config: argv['prefetch-storages'],
+    dbPath: argv.db,
+    block: argv.block,
+    endpoint: argv.endpoint,
+  })
+
+  return { chain, fetchStorageWorker }
 }

packages/chopsticks/src/plugins/fetch-storage/cli.ts

@@ -0,0 +1,33 @@
+import { z } from 'zod'
+import _ from 'lodash'
+import type { Argv } from 'yargs'
+
+import { configSchema, getYargsOptions } from '../../schema/index.js'
+import { fetchStorages } from '../../utils/fetch-storages.js'
+
+const schema = z.object(_.pick(configSchema.shape, ['endpoint', 'block', 'db']))
+
+export const cli = (y: Argv) => {
+  y.command({
+    command: 'fetch-storages [items..]',
+    aliases: ['fetch-storage'],
+    describe: 'Fetch and save storages',
+    builder: (yargs) => yargs.options(getYargsOptions(schema.shape)),
+    handler: async (argv) => {
+      const config = schema.parse(argv)
+      if (!argv.items) throw new Error('fetch-storages items are required')
+
+      try {
+        await fetchStorages({
+          block: config.block,
+          endpoint: config.endpoint,
+          dbPath: config.db,
+          config: argv.items as any,
+        })
+        process.exit(0)
+      } catch (e) {
+        process.exit(1)
+      }
+    },
+  })
+}

packages/chopsticks/src/plugins/fetch-storage/index.ts

@@ -0,0 +1 @@
+export * from './cli.js'

packages/chopsticks/src/schema/index.ts

@@ -19,7 +19,7 @@ export const configSchema = z.object({
         z.string(),
         z
           .number()
-          .max(Number.MAX_SAFE_INTEGER, 'Number is too big, please make it a string if you are uing a hex string'),
+          .max(Number.MAX_SAFE_INTEGER, 'Number is too big, please make it a string if you are using a hex string'),
         z.null(),
       ],
       {
@@ -65,6 +65,12 @@ export const configSchema = z.object({
         'Produce extra block when queued messages are detected. Default to true. Set to false to disable it.',
     })
     .optional(),
+  'prefetch-storages': z
+    .any({
+      description:
+        'Storage key prefixes config for fetching storage, useful for testing big migrations, see README for examples',
+    })
+    .optional(),
 })
 
 export type Config = z.infer<typeof configSchema>

packages/chopsticks/src/schema/options.test.ts

@@ -74,6 +74,12 @@ it('get yargs options from zod schema', () => {
         "description": "Port to listen on",
         "type": "number",
       },
+      "prefetch-storages": {
+        "choices": undefined,
+        "demandOption": false,
+        "description": "Storage key prefixes config for fetching storage, useful for testing big migrations, see README for examples",
+        "type": undefined,
+      },
       "process-queued-messages": {
         "choices": undefined,
         "demandOption": false,

packages/chopsticks/src/schema/parse.test.ts

@@ -64,6 +64,6 @@ describe('Parsed options', () => {
         // eslint-disable-next-line @typescript-eslint/no-loss-of-precision
         block: 0xb10f03bbc183da4d26e27528d28f6a73ddaf182fb6400ca363b77d2411ea5b0c,
       }),
-    ).toThrowError(/you are uing a hex string/)
+    ).toThrowError(/you are using a hex string/)
   })
 })

packages/chopsticks/src/setup-with-server.ts

@@ -16,6 +16,7 @@ export const setupWithServer = async (argv: Config) => {
     listenPort,
     async close() {
       await context.chain.close()
+      await context.fetchStorageWorker?.terminate()
       await close()
     },
   }

packages/chopsticks/src/utils/fetch-storages-worker.js

@@ -0,0 +1,13 @@
+import * as Comlink from 'comlink'
+import { parentPort } from 'node:worker_threads'
+import nodeEndpoint from 'comlink/dist/umd/node-adapter.js'
+
+import { fetchStorages } from './fetch-storages.js'
+
+const api = {
+  startFetch: async ({ ...options }) => {
+    await fetchStorages({ ...options })
+  },
+}
+
+Comlink.expose(api, nodeEndpoint(parentPort))

packages/chopsticks/src/utils/fetch-storages.test.ts

@@ -0,0 +1,140 @@
+import { ApiPromise } from '@polkadot/api'
+import { HexString } from '@polkadot/util/types'
+import { Like } from 'typeorm'
+import { ProviderInterface } from '@polkadot/rpc-provider/types'
+import { SqliteDatabase } from '@acala-network/chopsticks-db'
+import { WsProvider } from '@polkadot/rpc-provider'
+import { afterAll, beforeAll, describe, expect, it } from 'vitest'
+import { resolve } from 'node:path'
+import { tmpdir } from 'node:os'
+import { xxhashAsHex } from '@polkadot/util-crypto'
+
+import { FetchStorageConfig, fetchStorages, getPrefixesFromConfig } from './fetch-storages.js'
+
+describe('fetch-storages', () => {
+  let api: ApiPromise
+  let provider: ProviderInterface
+  const endpoint = 'wss://acala-rpc.aca-api.network'
+
+  beforeAll(async () => {
+    provider = new WsProvider(endpoint, 30_000)
+    api = new ApiPromise({ provider })
+    await api.isReady
+  })
+
+  afterAll(async () => {
+    await api.disconnect()
+  })
+
+  it('get prefixes from config works', async () => {
+    const config: FetchStorageConfig = [
+      '0x123456',
+      'Balances',
+      'Tokens.Accounts',
+      {
+        System: 'Account',
+      },
+      {
+        Tokens: {
+          Accounts: ['5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY'],
+        },
+      },
+      {
+        Tokens: {
+          Accounts: [
+            '5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY',
+            {
+              token: 'DOT',
+            },
+          ],
+        },
+      },
+      {
+        'Tokens.Accounts': ['5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY'],
+      },
+      {
+        'Tokens.Accounts': [
+          '5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY',
+          {
+            token: 'DOT',
+          },
+        ],
+      },
+    ]
+
+    const prefixes = await getPrefixesFromConfig(config, api)
+
+    expect(prefixes).toEqual([
+      '0x123456',
+      '0xc2261276cc9d1f8598ea4b6a74b15c2f',
+      '0x99971b5749ac43e0235e41b0d37869188ee7418a6531173d60d1f6a82d8f4d51',
+      '0x26aa394eea5630e07c48ae0c9558cef7b99d880ec681799c0cf30e8886371da9',
+      '0x99971b5749ac43e0235e41b0d37869188ee7418a6531173d60d1f6a82d8f4d51de1e86a9a8c739864cf3cc5ec2bea59fd43593c715fdd31c61141abd04a99fd6822c8558854ccde39a5684e7a56da27d',
+      '0x99971b5749ac43e0235e41b0d37869188ee7418a6531173d60d1f6a82d8f4d51de1e86a9a8c739864cf3cc5ec2bea59fd43593c715fdd31c61141abd04a99fd6822c8558854ccde39a5684e7a56da27dc483de2de1246ea70002',
+      '0x99971b5749ac43e0235e41b0d37869188ee7418a6531173d60d1f6a82d8f4d51de1e86a9a8c739864cf3cc5ec2bea59fd43593c715fdd31c61141abd04a99fd6822c8558854ccde39a5684e7a56da27d',
+      '0x99971b5749ac43e0235e41b0d37869188ee7418a6531173d60d1f6a82d8f4d51de1e86a9a8c739864cf3cc5ec2bea59fd43593c715fdd31c61141abd04a99fd6822c8558854ccde39a5684e7a56da27dc483de2de1246ea70002',
+    ])
+  })
+
+  it('get prefixes from config throws', async () => {
+    expect(() => getPrefixesFromConfig(['Balancess'], api)).rejects.toThrow(/Cannot find pallet Balancess/)
+
+    expect(() => getPrefixesFromConfig(['System.Acount'], api)).rejects.toThrow(
+      /Cannot find storage Acount in pallet System/,
+    )
+
+    expect(() =>
+      getPrefixesFromConfig(
+        [
+          {
+            System: ['Account', 'BlockHash'],
+          },
+        ],
+        api,
+      ),
+    ).rejects.toThrow(/Unsupported fetch-storage config: System.Account,BlockHash/)
+  })
+
+  it('fetch prefixes works', async () => {
+    const blockHash = '0x3a9a2d71537ceedff1a3895d68456f4a870bb89ab649fd47c6cf9c4f9731d580' // 4,500,000
+    const dbPath = resolve(tmpdir(), 'fetch.db.sqlite')
+
+    await fetchStorages({
+      block: blockHash,
+      endpoint,
+      dbPath,
+      config: [
+        'System.Number',
+        'Tips',
+        {
+          Rewards: {
+            PoolInfos: [{ Loans: { Token: 'ACA' } }],
+          },
+        },
+      ],
+    })
+
+    const db = new SqliteDatabase(dbPath)
+
+    const systemNumberStorage = await db.queryStorage(
+      blockHash,
+      (xxhashAsHex('System', 128) + xxhashAsHex('Number', 128).slice(2)) as HexString,
+    )
+    expect(systemNumberStorage?.value).toEqual('0x20aa4400')
+
+    const datasource = await db.datasource
+    const keyValueTable = datasource.getRepository('KeyValuePair')
+
+    expect(await keyValueTable.count()).toEqual(5)
+
+    expect(await keyValueTable.countBy({ key: Like(`${xxhashAsHex('Tips', 128)}%`) })).toEqual(3)
+
+    const rewards = await keyValueTable.findBy({ key: Like(`${xxhashAsHex('Rewards', 128)}%`) })
+    expect(rewards.length).toEqual(1)
+    expect(rewards[0].value).toEqual(
+      '0xf45ce8eb6fcaa12109000000000000000800002333cc48e197963c000000000000000014e1339c9e79e7380000000000000000010000000195319d9b71330d010000000000000000bb59ad064bb0bd000000000000000000',
+    )
+
+    db.close()
+  })
+})