feat: dynamic retrieval of database connection configuration (#110)

test fix: Node.js >= 20 will return localhost ipv6 address

Author: fengmk2

package.json

@@ -22,13 +22,13 @@
   "devDependencies": {
     "@eggjs/tsconfig": "^1.3.2",
     "@types/mocha": "^10.0.1",
-    "@types/node": "^18.14.6",
+    "@types/node": "^20.2.5",
     "egg-bin": "^6.1.2",
     "eslint": "^8.29.0",
     "eslint-config-egg": "^12.1.0",
     "git-contributor": "^2.0.0",
     "mm": "^3.3.0",
-    "typescript": "^4.9.5"
+    "typescript": "^5.1.3"
   },
   "homepage": "https://github.com/ali-sdk/ali-rds",
   "repository": {

src/PoolConfig.ts

@@ -0,0 +1,19 @@
+import MySQLPoolConfig from 'mysql/lib/PoolConfig';
+import type { PoolConfig, ConnectionConfig } from 'mysql';
+import type { GetConnectionConfig } from './types';
+
+export class RDSPoolConfig extends MySQLPoolConfig {
+  #getConnectionConfig: GetConnectionConfig;
+
+  constructor(options: PoolConfig, getConnectionConfig: GetConnectionConfig) {
+    super(options);
+    this.#getConnectionConfig = getConnectionConfig;
+  }
+
+  newConnectionConfig(): ConnectionConfig {
+    return {
+      ...super.newConnectionConfig(),
+      ...this.#getConnectionConfig(),
+    };
+  }
+}

src/client.ts

@@ -6,6 +6,7 @@ import type { PoolConnectionPromisify, RDSClientOptions, TransactionContext, Tra
 import { Operator } from './operator';
 import { RDSConnection } from './connection';
 import { RDSTransaction } from './transaction';
+import { RDSPoolConfig } from './PoolConfig';
 import literals from './literals';
 
 interface PoolPromisify extends Omit<Pool, 'query'> {
@@ -35,7 +36,18 @@ export class RDSClient extends Operator {
   constructor(options: RDSClientOptions) {
     super();
     const { connectionStorage, connectionStorageKey, ...mysqlOptions } = options;
-    this.#pool = mysql.createPool(mysqlOptions) as unknown as PoolPromisify;
+    // get connection options from getConnectionConfig method every time
+    if (mysqlOptions.getConnectionConfig) {
+      // eslint-disable-next-line @typescript-eslint/no-var-requires
+      const Pool = require('mysql/lib/Pool');
+      this.#pool = new Pool({
+        config: new RDSPoolConfig(mysqlOptions, mysqlOptions.getConnectionConfig),
+      });
+      // override _needsChangeUser to return false
+      (this.#pool as any)._needsChangeUser = () => false;
+    } else {
+      this.#pool = mysql.createPool(mysqlOptions) as unknown as PoolPromisify;
+    }
     [
       'query',
       'getConnection',

src/types.ts

@@ -1,10 +1,13 @@
 import { AsyncLocalStorage } from 'async_hooks';
-import type { PoolConnection, PoolConfig } from 'mysql';
+import type { PoolConnection, PoolConfig, ConnectionConfig } from 'mysql';
 import { RDSTransaction } from './transaction';
 
+export type GetConnectionConfig = () => ConnectionConfig;
+
 export interface RDSClientOptions extends PoolConfig {
   connectionStorageKey?: string;
   connectionStorage?: AsyncLocalStorage<Record<PropertyKey, RDSTransaction>>;
+  getConnectionConfig?: GetConnectionConfig
 }
 
 export interface PoolConnectionPromisify extends Omit<PoolConnection, 'query'> {

test/PoolConfig.test.ts

@@ -0,0 +1,84 @@
+import { strict as assert } from 'node:assert';
+import fs from 'node:fs/promises';
+import path from 'node:path';
+import mm from 'mm';
+import config from './config';
+import { RDSClient } from '../src/client';
+
+describe('test/PoolConfig.test.ts', () => {
+  const prefix = 'prefix-PoolConfig' + process.version + '-';
+  const table = 'ali-sdk-test-user';
+  let db: RDSClient;
+  let index = 0;
+  let newConnectionCount = 0;
+  before(async () => {
+    db = new RDSClient({
+      // test getConnectionConfig
+      connectionLimit: 2,
+      getConnectionConfig() {
+        console.log('get connection config index: %d', ++index);
+        return config;
+      },
+    });
+    db.pool.on('acquire', conn => {
+      console.log('acquire connection %o', conn.threadId);
+    });
+    db.pool.on('connection', conn => {
+      newConnectionCount++;
+      console.log('new connection %o', conn.threadId);
+    });
+    db.pool.on('enqueue', () => {
+      console.log('Waiting for available connection slot');
+    });
+    db.pool.on('release', conn => {
+      console.log('release connection %o', conn.threadId);
+    });
+
+    try {
+      const sql = await fs.readFile(path.join(__dirname, 'rds_init.sql'), 'utf-8');
+      await db.query(sql);
+    } catch (err) {
+      console.log('init table error: %s', err);
+    }
+    await db.query('delete from ?? where name like ?', [ table, prefix + '%' ]);
+  });
+
+  after(async () => {
+    return await db.end();
+  });
+
+  afterEach(() => {
+    mm.restore();
+  });
+
+  describe('new RDSClient(options.getConnectionConfig)', () => {
+    it('should get connection config from newConnectionConfig()', async () => {
+      assert.equal(db.pool.config.connectionConfig.database, undefined);
+      assert.equal(index, 1);
+      assert.equal((db.pool.config as any).newConnectionConfig().database, 'test');
+      assert.equal(index, 2);
+    });
+
+    it('should connect rds success', async () => {
+      const rows = await db.query('show tables');
+      // console.log(rows);
+      assert(rows);
+      assert(Array.isArray(rows));
+      assert.equal(index, 2);
+      const results = await Promise.all([
+        db.query('show tables'),
+        db.query('show tables'),
+        db.query('show tables'),
+      ]);
+      assert.equal(results.length, 3);
+      assert(results[0]);
+      assert(Array.isArray(results[0]));
+      assert(results[1]);
+      assert(Array.isArray(results[1]));
+      assert(results[2]);
+      assert(Array.isArray(results[2]));
+      assert.equal(index, 3);
+      assert.equal(newConnectionCount, 2);
+    });
+  });
+});

test/client.test.ts

@@ -172,6 +172,7 @@ describe('test/client.test.ts', () => {
     it('should beginTransaction error', async () => {
       const failDB = new RDSClient({
         port: 12312,
+        host: '127.0.0.1',
       });
       await assert.rejects(async () => {
         await failDB.beginTransaction();
@@ -339,6 +340,7 @@ describe('test/client.test.ts', () => {
     it('should beginTransactionScope() error', async () => {
       const failDB = new RDSClient({
         port: 12312,
+        host: '127.0.0.1',
       });
       await assert.rejects(async () => {
         await failDB.beginTransactionScope(async () => {
@@ -1321,6 +1323,7 @@ describe('test/client.test.ts', () => {
     it('should throw error when mysql connect fail', async () => {
       const db = new RDSClient({
         port: 33061,
+        host: '127.0.0.1',
       });
       try {
         await db.getConnection();

test/config.ts

@@ -1,5 +1,6 @@
 export default {
-  host: 'localhost',
+  // host: 'localhost',
+  host: '127.0.0.1',
   port: 3306,
   user: 'root',
   password: '',