From 2afff1b41fcd6c8ecc76d31122ee4d3155cf6973 Mon Sep 17 00:00:00 2001
From: Alex D <pr0fedt@gmail.com>
Date: Mon, 1 Mar 2021 10:42:07 +0300
Subject: [PATCH] Use derive customAccount and customLocks for balance
 overrides (#3248)

* modified api-derive balances to allow use of custom balance storages via derives; current implementation only works with standard call signatures

* type predicate as filter value

* renamed derive method name
---
 packages/api-derive/src/balances/account.ts | 12 +++++++++++-
 packages/api-derive/src/balances/all.ts     | 14 ++++++++++++--
 2 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/packages/api-derive/src/balances/account.ts b/packages/api-derive/src/balances/account.ts
index 47a06b123d..766968d6ee 100644
--- a/packages/api-derive/src/balances/account.ts
+++ b/packages/api-derive/src/balances/account.ts
@@ -54,8 +54,18 @@ function queryBalancesFree (api: ApiInterfaceRx, accountId: AccountId): Observab
   );
 }
 
+type DeriveCustomAccount = ApiInterfaceRx['derive'] & { [custom: string]: {
+  customAccount?: ApiInterfaceRx['query']['balances']['account']
+} }
+
 function queryBalancesAccount (api: ApiInterfaceRx, accountId: AccountId, modules: string[] = ['balances']): Observable<Result> {
-  const balances = modules.map((m): QueryableStorageMultiArg<'rxjs'> => [api.query[m].account, accountId]);
+  const balances = modules.map(
+    (m): QueryableStorageMultiArg<'rxjs'> => [
+      (api.derive as DeriveCustomAccount)[m]?.customAccount ?? api.query[m].account,
+      accountId
+    ]
+  );
+
   const extract = (data: AccountData[]) =>
     data.map(({ feeFrozen, free, miscFrozen, reserved }): BalanceResult => [free, reserved, feeFrozen, miscFrozen]);
 
diff --git a/packages/api-derive/src/balances/all.ts b/packages/api-derive/src/balances/all.ts
index 7c1629188d..ae3aa9f4a6 100644
--- a/packages/api-derive/src/balances/all.ts
+++ b/packages/api-derive/src/balances/all.ts
@@ -112,11 +112,21 @@ function queryOld (api: ApiInterfaceRx, accountId: AccountId): Observable<Result
   );
 }
 
+const isNonNullable = <T>(nullable: T): nullable is NonNullable<T> => !!nullable;
+
+type DeriveCustomLocks = ApiInterfaceRx['derive'] & { [custom: string]: {
+  customLocks?: ApiInterfaceRx['query']['balances']['locks']
+} }
+
 // current (balances, vesting)
 function queryCurrent (api: ApiInterfaceRx, accountId: AccountId, balanceInstances: string[] = ['balances']): Observable<ResultBalance> {
-  const lockCalls = balanceInstances.map((m) => api.query[m].locks);
+  const lockCalls = balanceInstances.map(
+    (m): ApiInterfaceRx['query']['balances']['locks'] | undefined =>
+      (api.derive as DeriveCustomLocks)[m]?.customLocks ?? api.query[m as 'balances']?.locks
+  );
+
   const lockEmpty = lockCalls.map((c) => !c);
-  const lockQueries = lockCalls.filter((c) => c).map((c): QueryableStorageMultiArg<'rxjs'> => [c, accountId]);
+  const lockQueries = lockCalls.filter(isNonNullable).map((c): QueryableStorageMultiArg<'rxjs'> => [c, accountId]);
 
   return (
     api.query.vesting?.vesting