feat: calculate totalBalance and check skip data and type in gather i…

…nputs
nervosnetwork · Sep 3, 2019 · f455545 · f455545
1 parent 5746438
commit f455545
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 16 deletions.
diff --git a/packages/neuron-wallet/src/database/address/dao.ts b/packages/neuron-wallet/src/database/address/dao.ts
@@ -51,6 +51,8 @@ export default class AddressDao {
   // sentBalance means balance of OutputStatus.Sent cells (sent to me but not committed)
   // pendingBalance means balance of OutputStatus.Pending cells (sent from me, but not committed)
   // so the final balance is (liveBalance + sentBalance - pendingBalance)
+  // balance is the balance which skipped data and type
+  // totalBalance means balance with data and type
   public static updateTxCountAndBalance = async (address: string): Promise<AddressEntity[]> => {
     const addressEntities = await getConnection()
       .getRepository(AddressEntity)
@@ -67,9 +69,12 @@ export default class AddressDao {
         const addressEntity = entity
         addressEntity.txCount = txCount
         const lockHashes: string[] = await LockUtils.addressToAllLockHashes(addressEntity.address)
-        addressEntity.liveBalance = await CellsService.getBalance(lockHashes, OutputStatus.Live)
-        addressEntity.sentBalance = await CellsService.getBalance(lockHashes, OutputStatus.Sent)
-        addressEntity.pendingBalance = await CellsService.getBalance(lockHashes, OutputStatus.Pending)
+        addressEntity.liveBalance = await CellsService.getBalance(lockHashes, OutputStatus.Live, true)
+        addressEntity.sentBalance = await CellsService.getBalance(lockHashes, OutputStatus.Sent, true)
+        addressEntity.pendingBalance = await CellsService.getBalance(lockHashes, OutputStatus.Pending, true)
+        const totalLiveBalance = await CellsService.getBalance(lockHashes, OutputStatus.Live, false)
+        const totalSentBalance = await CellsService.getBalance(lockHashes, OutputStatus.Sent, false)
+        addressEntity.totalBalance = (BigInt(totalLiveBalance) - BigInt(totalSentBalance)).toString()
         return addressEntity
       })
     )

diff --git a/packages/neuron-wallet/src/services/cells.ts b/packages/neuron-wallet/src/services/cells.ts
@@ -3,6 +3,7 @@ import OutputEntity from 'database/chain/entities/output'
 import { Cell, OutPoint, Input } from 'types/cell-types'
 import { CapacityNotEnough } from 'exceptions'
 import { OutputStatus } from './tx/params'
+import SkipDataAndType from './skip-data-and-type'
 
 export const MIN_CELL_CAPACITY = '6100000000'
 
@@ -11,16 +12,27 @@ export const MIN_CELL_CAPACITY = '6100000000'
 /* eslint no-restricted-syntax: "warn" */
 export default class CellsService {
   // exclude hasData = true and typeScript != null
-  public static getBalance = async (lockHashes: string[], status: OutputStatus): Promise<string> => {
+  public static getBalance = async (
+    lockHashes: string[],
+    status: OutputStatus,
+    skipDataAndType: boolean
+  ): Promise<string> => {
+    const queryParams = {
+      lockHash: In(lockHashes),
+      status,
+    }
+
+    if (skipDataAndType) {
+      Object.assign(queryParams, {
+        hasData: false,
+        typeScript: null,
+      })
+    }
+
     const cells: OutputEntity[] = await getConnection()
       .getRepository(OutputEntity)
       .find({
-        where: {
-          lockHash: In(lockHashes),
-          hasData: false,
-          typeScript: null,
-          status,
-        },
+        where: queryParams,
       })
 
     const capacity: bigint = cells.map(c => BigInt(c.capacity)).reduce((result, c) => result + c, BigInt(0))
@@ -70,16 +82,23 @@ export default class CellsService {
       throw new Error(`capacity can't be less than ${MIN_CELL_CAPACITY}`)
     }
 
+    const queryParams = {
+      lockHashes: In(lockHashes),
+      status: OutputStatus.Live,
+    }
+    const skipDataAndType = SkipDataAndType.getInstance().get()
+    if (skipDataAndType) {
+      Object.assign(queryParams, {
+        hasData: false,
+        typeScript: null,
+      })
+    }
+
     // only live cells, skip which has data or type
     const cellEntities: OutputEntity[] = await getConnection()
       .getRepository(OutputEntity)
       .find({
-        where: {
-          lockHash: In(lockHashes),
-          status: 'live',
-          hasData: false,
-          typeScript: null,
-        },
+        where: queryParams,
       })
     cellEntities.sort((a, b) => {
       const result = BigInt(a.capacity) - BigInt(b.capacity)