@ocap/statedb-fs 1.20.16 → 1.21.1

package/lib/db.js

@@ -4,6 +4,7 @@ const Table = require('./table/base');
 const Account = require('./table/account');
 const Token = require('./table/token');
 const Rollup = require('./table/rollup');
+const Balance = require('./table/balance');
 
 const { name, version } = require('../package.json');
 
@@ -14,17 +15,36 @@ class FsStateDB extends StateDB {
     this.name = name;
     this.version = version;
 
-    this.account = new Account('account', dataDir, 'address');
-    this.asset = new Table('asset', dataDir, 'address');
-    this.factory = new Table('factory', dataDir, 'address');
-    this.delegation = new Table('delegation', dataDir, 'address');
-    this.tx = new Table('tx', dataDir, 'hash');
-    this.token = new Token('token', dataDir, 'address');
-    this.chain = new Table('chain', dataDir, 'address');
-    this.stake = new Table('stake', dataDir, 'address');
-    this.rollup = new Rollup('rollup', dataDir, 'address');
-    this.rollupBlock = new Table('rollupBlock', dataDir, 'hash');
-    this.evidence = new Table('evidence', dataDir, 'hash');
+    this.balance = new Balance({ name: 'balance', dataDir, uniqIndex: ['address', 'tokenAddress'] });
+    this.account = new Account({
+      name: 'account',
+      dataDir,
+      uniqIndex: 'address',
+      balanceTable: this.balance,
+      syncBalance: true,
+    });
+    this.factory = new Table({
+      name: 'factory',
+      dataDir,
+      uniqIndex: 'address',
+      syncBalance: true,
+      balanceTable: this.balance,
+    });
+    this.stake = new Table({
+      name: 'stake',
+      dataDir,
+      uniqIndex: 'address',
+      syncBalance: true,
+      balanceTable: this.balance,
+    });
+    this.asset = new Table({ name: 'asset', dataDir, uniqIndex: 'address' });
+    this.delegation = new Table({ name: 'delegation', dataDir, uniqIndex: 'address' });
+    this.tx = new Table({ name: 'tx', dataDir, uniqIndex: 'hash' });
+    this.token = new Token({ name: 'token', dataDir, uniqIndex: 'address' });
+    this.chain = new Table({ name: 'chain', dataDir, uniqIndex: 'address' });
+    this.rollup = new Rollup({ name: 'rollup', dataDir, uniqIndex: 'address' });
+    this.rollupBlock = new Table({ name: 'rollupBlock', dataDir, uniqIndex: 'hash' });
+    this.evidence = new Table({ name: 'evidence', dataDir, uniqIndex: 'hash' });
 
     this.attachReadyListeners();
   }

package/lib/table/account.js

@@ -1,29 +1,19 @@
 const { ensureChecksumAddress } = require('@ocap/state/lib/states/account');
-const debug = require('debug')(require('../../package.json').name);
 const FsTable = require('./base');
 
 class AccountTable extends FsTable {
   async _get(address, { traceMigration = true } = {}) {
-    const current = await this.collection.by(this.uniqIndex, ensureChecksumAddress(address));
+    const current = await super._get(ensureChecksumAddress(address));
     if (current && traceMigration && Array.isArray(current.migratedTo) && current.migratedTo.length) {
-      return this.collection.by(this.uniqIndex, current.migratedTo[0]);
+      return this._get(current.migratedTo[0]);
     }
 
     return current;
   }
 
-  async _create(key, attrs) {
-    const doc = await this._get(key);
-    if (doc) {
-      throw new Error(`${this.name} already exists: ${key}`);
-    }
-
-    debug(`insert ${this.name}`, attrs);
-    const newDoc = { [this.uniqIndex]: key, ...attrs };
-    newDoc[this.uniqIndex] = ensureChecksumAddress(newDoc[this.uniqIndex]);
-    const result = await this.collection.insert(newDoc);
-
-    return result;
+  _create(key, attrs = {}, ctx = {}) {
+    const address = ensureChecksumAddress(key);
+    return super._create(address, { ...attrs, address }, ctx);
   }
 }
 

package/lib/table/balance.js

@@ -0,0 +1,41 @@
+const FsTable = require('./base');
+
+class BalanceTable extends FsTable {
+  async getBalance(address) {
+    const tokens = await this.collection.find({ address });
+
+    return (tokens || []).reduce((acc, token) => {
+      acc[token.tokenAddress] = token.balance;
+      return acc;
+    }, {});
+  }
+
+  async updateBalance({ address, tokens, context = {} }, ctx) {
+    if (!Object.keys(tokens).length) return {};
+
+    const tokenBalances = await this.getBalance(address);
+    const updatedTokens = Object.keys(tokens)
+      .filter(
+        (token) => (tokenBalances[token] && tokenBalances[token] !== '0') || (tokens[token] && tokens[token] !== '0')
+      )
+      .filter((token) => tokenBalances[token] !== tokens[token]);
+
+    await Promise.all(
+      updatedTokens.map(async (token) => {
+        const key = { address, tokenAddress: token };
+
+        if (tokenBalances[token]) {
+          await this.update(key, { ...key, balance: tokens[token], context }, ctx);
+        } else {
+          await this.create(key, { ...key, balance: tokens[token], context }, ctx);
+        }
+
+        tokenBalances[token] = tokens[token];
+      })
+    );
+
+    return tokens;
+  }
+}
+
+module.exports = BalanceTable;

package/lib/table/base.js

@@ -8,13 +8,27 @@ const FsAdapter = require('lokijs/src/loki-fs-structured-adapter');
 const debug = require('debug')(require('../../package.json').name);
 
 class FsTable extends StateDBTable {
-  constructor(name, dataDir, uniqIndex) {
-    super();
+  /**
+   * @param {Object} param
+   * @param {string} param.name db name
+   * @param {string|string[]} param.uniqIndex primary keys
+   * @param {string} param.dataDir data directory for db
+   * @param {boolean} param.syncBalance sync balance table when create / update / get
+   * @param {BalanceTable} param.balanceTable balance table
+   */
+  constructor({ name, dataDir, uniqIndex, balanceTable, syncBalance = false }) {
+    super(uniqIndex);
 
     this.name = name;
     this.dataDir = dataDir;
     this.uniqIndex = uniqIndex;
     this.collection = null;
+    this.balanceTable = balanceTable;
+    this.syncBalance = syncBalance;
+
+    if (this.syncBalance && !this.balanceTable) {
+      throw new Error('balanceTable is required when syncBalance is true');
+    }
 
     const adapter = new FsAdapter();
     const db = new Lokijs(path.join(dataDir, `${name}.db`), {
@@ -26,7 +40,7 @@ class FsTable extends StateDBTable {
         this.collection = db.getCollection(name);
 
         if (this.collection === null) {
-          this.collection = db.addCollection(name, { unique: [uniqIndex], clone: true });
+          this.collection = db.addCollection(name, { unique: [this.primaryKey], clone: true });
         }
 
         this.markReady();
@@ -35,18 +49,48 @@ class FsTable extends StateDBTable {
   }
 
   async _create(key, attrs) {
-    const doc = await this._get(key);
+    const id = this.generatePrimaryKey(key);
+    // Don't call this._get here cause _get may be overwritten
+    const doc = await FsTable.prototype._get.call(this, id);
     if (doc) {
       throw new Error(`${this.name} already exists: ${key}`);
     }
+
     debug(`insert ${this.name}`, attrs);
-    const result = await this.collection.insert({ [this.uniqIndex]: key, ...attrs });
+
+    const insertAttrs = { ...attrs };
+    // insert without tokens cause we've migrate token to balance table
+    if (this.syncBalance) {
+      delete insertAttrs.tokens;
+    }
+
+    const result = await this.collection.insert({ [this.primaryKey]: id, ...insertAttrs });
+
+    if (this.syncBalance && attrs.tokens) {
+      debug(`update balance for ${this.name}`, attrs);
+
+      result.tokens = await this.balanceTable.updateBalance({
+        address: attrs.address,
+        tokens: attrs.tokens,
+        context: attrs.context,
+      });
+    }
 
     return result;
   }
 
-  _get(key) {
-    return key ? this.collection.by(this.uniqIndex, key) : null;
+  async _get(key) {
+    if (!key) return null;
+
+    const id = this.generatePrimaryKey(key);
+    const result = await this.collection.by(this.primaryKey, id);
+
+    if (result && this.syncBalance) {
+      const balance = await this.balanceTable.getBalance(result.address);
+      result.tokens = { ...(result.tokens || {}), ...balance };
+    }
+
+    return result;
   }
 
   _history() {
@@ -54,28 +98,49 @@ class FsTable extends StateDBTable {
   }
 
   async _update(key, updates) {
-    const doc = await this.collection.by(this.uniqIndex, key);
+    const id = this.generatePrimaryKey(key);
+    // Don't call this._get here cause _get may be overwritten
+    const doc = await FsTable.prototype._get.call(this, id);
     if (!doc) {
       throw new Error(`${this.name} does not exists: ${key}`);
     }
 
-    debug(`update ${this.name}`, { key, updates });
     Object.assign(doc, updates);
-    await this.collection.update(doc);
+
+    const updateAttrs = { ...doc };
+    // insert without tokens cause we've migrate token to balance table
+    if (this.syncBalance) {
+      delete updateAttrs.tokens;
+    }
+
+    await this.collection.update(updateAttrs);
+
+    if (this.syncBalance && doc.tokens) {
+      debug(`update balance for ${this.name}`, doc);
+
+      doc.tokens = await this.balanceTable.updateBalance({
+        address: doc.address,
+        tokens: doc.tokens,
+        context: doc.context,
+      });
+    }
 
     return doc;
   }
 
   updateOrCreate(exist, state, ctx) {
-    if (!state[this.uniqIndex]) {
+    const id = this.generatePrimaryKey(state);
+    const attrs = omit(state, this.primaryKey);
+
+    if (!id) {
       throw new Error('Cannot update or create without uniq index');
     }
 
     if (exist) {
-      return this.update(state[this.uniqIndex], omit(state, [this.uniqIndex]), ctx);
+      return this.update(id, attrs, ctx);
     }
 
-    return this.create(state[this.uniqIndex], omit(state, [this.uniqIndex]), ctx);
+    return this.create(id, attrs, ctx);
   }
 
   _reset() {

package/package.json

@@ -1,7 +1,7 @@
 {
   "name": "@ocap/statedb-fs",
   "description": "OCAP statedb adapter that uses fs as backend",
-  "version": "1.20.16",
+  "version": "1.21.1",
   "author": "wangshijun <shijun@arcblock.io> (https://www.arcblock.io)",
   "bugs": {
     "url": "https://github.com/ArcBlock/blockchain/issues",
@@ -36,9 +36,9 @@
     "debug": "^4.3.6",
     "lodash": "^4.17.21",
     "lokijs": "^1.5.12",
-    "@ocap/statedb": "1.20.16",
-    "@ocap/state": "1.20.16",
-    "@ocap/util": "1.20.16"
+    "@ocap/state": "1.21.1",
+    "@ocap/statedb": "1.21.1",
+    "@ocap/util": "1.21.1"
   },
   "scripts": {
     "lint": "eslint tests lib",