From c94f3f89e313e881a4686201a83806e92b16d823 Mon Sep 17 00:00:00 2001
From: Oskar Dudycz <oskar.dudycz@gmail.com>
Date: Thu, 8 Aug 2024 10:28:27 +0200
Subject: [PATCH] Adjusted transaction typing in Dumbo to include DBClient type

Thanks to that, transaction's connection will be correctly resolved
---
 .../dumbo/src/core/connections/connection.ts       |  4 +++-
 .../dumbo/src/core/connections/transaction.ts      | 14 ++++++++------
 src/packages/dumbo/src/postgres/index.ts           |  2 ++
 3 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/src/packages/dumbo/src/core/connections/connection.ts b/src/packages/dumbo/src/core/connections/connection.ts
index c4021e7..b198daf 100644
--- a/src/packages/dumbo/src/core/connections/connection.ts
+++ b/src/packages/dumbo/src/core/connections/connection.ts
@@ -43,7 +43,9 @@ export type CreateConnectionOptions<
   close: (client: DbClient) => Promise<void>;
   initTransaction: (
     connection: () => ConnectionType,
-  ) => (client: Promise<DbClient>) => DatabaseTransaction<ConnectorType>;
+  ) => (
+    client: Promise<DbClient>,
+  ) => DatabaseTransaction<ConnectorType, DbClient>;
   executor: () => Executor;
 };
 
diff --git a/src/packages/dumbo/src/core/connections/transaction.ts b/src/packages/dumbo/src/core/connections/transaction.ts
index 758af1c..e69e550 100644
--- a/src/packages/dumbo/src/core/connections/transaction.ts
+++ b/src/packages/dumbo/src/core/connections/transaction.ts
@@ -14,12 +14,13 @@ export interface DatabaseTransaction<
 
 export interface DatabaseTransactionFactory<
   ConnectorType extends string = string,
+  DbClient = unknown,
 > {
-  transaction: () => DatabaseTransaction<ConnectorType>;
+  transaction: () => DatabaseTransaction<ConnectorType, DbClient>;
 
   withTransaction: <Result = never>(
     handle: (
-      transaction: DatabaseTransaction<ConnectorType>,
+      transaction: DatabaseTransaction<ConnectorType, DbClient>,
     ) => Promise<TransactionResult<Result> | Result>,
   ) => Promise<Result>;
 }
@@ -38,11 +39,12 @@ const toTransactionResult = <Result>(
 
 export const executeInTransaction = async <
   ConnectorType extends string = string,
+  DbClient = unknown,
   Result = void,
 >(
-  transaction: DatabaseTransaction<ConnectorType>,
+  transaction: DatabaseTransaction<ConnectorType, DbClient>,
   handle: (
-    transaction: DatabaseTransaction<ConnectorType>,
+    transaction: DatabaseTransaction<ConnectorType, DbClient>,
   ) => Promise<TransactionResult<Result> | Result>,
 ): Promise<Result> => {
   await transaction.begin();
@@ -67,8 +69,8 @@ export const transactionFactoryWithDbClient = <
   connect: () => Promise<DbClient>,
   initTransaction: (
     client: Promise<DbClient>,
-  ) => DatabaseTransaction<ConnectorType>,
-): DatabaseTransactionFactory<ConnectorType> => ({
+  ) => DatabaseTransaction<ConnectorType, DbClient>,
+): DatabaseTransactionFactory<ConnectorType, DbClient> => ({
   transaction: () => initTransaction(connect()),
   withTransaction: (handle) =>
     executeInTransaction(initTransaction(connect()), handle),
diff --git a/src/packages/dumbo/src/postgres/index.ts b/src/packages/dumbo/src/postgres/index.ts
index 6c1dbdf..2b5fbf4 100644
--- a/src/packages/dumbo/src/postgres/index.ts
+++ b/src/packages/dumbo/src/postgres/index.ts
@@ -1,6 +1,7 @@
 export * from './core';
 export * from './pg';
 import {
+  type NodePostgresConnection,
   type NodePostgresConnector,
   type NodePostgresPool,
   type NodePostgresPoolOptions,
@@ -10,4 +11,5 @@ import {
 export type PostgresConnector = NodePostgresConnector;
 export type PostgresPoolOptions = NodePostgresPoolOptions;
 export type PostgresPool = NodePostgresPool;
+export type PostgresConnection = NodePostgresConnection;
 export const postgresPool = nodePostgresPool;