From de94c41a59579c91b95e81bcc9907263942bbe0e Mon Sep 17 00:00:00 2001
From: Oskar Dudycz <oskar.dudycz@gmail.com>
Date: Fri, 12 Jul 2024 12:43:52 +0200
Subject: [PATCH] Added option to inject PoolClient to PongoClient

That will enable running transactions and sharing connection with Emmett projections
---
 src/packages/pongo/src/main/dbClient.ts       |  9 +++-----
 src/packages/pongo/src/main/index.ts          |  2 +-
 .../src/main/{client.ts => pongoClient.ts}    |  4 ++--
 src/packages/pongo/src/postgres/client.ts     | 23 +++++++++++++------
 .../pongo/src/postgres/postgresCollection.ts  |  2 +-
 5 files changed, 23 insertions(+), 17 deletions(-)
 rename src/packages/pongo/src/main/{client.ts => pongoClient.ts} (87%)

diff --git a/src/packages/pongo/src/main/dbClient.ts b/src/packages/pongo/src/main/dbClient.ts
index 6e79ade..6cf70f9 100644
--- a/src/packages/pongo/src/main/dbClient.ts
+++ b/src/packages/pongo/src/main/dbClient.ts
@@ -1,4 +1,4 @@
-import { postgresClient } from '../postgres';
+import { postgresClient, type PongoClientOptions } from '../postgres';
 import type { PongoCollection } from './typing/operations';
 
 export interface DbClient {
@@ -7,10 +7,7 @@ export interface DbClient {
   collection: <T>(name: string) => PongoCollection<T>;
 }
 
-export const getDbClient = (
-  connectionString: string,
-  database?: string,
-): DbClient => {
+export const getDbClient = (options: PongoClientOptions): DbClient => {
   // This is the place where in the future could come resolution of other database types
-  return postgresClient(connectionString, database);
+  return postgresClient(options);
 };
diff --git a/src/packages/pongo/src/main/index.ts b/src/packages/pongo/src/main/index.ts
index cbc1b7c..a88ac3a 100644
--- a/src/packages/pongo/src/main/index.ts
+++ b/src/packages/pongo/src/main/index.ts
@@ -1,3 +1,3 @@
-export * from './client';
 export * from './dbClient';
+export * from './pongoClient';
 export * from './typing';
diff --git a/src/packages/pongo/src/main/client.ts b/src/packages/pongo/src/main/pongoClient.ts
similarity index 87%
rename from src/packages/pongo/src/main/client.ts
rename to src/packages/pongo/src/main/pongoClient.ts
index 9fbf2c1..e415423 100644
--- a/src/packages/pongo/src/main/client.ts
+++ b/src/packages/pongo/src/main/pongoClient.ts
@@ -6,7 +6,7 @@ export const pongoClient = (connectionString: string): PongoClient => {
   const defaultDbName = getDatabaseNameOrDefault(connectionString);
   const dbClients: Map<string, DbClient> = new Map();
 
-  const dbClient = getDbClient(connectionString);
+  const dbClient = getDbClient({ connectionString });
   dbClients.set(defaultDbName, dbClient);
 
   const pongoClient: PongoClient = {
@@ -25,7 +25,7 @@ export const pongoClient = (connectionString: string): PongoClient => {
       return (
         dbClients.get(dbName) ??
         dbClients
-          .set(dbName, getDbClient(connectionString, dbName))
+          .set(dbName, getDbClient({ connectionString, database: dbName }))
           .get(dbName)!
       );
     },
diff --git a/src/packages/pongo/src/postgres/client.ts b/src/packages/pongo/src/postgres/client.ts
index 23b80d3..a7d1fdf 100644
--- a/src/packages/pongo/src/postgres/client.ts
+++ b/src/packages/pongo/src/postgres/client.ts
@@ -1,16 +1,25 @@
 import { endPool, getPool } from '@event-driven-io/dumbo';
+import pg from 'pg';
 import { type DbClient } from '../main';
 import { postgresCollection } from './postgresCollection';
 
-export const postgresClient = (
-  connectionString: string,
-  database?: string,
-): DbClient => {
-  const pool = getPool({ connectionString, database });
+export type PongoClientOptions = {
+  connectionString: string;
+  database?: string | undefined;
+  client?: pg.PoolClient;
+};
+
+export const postgresClient = (options: PongoClientOptions): DbClient => {
+  const { connectionString, database, client } = options;
+  const managesPoolLifetime = !client;
+  const clientOrPool = client ?? getPool({ connectionString, database });
 
   return {
     connect: () => Promise.resolve(),
-    close: () => endPool({ connectionString, database }),
-    collection: <T>(name: string) => postgresCollection<T>(name, pool),
+    close: () =>
+      managesPoolLifetime
+        ? endPool({ connectionString, database })
+        : Promise.resolve(),
+    collection: <T>(name: string) => postgresCollection<T>(name, clientOrPool),
   };
 };
diff --git a/src/packages/pongo/src/postgres/postgresCollection.ts b/src/packages/pongo/src/postgres/postgresCollection.ts
index a93c76f..1dc5538 100644
--- a/src/packages/pongo/src/postgres/postgresCollection.ts
+++ b/src/packages/pongo/src/postgres/postgresCollection.ts
@@ -17,7 +17,7 @@ import { buildUpdateQuery } from './update';
 
 export const postgresCollection = <T>(
   collectionName: string,
-  pool: pg.Pool,
+  pool: pg.Pool | pg.PoolClient,
 ): PongoCollection<T> => {
   const execute = (sql: SQL) => executeSQL(pool, sql);
   const SqlFor = collectionSQLBuilder(collectionName);