Ensured that there's a single active transaction within the connection

Author: oskardudycz

src/packages/dumbo/src/core/connections/connection.ts

@@ -45,7 +45,10 @@ export type CreateConnectionOptions<
   close: (client: DbClient) => Promise<void>;
   initTransaction: (
     connection: () => ConnectionType,
-  ) => (client: Promise<DbClient>) => DatabaseTransaction<Connector, DbClient>;
+  ) => (
+    client: Promise<DbClient>,
+    options?: { close: (client: DbClient, error?: unknown) => Promise<void> },
+  ) => DatabaseTransaction<Connector, DbClient>;
   executor: () => Executor;
 };
 

src/packages/dumbo/src/core/connections/transaction.ts

@@ -70,12 +70,27 @@ export const transactionFactoryWithDbClient = <
   connect: () => Promise<DbClient>,
   initTransaction: (
     client: Promise<DbClient>,
+    options?: { close: (client: DbClient, error?: unknown) => Promise<void> },
   ) => DatabaseTransaction<Connector, DbClient>,
-): DatabaseTransactionFactory<Connector, DbClient> => ({
-  transaction: () => initTransaction(connect()),
-  withTransaction: (handle) =>
-    executeInTransaction(initTransaction(connect()), handle),
-});
+): DatabaseTransactionFactory<Connector, DbClient> => {
+  let currentTransaction: DatabaseTransaction<Connector, DbClient> | undefined =
+    undefined;
+
+  const getOrInitCurrentTransaction = () =>
+    currentTransaction ??
+    (currentTransaction = initTransaction(connect(), {
+      close: () => {
+        currentTransaction = undefined;
+        return Promise.resolve();
+      },
+    }));
+
+  return {
+    transaction: getOrInitCurrentTransaction,
+    withTransaction: (handle) =>
+      executeInTransaction(getOrInitCurrentTransaction(), handle),
+  };
+};
 
 const wrapInConnectionClosure = async <
   ConnectionType extends Connection = Connection,

src/packages/dumbo/src/storage/postgresql/pg/connections/transaction.ts

@@ -30,15 +30,19 @@ export const nodePostgresTransaction =
     commit: async () => {
       const client = await getClient;
 
-      await client.query('COMMIT');
-
-      if (options?.close) await options?.close(client);
+      try {
+        await client.query('COMMIT');
+      } finally {
+        if (options?.close) await options?.close(client);
+      }
     },
     rollback: async (error?: unknown) => {
       const client = await getClient;
-      await client.query('ROLLBACK');
-
-      if (options?.close) await options?.close(client, error);
+      try {
+        await client.query('ROLLBACK');
+      } finally {
+        if (options?.close) await options?.close(client, error);
+      }
     },
     execute: sqlExecutor(nodePostgresSQLExecutor(), {
       connect: () => getClient,

src/packages/dumbo/src/storage/sqlite/core/connections/index.ts

@@ -18,7 +18,6 @@ export type SQLiteClient = {
 };
 
 export type SQLitePoolClient = {
-  close(): Promise<void>;
   release: () => void;
   command: (sql: string, values?: Parameters[]) => Promise<void>;
   query: <T>(sql: string, values?: Parameters[]) => Promise<T[]>;
@@ -47,20 +46,20 @@ export type SQLitePoolConnectionOptions<
   ConnectorType extends SQLiteConnectorType = SQLiteConnectorType,
 > = {
   connector: ConnectorType;
-  allowNestedTransactions: boolean;
   type: 'PoolClient';
   connect: Promise<SQLitePoolClient>;
   close: (client: SQLitePoolClient) => Promise<void>;
+  allowNestedTransactions: boolean;
 };
 
 export type SQLiteClientConnectionOptions<
   ConnectorType extends SQLiteConnectorType = SQLiteConnectorType,
 > = {
   connector: ConnectorType;
-  allowNestedTransactions: boolean;
   type: 'Client';
   connect: Promise<SQLiteClient>;
   close: (client: SQLiteClient) => Promise<void>;
+  allowNestedTransactions: boolean;
 };
 
 export type SQLiteClientConnection<
@@ -84,20 +83,12 @@ export const sqliteClientConnection = <
 ): SQLiteClientConnection<ConnectorType> => {
   const { connect, close, allowNestedTransactions } = options;
 
-  const transactionCounter = transactionNestingCounter();
-
   return createConnection({
     connector: options.connector,
     connect,
     close,
-    initTransaction: (connection) => {
-      return sqliteTransaction(
-        options.connector,
-        connection,
-        transactionCounter,
-        allowNestedTransactions,
-      );
-    },
+    initTransaction: (connection) =>
+      sqliteTransaction(options.connector, connection, allowNestedTransactions),
     executor: () => sqliteSQLExecutor(options.connector),
   });
 };
@@ -139,7 +130,6 @@ export const sqlitePoolClientConnection = <
 ): SQLitePoolClientConnection<ConnectorType> => {
   const { connect, close, allowNestedTransactions } = options;
 
-  const transactionCounter = transactionNestingCounter();
   return createConnection({
     connector: options.connector,
     connect,
@@ -148,7 +138,6 @@ export const sqlitePoolClientConnection = <
       sqliteTransaction(
         options.connector,
         connection,
-        transactionCounter,
         allowNestedTransactions ?? false,
       ),
     executor: () => sqliteSQLExecutor(options.connector),

src/packages/dumbo/src/storage/sqlite/core/transactions/index.ts

@@ -5,9 +5,9 @@ import {
   type DatabaseTransaction,
 } from '../../../../core';
 import { sqliteSQLExecutor } from '../../core/execute';
-import type {
-  SQLiteClientOrPoolClient,
-  TransactionNestingCounter,
+import {
+  transactionNestingCounter,
+  type SQLiteClientOrPoolClient,
 } from '../connections';
 
 export type SQLiteTransaction<
@@ -21,64 +21,71 @@ export const sqliteTransaction =
   >(
     connector: ConnectorType,
     connection: () => Connection<ConnectorType, DbClient>,
-    transactionCounter: TransactionNestingCounter,
     allowNestedTransactions: boolean,
   ) =>
   (
     getClient: Promise<DbClient>,
     options?: { close: (client: DbClient, error?: unknown) => Promise<void> },
-  ): DatabaseTransaction<ConnectorType, DbClient> => ({
-    connection: connection(),
-    connector,
-    begin: async function () {
-      const client = await getClient;
+  ): DatabaseTransaction<ConnectorType, DbClient> => {
+    const transactionCounter = transactionNestingCounter();
+    return {
+      connection: connection(),
+      connector,
+      begin: async function () {
+        const client = await getClient;
+
+        if (allowNestedTransactions) {
+          if (transactionCounter.level >= 1) {
+            transactionCounter.increment();
+            await client.query(
+              `SAVEPOINT transaction${transactionCounter.level}`,
+            );
+            return;
+          }
 
-      if (allowNestedTransactions) {
-        if (transactionCounter.level >= 1) {
           transactionCounter.increment();
-          await client.query(
-            `SAVEPOINT transaction${transactionCounter.level}`,
-          );
-          return;
         }
 
-        transactionCounter.increment();
-      }
+        await client.query('BEGIN TRANSACTION');
+      },
+      commit: async function () {
+        const client = await getClient;
 
-      await client.query('BEGIN TRANSACTION');
-    },
-    commit: async function () {
-      const client = await getClient;
+        try {
+          if (allowNestedTransactions) {
+            if (transactionCounter.level > 1) {
+              await client.query(
+                `RELEASE transaction${transactionCounter.level}`,
+              );
+              transactionCounter.decrement();
 
-      if (allowNestedTransactions) {
-        if (transactionCounter.level > 1) {
-          await client.query(`RELEASE transaction${transactionCounter.level}`);
-          transactionCounter.decrement();
+              return;
+            }
 
-          return;
+            transactionCounter.reset();
+          }
+          await client.query('COMMIT');
+        } finally {
+          if (options?.close) await options?.close(client);
         }
+      },
+      rollback: async function (error?: unknown) {
+        const client = await getClient;
+        try {
+          if (allowNestedTransactions) {
+            if (transactionCounter.level > 1) {
+              transactionCounter.decrement();
+              return;
+            }
+          }
 
-        transactionCounter.reset();
-      }
-      await client.query('COMMIT');
-
-      if (options?.close) await options?.close(client);
-    },
-    rollback: async function (error?: unknown) {
-      const client = await getClient;
-
-      if (allowNestedTransactions) {
-        if (transactionCounter.level > 1) {
-          transactionCounter.decrement();
-          return;
+          await client.query('ROLLBACK');
+        } finally {
+          if (options?.close) await options?.close(client, error);
         }
-      }
-
-      await client.query('ROLLBACK');
-
-      if (options?.close) await options?.close(client, error);
-    },
-    execute: sqlExecutor(sqliteSQLExecutor(connector), {
-      connect: () => getClient,
-    }),
-  });
+      },
+      execute: sqlExecutor(sqliteSQLExecutor(connector), {
+        connect: () => getClient,
+      }),
+    };
+  };