feat(postgres): added pool support

Author: rmolinamir

lib/postgres/src/__tests__/drizzle-orm.module.test.ts

@@ -2,13 +2,37 @@ import { PgDatabase } from 'drizzle-orm/pg-core';
 import { afterAll, beforeAll, describe, expect, test } from 'vitest';
 import { Test, TestingModule } from '@nestjs/testing';
 import { DrizzleOrmModule } from '../nestjs/drizzle-orm.module.js';
+import { Pool } from 'pg';
 
-describe('PgTransactionManager', () => {
+describe('DrizzleOrmModule', () => {
   let module: TestingModule;
 
   beforeAll(async () => {
     module = await Test.createTestingModule({
-      imports: [DrizzleOrmModule.forRoot({ client: globalThis.__pgClient })]
+      imports: [DrizzleOrmModule.forRoot({ pg: globalThis.__pgClient })]
+    }).compile();
+
+    await module.init();
+  });
+
+  afterAll(async () => {
+    await module?.close();
+  });
+
+  test('return an instance of PgDatabase', () => {
+    const db = module.get(PgDatabase);
+    expect(db).toBeInstanceOf(PgDatabase);
+  });
+});
+
+describe('DrizzleOrmModule with a pool', () => {
+  let module: TestingModule;
+
+  beforeAll(async () => {
+    module = await Test.createTestingModule({
+      imports: [
+        DrizzleOrmModule.forRoot({ pg: new Pool(globalThis.__dbCredentials) })
+      ]
     }).compile();
 
     await module.init();

lib/postgres/src/__tests__/pg-client.module.test.ts

@@ -37,4 +37,22 @@ describe('PgClientModule', () => {
 
     await module.close();
   });
+
+  test('module connects clients on bootstrap and disconnects on shutdown', async () => {
+    const module = await Test.createTestingModule({
+      imports: [PgClientModule.forRoot(globalThis.__dbCredentials)]
+    }).compile();
+
+    await module.init();
+
+    const client = module.get(pg.Client);
+
+    const res = await client.query('SELECT NOW()');
+
+    expect(res.rows[0].now).toBeDefined();
+
+    await module.close();
+
+    expect(() => client.query('SELECT NOW()')).rejects.toThrow();
+  });
 });

lib/postgres/src/__tests__/pg-pool.module.test.ts

@@ -0,0 +1,70 @@
+import pg from 'pg';
+import { describe, expect, test } from 'vitest';
+import { Test } from '@nestjs/testing';
+import { PgPoolModule } from '../nestjs/pg-pool.module.js';
+
+describe('PgPoolModule', () => {
+  test('module returns an instance of Pool using forRoot', async () => {
+    const module = await Test.createTestingModule({
+      imports: [
+        PgPoolModule.forRoot({
+          port: 5432
+        })
+      ]
+    }).compile();
+
+    const pool = module.get(pg.Pool);
+
+    expect(pool).toBeInstanceOf(pg.Pool);
+
+    await module.close();
+  });
+
+  test('module returns an instance of Client using forRootAsync', async () => {
+    const module = await Test.createTestingModule({
+      imports: [
+        PgPoolModule.forRootAsync({
+          useFactory: () => ({
+            port: 5432
+          })
+        })
+      ]
+    }).compile();
+
+    const pool = module.get(pg.Pool);
+
+    expect(pool).toBeInstanceOf(pg.Pool);
+
+    await module.close();
+  });
+
+  test('module disconnects pools on shutdown', async () => {
+    const module = await Test.createTestingModule({
+      imports: [PgPoolModule.forRoot(globalThis.__dbCredentials)]
+    }).compile();
+
+    await module.init();
+
+    const pool = module.get(pg.Pool);
+
+    const poolRes = await pool.query('SELECT NOW()');
+
+    expect(poolRes.rows[0].now).toBeDefined();
+
+    const client = await pool.connect();
+
+    const clientRes = await client.query('SELECT NOW()');
+
+    expect(clientRes.rows[0].now).toBeDefined();
+
+    client.release();
+
+    expect(() => client.query('SELECT NOW()')).rejects.toThrow();
+
+    await module.close();
+
+    expect(() => pool.connect()).rejects.toThrow();
+
+    expect(() => pool.query('SELECT NOW()')).rejects.toThrow();
+  });
+});

lib/postgres/src/__tests__/pg-transaction-manager.module.test.ts

@@ -24,7 +24,7 @@ describe('PgTransactionManagerModule', () => {
   beforeAll(async () => {
     module = await Test.createTestingModule({
       imports: [
-        DrizzleOrmModule.forRoot({ client: globalThis.__pgClient }),
+        DrizzleOrmModule.forRoot({ pg: globalThis.__pgClient }),
         PgTransactionManagerModule
       ]
     }).compile();

lib/postgres/src/__tests__/pg-transaction-manager.test.ts

@@ -17,6 +17,7 @@ import { PgTransactionManager } from '../pg-transaction-manager.js';
 import { testTable2 } from './test.table.js';
 import { InvalidOperationException } from '@ddd-framework/core';
 import { expectEntries } from './assertions.js';
+import { Pool } from 'pg';
 
 describe('PgTransactionManager', () => {
   let module: TestingModule;
@@ -25,7 +26,7 @@ describe('PgTransactionManager', () => {
 
   beforeAll(async () => {
     module = await Test.createTestingModule({
-      imports: [DrizzleOrmModule.forRoot({ client: globalThis.__pgClient })]
+      imports: [DrizzleOrmModule.forRoot({ pg: globalThis.__pgClient })]
     }).compile();
 
     await module.init();
@@ -170,3 +171,83 @@ describe('PgTransactionManager', () => {
     await expectEntries(db, testTable2, [{ id: firstId }, { id: secondId }]);
   });
 });
+
+describe('PgTransactionManager with a pool', () => {
+  let module: TestingModule;
+  let db: NodePgDatabase;
+  let manager: PgTransactionManager;
+
+  beforeAll(async () => {
+    module = await Test.createTestingModule({
+      imports: [
+        DrizzleOrmModule.forRoot({ pg: new Pool(globalThis.__dbCredentials) })
+      ]
+    }).compile();
+
+    await module.init();
+
+    db = module.get(PgDatabase);
+
+    // This would be returned by a dependency injection container.
+    manager = new PgTransactionManager(db);
+  });
+
+  beforeEach(async () => {
+    await db.delete(testTable2);
+  });
+
+  afterEach(async () => {
+    vi.clearAllMocks();
+  });
+
+  afterAll(async () => {
+    await module?.close();
+  });
+
+  test('starts transaction and implicitly commits', async () => {
+    const firstId = Uuid.generate();
+    const secondId = Uuid.generate();
+
+    await manager.startTransaction(async (transaction) => {
+      await transaction.context
+        .insert(testTable2)
+        .values([{ id: firstId }, { id: secondId }]);
+    });
+
+    await expectEntries(db, testTable2, [{ id: firstId }, { id: secondId }]);
+  });
+
+  test('returns result from transaction', async () => {
+    const firstId = Uuid.generate();
+    const secondId = Uuid.generate();
+
+    const result = await manager.startTransaction(async (transaction) => {
+      await transaction.context
+        .insert(testTable2)
+        .values([{ id: firstId }, { id: secondId }]);
+
+      return true;
+    });
+
+    expect(result).toBe(true);
+
+    await expectEntries(db, testTable2, [{ id: firstId }, { id: secondId }]);
+  });
+
+  test('transactions are automatically rolled back when an error occurs', async () => {
+    const firstId = Uuid.generate();
+    const secondId = Uuid.generate();
+
+    await expect(() =>
+      manager.startTransaction(async (transaction) => {
+        await transaction.context
+          .insert(testTable2)
+          .values([{ id: firstId }, { id: secondId }]);
+
+        throw new InvalidOperationException('An error occurred');
+      })
+    ).rejects.toThrow('An error occurred');
+
+    await expectEntries(db, testTable2, []);
+  });
+});

lib/postgres/src/__tests__/pg-transaction.test.ts

@@ -22,7 +22,7 @@ describe('PgTransaction', () => {
 
   beforeAll(async () => {
     module = await Test.createTestingModule({
-      imports: [DrizzleOrmModule.forRoot({ client: globalThis.__pgClient })]
+      imports: [DrizzleOrmModule.forRoot({ pg: globalThis.__pgClient })]
     }).compile();
 
     await module.init();

lib/postgres/src/nestjs/drizzle-orm.module.ts

@@ -14,7 +14,7 @@ import {
 } from '@nestjs/common';
 
 export type DrizzleOrmModuleOptions = {
-  client: pg.Client;
+  pg: pg.Client | pg.Pool;
   options?: DrizzleConfig;
 };
 
@@ -43,12 +43,12 @@ export class DrizzleOrmModule {
    * Configurations are only loaded once because the internal module is global.
    */
   public static forRoot({
-    client,
+    pg,
     options
   }: DrizzleOrmModuleOptions): DynamicModule {
     const dbProvider: Provider<PgDatabase<PgQueryResultHKT>> = {
       provide: PgDatabase,
-      useValue: drizzle(client, options)
+      useValue: drizzle(pg, options)
     };
 
     return {
@@ -69,9 +69,9 @@ export class DrizzleOrmModule {
     const dbProvider = {
       provide: PgDatabase,
       useFactory: async ({
-        options = {},
-        client
-      }: DrizzleOrmModuleOptions): Promise<unknown> => drizzle(client, options),
+        pg,
+        options = {}
+      }: DrizzleOrmModuleOptions): Promise<unknown> => drizzle(pg, options),
       inject: [this.OPTIONS_PROVIDER_TOKEN]
     };