feat: support custom query lifecricle (#104)

ali-sdk · Mar 5, 2023 · 5941c69 · 5941c69
1 parent e7d488a
commit 5941c69
Show file tree

Hide file tree

Showing 6 changed files with 177 additions and 3 deletions.
diff --git a/README.md b/README.md
@@ -356,6 +356,19 @@ const rows = await db.query('SELECT * FROM your_table WHERE id=:id', { id: 123 }
 console.log(rows);
 ```
 
+### Custom query lifecricle
+
+```ts
+db.beforeQuery((sql: string) => {
+  // change sql string
+  return `/* add custom format here */ ${sql}`;
+});
+
+db.afterQuery((sql: string, result: any, execDuration: number, err?: Error) => {
+  // handle logger here
+});
+```
+
 ## APIs
 
 - `*` Meaning this function is yieldable.

diff --git a/src/client.ts b/src/client.ts
@@ -57,8 +57,15 @@ export class RDSClient extends Operator {
 
   async getConnection() {
     try {
-      const conn = await this.#pool.getConnection();
-      return new RDSConnection(conn);
+      const _conn = await this.#pool.getConnection();
+      const conn = new RDSConnection(_conn);
+      if (this.beforeQueryHandler) {
+        conn.beforeQuery(this.beforeQueryHandler);
+      }
+      if (this.afterQueryHandler) {
+        conn.afterQuery(this.afterQueryHandler);
+      }
+      return conn;
     } catch (err) {
       if (err.name === 'Error') {
         err.name = 'RDSClientGetConnectionError';
@@ -80,7 +87,14 @@ export class RDSClient extends Operator {
       conn.release();
       throw err;
     }
-    return new RDSTransaction(conn);
+    const tran = new RDSTransaction(conn);
+    if (this.beforeQueryHandler) {
+      tran.beforeQuery(this.beforeQueryHandler);
+    }
+    if (this.afterQueryHandler) {
+      tran.afterQuery(this.afterQueryHandler);
+    }
+    return tran;
   }
 
   /**

diff --git a/src/operator.ts b/src/operator.ts
@@ -2,6 +2,7 @@ import { debuglog } from 'node:util';
 import { SqlString } from './sqlstring';
 import literals from './literals';
 import {
+  AfterQueryHandler, BeforeQueryHandler,
   DeleteResult,
   InsertOption, InsertResult,
   LockResult, LockTableOption,
@@ -15,8 +16,19 @@ const debug = debuglog('ali-rds:operator');
  * Operator Interface
  */
 export abstract class Operator {
+  protected beforeQueryHandler?: BeforeQueryHandler;
+  protected afterQueryHandler?: AfterQueryHandler;
+
   get literals() { return literals; }
 
+  beforeQuery(beforeQueryHandler: BeforeQueryHandler) {
+    this.beforeQueryHandler = beforeQueryHandler;
+  }
+
+  afterQuery(afterQueryHandler: AfterQueryHandler) {
+    this.afterQueryHandler = afterQueryHandler;
+  }
+
   escape(value: any, stringifyObjects?: boolean, timeZone?: string): string {
     return SqlString.escape(value, stringifyObjects, timeZone);
   }
@@ -45,17 +57,33 @@ export abstract class Operator {
     if (values) {
       sql = this.format(sql, values);
     }
+    if (this.beforeQueryHandler) {
+      const newSql = this.beforeQueryHandler(sql);
+      if (newSql) {
+        sql = newSql;
+      }
+    }
     debug('query %o', sql);
+    let execDuration: number;
+    const queryStart = Date.now();
     try {
       const rows = await this._query(sql);
+      execDuration = Date.now() - queryStart;
+      if (this.afterQueryHandler) {
+        this.afterQueryHandler(sql, rows, execDuration);
+      }
       if (Array.isArray(rows)) {
         debug('query get %o rows', rows.length);
       } else {
         debug('query result: %o', rows);
       }
       return rows;
     } catch (err) {
+      execDuration = Date.now() - queryStart;
       err.stack = `${err.stack}\n    sql: ${sql}`;
+      if (this.afterQueryHandler) {
+        this.afterQueryHandler(sql, null, execDuration, err);
+      }
       debug('query error: %o', err);
       throw err;
     }

diff --git a/src/types.ts b/src/types.ts
@@ -50,3 +50,6 @@ export type LockTableOption = {
   lockType: string;
   tableAlias: string;
 };
+
+export type BeforeQueryHandler = (sql: string) => string | undefined | void;
+export type AfterQueryHandler = (sql: string, result: any, execDuration: number, err?: Error) => void;
diff --git a/test/client.test.ts b/test/client.test.ts
@@ -1203,4 +1203,59 @@ describe('test/client.test.ts', () => {
       await db2.end();
     });
   });
+
+  describe('query lifecricle work', () => {
+    it('should work on client and transactions', async () => {
+      const db = new RDSClient(config);
+      let count = 0;
+      let lastSql = '';
+      db.beforeQuery(sql => {
+        count++;
+        lastSql = sql;
+      });
+      let lastArgs: any;
+      db.afterQuery((...args) => {
+        lastArgs = args;
+      });
+      await db.query('select * from ?? limit 10', [ table ]);
+      assert.equal(lastSql, 'select * from `ali-sdk-test-user` limit 10');
+      assert.equal(lastArgs[0], lastSql);
+      assert.equal(Array.isArray(lastArgs[1]), true);
+      assert.equal(count, 1);
+
+      await db.beginTransactionScope(async conn => {
+        await conn.query(`insert into ??(name, email, gmt_create, gmt_modified)
+          values(?, ?, now(), now())`,
+        [ table, prefix + 'beginTransactionScope1', prefix + '[email protected]' ]);
+      });
+      assert.equal(lastSql, 'insert into `ali-sdk-test-user`(name, email, gmt_create, gmt_modified)\n' +
+        `          values('${prefix}beginTransactionScope1', '${prefix}[email protected]', now(), now())`);
+      assert.equal(lastArgs[0], lastSql);
+      assert.equal(lastArgs[1].affectedRows, 1);
+      assert.equal(count, 2);
+
+      await db.beginDoomedTransactionScope(async conn => {
+        await conn.query(`insert into ??(name, email, gmt_create, gmt_modified)
+          values(?, ?, now(), now())`,
+        [ table, prefix + 'beginDoomedTransactionScope1', prefix + '[email protected]' ]);
+      });
+      assert.equal(lastSql, 'insert into `ali-sdk-test-user`(name, email, gmt_create, gmt_modified)\n' +
+        `          values('${prefix}beginDoomedTransactionScope1', '${prefix}[email protected]', now(), now())`);
+      assert.equal(lastArgs[0], lastSql);
+      assert.equal(lastArgs[1].affectedRows, 1);
+      assert.equal(count, 3);
+
+      const conn = await db.getConnection();
+      await conn.beginTransaction();
+      await conn.query(`insert into ??(name, email, gmt_create, gmt_modified)
+        values(?, ?, now(), now())`,
+      [ table, prefix + 'transaction1', prefix + '[email protected]' ]);
+      await conn.commit();
+      assert.equal(lastSql, 'insert into `ali-sdk-test-user`(name, email, gmt_create, gmt_modified)\n' +
+        `        values('${prefix}transaction1', '${prefix}[email protected]', now(), now())`);
+      assert.equal(lastArgs[0], lastSql);
+      assert.equal(lastArgs[1].affectedRows, 1);
+      assert.equal(count, 4);
+    });
+  });
 });
diff --git a/test/operator.test.ts b/test/operator.test.ts
@@ -65,4 +65,65 @@ describe('test/operator.test.ts', () => {
       }
     });
   });
+
+  describe('beforeQuery(), afterQuery()', () => {
+    class CustomOperator extends Operator {
+      protected async _query(sql: string): Promise<any> {
+        // console.log(sql);
+        if (sql === 'error') throw new Error('mock error');
+        return { sql };
+      }
+    }
+
+    it('should override query sql', async () => {
+      const op = new CustomOperator();
+      op.beforeQuery(sql => {
+        return `hello ${sql}`;
+      });
+      const result = await op.query('foo');
+      assert.equal(result.sql, 'hello foo');
+    });
+
+    it('should not override query sql', async () => {
+      const op = new CustomOperator();
+      op.beforeQuery(sql => {
+        assert(sql);
+      });
+      const result = await op.query('foo');
+      assert.equal(result.sql, 'foo');
+    });
+
+    it('should get query result on after hook', async () => {
+      const op = new CustomOperator();
+      op.afterQuery((sql, result, execDuration, err) => {
+        assert.equal(sql, 'foo');
+        assert.deepEqual(result, { sql });
+        assert.equal(typeof execDuration, 'number');
+        assert(execDuration >= 0);
+        assert.equal(err, undefined);
+      });
+      const result = await op.query('foo');
+      assert.equal(result.sql, 'foo');
+    });
+
+    it('should get query error on after hook', async () => {
+      const op = new CustomOperator();
+      op.afterQuery((sql, result, execDuration, err) => {
+        assert.equal(sql, 'error');
+        assert.equal(result, null);
+        assert.equal(typeof execDuration, 'number');
+        assert(execDuration >= 0);
+        assert(err instanceof Error);
+        assert.equal(err.message, 'mock error');
+        assert.match(err.stack!, /sql: error/);
+      });
+      await assert.rejects(async () => {
+        await op.query('error');
+      }, (err: any) => {
+        assert.equal(err.message, 'mock error');
+        assert.match(err.stack, /sql: error/);
+        return true;
+      });
+    });
+  });
 });