feat: add unlock/lock tables

lib/client.js

@@ -100,7 +100,7 @@ class RDSClient extends Operator {
 
   /**
    * doomed to be rollbacked after transaction scope
-   * useful on writing test that depend on database
+   * useful on writing tests which are related with database
    *
    * @param {Function} scope - scope with code
    * @param {Object} [ctx] - transaction env context, like koa's ctx.

lib/operator.js

@@ -367,6 +367,98 @@ class Operator {
     }
     return ' LIMIT ' + offset + ', ' + limit;
   }
+
+  /**
+   * Lock tables.
+   * @param {object[]} tables table lock descriptions.
+   * @description
+   * LOCK TABLES
+   *   tbl_name [[AS] alias] lock_type
+   *   [, tbl_name [[AS] alias] lock_type] ...
+   * lock_type: {
+   *   READ [LOCAL]
+   *   | [LOW_PRIORITY] WRITE
+   * }
+   * For more details:
+   * https://dev.mysql.com/doc/refman/8.0/en/lock-tables.html
+   * @example
+   * await locks([{ tableName: 'posts', lockType: 'READ', tableAlias: 't' }]);
+   */
+  async locks(tables) {
+    const sql = this._locks(tables);
+    debug('lock tables \n=> %j', sql);
+    return await this.query(sql);
+  }
+
+  /**
+   * Lock a single table.
+   * @param {string} tableName
+   * @param {string} lockType
+   * @param {string} tableAlias
+   * @description
+   * LOCK TABLES
+   *   tbl_name [[AS] alias] lock_type
+   *   [, tbl_name [[AS] alias] lock_type] ...
+   * lock_type: {
+   *   READ [LOCAL]
+   *   | [LOW_PRIORITY] WRITE
+   * }
+   * For more details:
+   * https://dev.mysql.com/doc/refman/8.0/en/lock-tables.html
+   * @example
+   * await lockOne('posts_table', 'READ', 't'); // LOCK TABLS 'posts_table' AS t READ
+   */
+  async lockOne(tableName, lockType, tableAlias) {
+    const sql = this._locks([{
+      tableName,
+      lockType,
+      tableAlias,
+    }]);
+    debug('lock one table \n=> %j', sql);
+    return await this.query(sql);
+  }
+
+  _locks(tableLocks) {
+    if (tableLocks.length === 0) {
+      throw new Error('Cannot lock empty tables.');
+    }
+    let sql = 'LOCK TABLES ';
+    for (let i = 0; i < tableLocks.length; i++) {
+      const table = tableLocks[i];
+      const { tableName, lockType, tableAlias } = table;
+      if (!tableName) {
+        throw new Error('No table_name provided while trying to lock table');
+      }
+      if (!lockType) {
+        throw new Error('No lock_type provided while trying to lock table `' + tableName + '`');
+      }
+      if ([ 'READ', 'WRITE', 'READ LOCAL', 'LOW_PRIORITY WRITE' ].indexOf(lockType.toUpperCase()) < 0) {
+        throw new Error('lock_type provided while trying to lock table `' + tableName +
+        '` must be one of the following(CASE INSENSITIVE):\n`READ` | `WRITE` | `READ LOCAL` | `LOW_PRIORITY WRITE`');
+      }
+      if (i > 0) {
+        sql += ', ';
+      }
+      sql += ' ' + this.escapeId(tableName) + ' ';
+      if (tableAlias) {
+        sql += ' AS ' + this.escapeId(tableAlias) + ' ';
+      }
+      sql += ' ' + lockType;
+    }
+    return sql + ';';
+  }
+
+  /**
+   * To unlock all tables locked in current session.
+   * For more details:
+   * https://dev.mysql.com/doc/refman/8.0/en/lock-tables.html
+   * @example
+   * await unlock(); // unlock all tables.
+   */
+  async unlock() {
+    debug('unlock tables');
+    return await this.query('UNLOCK TABLES;');
+  }
 }
 
 module.exports = Operator;

test/client.test.js

@@ -97,6 +97,69 @@ describe('test/client.test.js', () => {
     });
   });
 
+  describe('locks([...]), lockOne(name, lockType, alias), unlock()', () => {
+    it('validate arguments', async () => {
+      await assert.rejects(async () => {
+        await db.locks([
+          { tableName: 'xxxx' },
+        ]);
+      }, new Error('No lock_type provided while trying to lock table `xxxx`'));
+
+      await assert.rejects(async () => {
+        await db.locks([
+          { lockType: 'READ' },
+        ]);
+      }, new Error('No table_name provided while trying to lock table'));
+    });
+
+    it('should lock a table', async () => {
+      const sql = db._locks([
+        { tableName: 'posts', lockType: 'READ' },
+      ]);
+      assert.equal(sql.replaceAll(/\s+/g, ' '), 'LOCK TABLES `posts` READ;');
+    });
+
+    it('should lock multiple tables', async () => {
+      const sql = db._locks([
+        { tableName: 'posts', lockType: 'READ' },
+        { tableName: 'posts2', lockType: 'WRITE' },
+        { tableName: 'posts3', lockType: 'WRITE', tableAlias: 't' },
+      ]);
+      assert.equal(sql.replaceAll(/\s+/g, ' '), 'LOCK TABLES `posts` READ, `posts2` WRITE, `posts3` AS `t` WRITE;');
+      await assert.rejects(async () => {
+        await db.locks([
+          { tableName: 'xxxx' },
+        ]);
+      }, new Error('No lock_type provided while trying to lock table `xxxx`'));
+    });
+
+    it('should prevent sql injection', async () => {
+      // identifier injection test.
+      const sql = db._locks([
+        { tableName: '(select * from others)', lockType: 'READ' },
+        { tableName: ';-- \nshow tables;', lockType: 'READ' },
+      ]);
+      assert.equal(sql.replaceAll(/\s+/g, ' '), 'LOCK TABLES `(select * from others)` READ, `;-- show tables;` READ;');
+      // illeagle lockType test.
+      await assert.rejects(async () => {
+        db._locks([
+          { tableName: 'some table', lockType: '(show tables;)--' },
+        ]);
+      });
+    });
+
+    it('should unlock tables', async () => {
+      await db.lockOne('ali-sdk-test-user', 'READ', 't');
+      // error thrown: when table locked with alias, you can only query with the alias.
+      await assert.rejects(async () => {
+        await db.query('select * from `ali-sdk-test-user` limit 1;');
+      });
+      await db.unlock();
+      // recovered after unlock.
+      await db.query('select * from `ali-sdk-test-user` limit 1;');
+    });
+  });
+
   describe('transactions', () => {
     it('should beginTransaction error', async () => {
       const failDB = new RDSClient({
@@ -213,6 +276,24 @@ describe('test/client.test.js', () => {
       assert.equal(rows[1].name, prefix + 'beginTransaction2');
     });
 
+    it('should lock & unlock table during transaction', async () => {
+      const conn = await db.getConnection();
+      try {
+        await conn.beginTransaction();
+        await conn.lockOne('ali-sdk-test-user', 'READ', 't');
+        // error thrown: when table locked with alias, you can only query with the alias.
+        await assert.rejects(async () => {
+          await conn.query('select * from `ali-sdk-test-user` limit 1;');
+        });
+        await conn.unlock();
+        // recovered after unlock.
+        await conn.query('select * from `ali-sdk-test-user` limit 1;');
+      } catch (err) {
+        conn.release();
+        throw err;
+      }
+    });
+
     it('should rollback when query fail', async () => {
       const conn = await db.getConnection();
       try {