feature(extension-driver-canner): connect to datasource using the connection info of the API requester

Author: onlyjackfrost

packages/extension-driver-canner/src/lib/cannerDataSource.ts

@@ -24,7 +24,11 @@ export interface PGOptions extends PoolConfig {
 @VulcanExtensionId('canner')
 export class CannerDataSource extends DataSource<any, PGOptions> {
   private logger = this.getLogger();
-  private poolMapping = new Map<string, { pool: Pool; options?: PGOptions }>();
+  protected poolMapping = new Map<
+    string,
+    { pool: Pool; options?: PGOptions }
+  >();
+  protected UserPool = new Map<string, Pool>();
 
   public override async onActivate() {
     const profiles = this.getProfiles().values();
@@ -108,14 +112,14 @@ export class CannerDataSource extends DataSource<any, PGOptions> {
     bindParams,
     profileName,
     operations,
+    headers,
   }: ExecuteOptions): Promise<DataResult> {
-    if (!this.poolMapping.has(profileName)) {
-      throw new InternalError(`Profile instance ${profileName} not found`);
-    }
-    const { pool, options } = this.poolMapping.get(profileName)!;
-    this.logger.debug(`Acquiring connection from ${profileName}`);
-    const client = await pool.connect();
     this.logger.debug(`Acquired connection from ${profileName}`);
+    const { options } = this.poolMapping.get(profileName)!;
+    const auth = headers?.['authorization'];
+    const password = auth?.trim().split(' ')[1];
+    const pool = this.getPool(profileName, password);
+    const client = await pool.connect();
     try {
       const builtSQL = buildSQL(sql, operations);
       const cursor = client.query(
@@ -150,6 +154,38 @@ export class CannerDataSource extends DataSource<any, PGOptions> {
     }
   }
 
+  // use protected to make it testable
+  protected getPool(profileName: string, password?: string): Pool {
+    if (!this.poolMapping.has(profileName)) {
+      throw new InternalError(`Profile instance ${profileName} not found`);
+    }
+    const { pool: defaultPool, options: poolOptions } =
+      this.poolMapping.get(profileName)!;
+    this.logger.debug(`Acquiring connection from ${profileName}`);
+    if (!password) {
+      return defaultPool;
+    }
+    const database = poolOptions?.database || '';
+    const userPoolKey = this.getUserPoolKey(password, database);
+    if (this.UserPool.has(userPoolKey)) {
+      const userPool = this.UserPool.get(userPoolKey);
+      if (!userPool) {
+        throw new InternalError(
+          `User pool ${userPoolKey} is not a Pool instance`
+        );
+      }
+      return userPool;
+    }
+    const pool = new Pool({ ...poolOptions, password: password });
+    this.UserPool.set(userPoolKey, pool);
+    return pool;
+  }
+
+  // use protected to make it testable
+  protected getUserPoolKey(pat: string, database?: string) {
+    return `${pat}-${database}`;
+  }
+
   private async getResultFromCursor(
     cursor: Cursor,
     options: PGOptions = {}

packages/extension-driver-canner/test/cannerDataSource.spec.ts

@@ -1,5 +1,6 @@
 import { CannerServer } from './cannerServer';
 import { CannerDataSource, PGOptions } from '../src';
+import { MockCannerDataSource } from './mock';
 import { ExportOptions, InternalError, streamToArray } from '@vulcan-sql/core';
 import { Writable } from 'stream';
 import * as sinon from 'ts-sinon';
@@ -8,7 +9,9 @@ import { CannerAdapter } from '../src/lib/cannerAdapter';
 
 const pg = new CannerServer();
 let dataSource: CannerDataSource;
+let mockDataSource: MockCannerDataSource;
 
+const directory = 'tmp_test_canner';
 // restore all sinon mock/stub before each test
 beforeEach(() => {
   sinon.default.restore();
@@ -42,22 +45,22 @@ it('Data source should throw error when activating if any profile is invalid', a
 
 // export method should be executed successfully
 it('Data source should export successfully', async () => {
-  fs.mkdirSync('tmp', { recursive: true });
+  fs.mkdirSync(directory, { recursive: true });
   dataSource = new CannerDataSource({}, '', [pg.getProfile('profile1')]);
   await dataSource.activate();
 
   // Act, Assert
   await expect(
     dataSource.export({
       sql: 'select 1',
-      directory: 'tmp',
+      directory,
       profileName: 'profile1',
     } as ExportOptions)
   ).resolves.not.toThrow();
-  expect(fs.readdirSync('tmp').length).toBe(1);
+  expect(fs.readdirSync(directory).length).toBe(1);
 
   // clean up
-  fs.rmSync('tmp', { recursive: true, force: true });
+  fs.rmSync(directory, { recursive: true, force: true });
 }, 100000);
 
 it('Data source should throw error when fail to export data', async () => {
@@ -73,22 +76,22 @@ it('Data source should throw error when fail to export data', async () => {
       );
     });
 
-  fs.mkdirSync('tmp', { recursive: true });
+  fs.mkdirSync(directory, { recursive: true });
   dataSource = new CannerDataSource({}, '', [pg.getProfile('profile1')]);
   await dataSource.activate();
 
   // Act, Assert
   await expect(
     dataSource.export({
       sql: 'select 1',
-      directory: 'tmp',
+      directory,
       profileName: 'profile1',
     } as ExportOptions)
   ).rejects.toThrow();
-  expect(fs.readdirSync('tmp').length).toBe(0);
+  expect(fs.readdirSync(directory).length).toBe(0);
 
   // clean up
-  fs.rmSync('tmp', { recursive: true, force: true });
+  fs.rmSync(directory, { recursive: true, force: true });
 }, 100000);
 
 it('Data source should throw error when given directory is not exist', async () => {
@@ -100,7 +103,7 @@ it('Data source should throw error when given directory is not exist', async ()
   await expect(
     dataSource.export({
       sql: 'select 1',
-      directory: 'tmp',
+      directory: directory,
       profileName: 'profile1',
     } as ExportOptions)
   ).rejects.toThrow();
@@ -110,13 +113,13 @@ it('Data source should throw error when given profile name is not exist', async
   // Arrange
   dataSource = new CannerDataSource({}, '', [pg.getProfile('profile1')]);
   await dataSource.activate();
-  fs.mkdirSync('tmp', { recursive: true });
+  fs.mkdirSync(directory, { recursive: true });
 
   // Act, Assert
   await expect(
     dataSource.export({
       sql: 'select 1',
-      directory: 'tmp',
+      directory,
       profileName: 'profile not exist',
     } as ExportOptions)
   ).rejects.toThrow();
@@ -318,3 +321,55 @@ it('Data source should release connection when readable stream is destroyed', as
   expect(rows.length).toBe(1);
   // afterEach hook will timeout if any leak occurred.
 }, 300000);
+
+it('Should return the same pool when the profile is the same', async () => {
+  // Arrange
+  mockDataSource = new MockCannerDataSource({}, '', [
+    pg.getProfile('profile1'),
+  ]);
+  await mockDataSource.activate();
+  // Act
+  const pool1 = mockDataSource.getPool('profile1');
+  const pool2 = mockDataSource.getPool('profile1');
+  // Assert
+  expect(pool1 === pool2).toBeTruthy();
+}, 30000);
+
+it('Should return the same pool when the profile and authentication is the same', async () => {
+  // Arrange
+  mockDataSource = new MockCannerDataSource({}, '', [
+    pg.getProfile('profile1'),
+  ]);
+  await mockDataSource.activate();
+  // Act
+  const pool1 = mockDataSource.getPool('profile1', 'the-same-authentication');
+  const pool2 = mockDataSource.getPool('profile1', 'the-same-authentication');
+  // Assert
+  expect(pool1 === pool2).toBeTruthy();
+}, 30000);
+
+it('Should return different pool if authentication exist in headers even the profile is the same', async () => {
+  // Arrange
+  mockDataSource = new MockCannerDataSource({}, '', [
+    pg.getProfile('profile1'),
+  ]);
+  await mockDataSource.activate();
+  // Act
+  const pool1 = mockDataSource.getPool('profile1');
+  const pool2 = mockDataSource.getPool('profile1', 'my-authentication');
+  // Assert
+  expect(pool1 == pool2).toBeFalsy();
+}, 30000);
+
+it('Should return different pool with different authentication even the profile is the same', async () => {
+  // Arrange
+  mockDataSource = new MockCannerDataSource({}, '', [
+    pg.getProfile('profile1'),
+  ]);
+  await mockDataSource.activate();
+  // Act
+  const pool1 = mockDataSource.getPool('profile1', 'authentication');
+  const pool2 = mockDataSource.getPool('profile1', 'differ-authentication');
+  // Assert
+  expect(pool1 === pool2).toBeFalsy();
+}, 30000);

packages/extension-driver-canner/test/mock/index.ts

@@ -0,0 +1 @@
+export * from './mockCannerDataSource';

packages/extension-driver-canner/test/mock/mockCannerDataSource.ts

@@ -0,0 +1,30 @@
+import { CannerDataSource } from '../../src';
+import { InternalError } from '@vulcan-sql/core';
+import { Pool } from 'pg';
+
+export class MockCannerDataSource extends CannerDataSource {
+  public override getPool(profileName: string, password?: string): Pool {
+    if (!this.poolMapping.has(profileName)) {
+      throw new InternalError(`Profile instance ${profileName} not found`);
+    }
+    const { pool: defaultPool, options: poolOptions } =
+      this.poolMapping.get(profileName)!;
+    if (!password) {
+      return defaultPool;
+    }
+    const database = poolOptions?.database || '';
+    const userPoolKey = this.getUserPoolKey(password, database);
+    if (this.UserPool.has(userPoolKey)) {
+      const userPool = this.UserPool.get(userPoolKey);
+      if (!userPool) {
+        throw new InternalError(
+          `User pool ${userPoolKey} is not a Pool instance`
+        );
+      }
+      return userPool;
+    }
+    const pool = new Pool({ ...poolOptions, password: password });
+    this.UserPool.set(userPoolKey, pool);
+    return pool;
+  }
+}

packages/extension-driver-duckdb/src/lib/duckdbDataSource.ts

@@ -93,10 +93,13 @@ export class DuckDBDataSource extends DataSource<any, DuckDBOptions> {
     bindParams,
     profileName,
     operations,
+    headers,
   }: ExecuteOptions): Promise<DataResult> {
     if (!this.dbMapping.has(profileName)) {
       throw new InternalError(`Profile instance ${profileName} not found`);
     }
+    console.log(`execute duckdb: ${sql}`);
+    console.log({ headers });
     const { db, configurationParameters, ...options } =
       this.dbMapping.get(profileName)!;
     const [firstDataSQL, restDataSQL] = buildSQL(sql, operations);