Fixed connection pool management to correctly handle database names

Thanks to that pongoClient.close can close connection pools without needing to call endAllPools.

PongoClient will track opened dbClients and close them all.

Author: oskardudycz

src/package-lock.json

@@ -69,6 +69,7 @@
     "@types/mongodb": "^4.0.7",
     "@types/node": "20.11.30",
     "@types/pg": "^8.11.6",
+    "@types/pg-connection-string": "^2.0.0",
     "@types/pg-format": "^1.0.5",
     "@types/uuid": "9.0.8",
     "@typescript-eslint/eslint-plugin": "7.9.0",
@@ -89,6 +90,7 @@
   },
   "peerDependencies": {
     "pg": "^8.12.0",
+    "pg-connection-string": "^2.6.4",
     "pg-format": "^1.0.4",
     "testcontainers": "^10.10.1"
   },

src/package.json

@@ -47,14 +47,19 @@
     "dist"
   ],
   "peerDependencies": {
-    "@types/uuid": "^9.0.8",
     "@types/pg": "^8.11.6",
+    "@types/pg-connection-string": "^2.0.0",
     "@types/pg-format": "^1.0.5",
+    "@types/uuid": "^9.0.8",
     "pg": "^8.12.0",
+    "pg-connection-string": "^2.6.4",
     "pg-format": "^1.0.4",
     "uuid": "^9.0.1"
   },
   "devDependencies": {
     "@types/node": "20.11.30"
+  },
+  "dependencies": {
+    "@event-driven-io/dumbo": "^0.1.0"
   }
 }

src/packages/dumbo/package.json

@@ -14,6 +14,6 @@ export const postgresClient = (
 
   return {
     connect: () => pool.connect(),
-    close: () => endPool(connectionString),
+    close: () => endPool({ connectionString, database }),
   };
 };

src/packages/dumbo/src/connections/client.ts

@@ -0,0 +1,6 @@
+import pgcs from 'pg-connection-string';
+
+export const defaultPostgreSqlDatabase = 'postgres';
+
+export const getDatabaseNameOrDefault = (connectionString: string) =>
+  pgcs.parse(connectionString).database ?? defaultPostgreSqlDatabase;

src/packages/dumbo/src/connections/connectionString.ts

@@ -1,2 +1,3 @@
 export * from './client';
+export * from './connectionString';
 export * from './pool';

src/packages/dumbo/src/connections/index.ts

@@ -1,6 +1,11 @@
 import pg from 'pg';
+import {
+  defaultPostgreSqlDatabase,
+  getDatabaseNameOrDefault,
+} from './connectionString';
 
 const pools: Map<string, pg.Pool> = new Map();
+const usageCounter: Map<string, number> = new Map();
 
 export const getPool = (
   connectionStringOrOptions: string | pg.PoolConfig,
@@ -15,22 +20,63 @@ export const getPool = (
       ? { connectionString }
       : connectionStringOrOptions;
 
-  //TODO: this should include database name resolution for key
+  const database =
+    poolOptions.database ??
+    (poolOptions.connectionString
+      ? getDatabaseNameOrDefault(poolOptions.connectionString)
+      : undefined);
+
+  const lookupKey = key(connectionString, database);
+
+  updatePoolUsageCounter(lookupKey, 1);
+
   return (
-    pools.get(connectionString) ??
-    pools.set(connectionString, new pg.Pool(poolOptions)).get(connectionString)!
+    pools.get(lookupKey) ??
+    pools.set(lookupKey, new pg.Pool(poolOptions)).get(lookupKey)!
   );
 };
 
-export const endPool = async (connectionString: string): Promise<void> => {
-  const pool = pools.get(connectionString);
-  if (pool) {
+export const endPool = async ({
+  connectionString,
+  database,
+  force,
+}: {
+  connectionString: string;
+  database?: string | undefined;
+  force?: boolean;
+}): Promise<void> => {
+  database = database ?? getDatabaseNameOrDefault(connectionString);
+  const lookupKey = key(connectionString, database);
+
+  const pool = pools.get(lookupKey);
+  if (pool && (updatePoolUsageCounter(lookupKey, -1) <= 0 || force === true)) {
+    await onEndPool(lookupKey, pool);
+  }
+};
+
+export const onEndPool = async (lookupKey: string, pool: pg.Pool) => {
+  try {
     await pool.end();
-    pools.delete(connectionString);
+    pools.delete(lookupKey);
+  } catch (error) {
+    console.log(`Error while closing the connection pool: ${lookupKey}`);
+    console.log(error);
   }
 };
 
 export const endAllPools = () =>
   Promise.all(
-    [...pools.keys()].map((connectionString) => endPool(connectionString)),
+    [...pools.entries()].map(([lookupKey, pool]) => onEndPool(lookupKey, pool)),
   );
+
+const key = (connectionString: string, database: string | undefined) =>
+  `${connectionString}|${database ?? defaultPostgreSqlDatabase}`;
+
+const updatePoolUsageCounter = (lookupKey: string, by: 1 | -1): number => {
+  const currentCounter = usageCounter.get(lookupKey) ?? 0;
+  const newCounter = currentCounter + by;
+
+  usageCounter.set(lookupKey, currentCounter + by);
+
+  return newCounter;
+};

src/packages/dumbo/src/connections/pool.ts

@@ -47,16 +47,20 @@
     "dist"
   ],
   "peerDependencies": {
-    "@types/uuid": "^9.0.8",
+    "@event-driven-io/dumbo": "^0.1.0",
+    "@types/mongodb": "^4.0.7",
     "@types/pg": "^8.11.6",
     "@types/pg-format": "^1.0.5",
-    "@types/mongodb": "^4.0.7",
-    "@event-driven-io/dumbo": "^0.1.0",
+    "@types/uuid": "^9.0.8",
     "pg": "^8.12.0",
     "pg-format": "^1.0.4",
     "uuid": "^9.0.1"
   },
   "devDependencies": {
     "@types/node": "20.11.30"
+  },
+  "dependencies": {
+    "@types/pg-connection-string": "^2.0.0",
+    "pg-connection-string": "^2.6.4"
   }
 }

src/packages/pongo/package.json

@@ -1,4 +1,3 @@
-import { endAllPools } from '@event-driven-io/dumbo';
 import {
   MongoDBContainer,
   type StartedMongoDBContainer,
@@ -60,7 +59,7 @@ void describe('MongoDB Compatibility Tests', () => {
 
   after(async () => {
     try {
-      await endAllPools();
+      await pongoClient.close();
       await postgres.stop();
     } catch (error) {
       console.log(error);

src/packages/pongo/src/e2e/compatibilityTest.e2e.spec.ts

@@ -1,17 +1,34 @@
-import { getDbClient } from './dbClient';
+import { getDatabaseNameOrDefault } from '@event-driven-io/dumbo';
+import { getDbClient, type DbClient } from './dbClient';
 import type { PongoClient, PongoDb } from './typing/operations';
 
 export const pongoClient = (connectionString: string): PongoClient => {
+  const defaultDbName = getDatabaseNameOrDefault(connectionString);
+  const dbClients: Map<string, DbClient> = new Map();
+
   const dbClient = getDbClient(connectionString);
+  dbClients.set(defaultDbName, dbClient);
 
   const pongoClient: PongoClient = {
     connect: async () => {
       await dbClient.connect();
       return pongoClient;
     },
-    close: () => dbClient.close(),
-    db: (dbName?: string): PongoDb =>
-      dbName ? getDbClient(connectionString, dbName) : dbClient,
+    close: async () => {
+      for (const db of dbClients.values()) {
+        await db.close();
+      }
+    },
+    db: (dbName?: string): PongoDb => {
+      if (!dbName) return dbClient;
+
+      return (
+        dbClients.get(dbName) ??
+        dbClients
+          .set(dbName, getDbClient(connectionString, dbName))
+          .get(dbName)!
+      );
+    },
   };
 
   return pongoClient;

src/packages/pongo/src/main/client.ts

@@ -10,7 +10,7 @@ export const postgresClient = (
 
   return {
     connect: () => Promise.resolve(),
-    close: () => endPool(connectionString),
+    close: () => endPool({ connectionString, database }),
     collection: <T>(name: string) => postgresCollection<T>(name, pool),
   };
 };