From 6170d10e626ebfed2fbaa763bfe5db26242054a6 Mon Sep 17 00:00:00 2001
From: Mateusz Kwasniewski <kwasniewski.mateusz@gmail.com>
Date: Wed, 31 Jul 2024 10:22:05 +0200
Subject: [PATCH] feat: rollback transaction wrapper (#7706)

---
 src/lib/db/transaction.ts | 24 ++++++++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/src/lib/db/transaction.ts b/src/lib/db/transaction.ts
index 2a7988d75d..0fa5905ac7 100644
--- a/src/lib/db/transaction.ts
+++ b/src/lib/db/transaction.ts
@@ -36,10 +36,15 @@ export type DeferredServiceFactory<S> = (db: Knex) => S;
 export type ServiceFactory<S> = (
     config: IUnleashConfig,
 ) => DeferredServiceFactory<S>;
+
 export type WithTransactional<S> = S & {
     transactional: <R>(fn: (service: S) => R) => Promise<R>;
 };
 
+export type WithRollback<S> = S & {
+    rollback: <R>(fn: (service: S) => R) => Promise<R>;
+};
+
 /**
  * @deprecated this is a temporary solution to deal with transactions at the store level.
  * Ideally, we should handle transactions at the service level (each service method should be transactional).
@@ -81,6 +86,25 @@ export function withTransactional<S>(
     return service;
 }
 
+export function withRollback<S>(
+    serviceFactory: (db: Knex) => S,
+    db: Knex,
+): WithRollback<S> {
+    const service = serviceFactory(db) as WithRollback<S>;
+
+    service.rollback = async <R>(fn: (service: S) => R) => {
+        const trx = await db.transaction({ isolationLevel: 'serializable' });
+        try {
+            const transactionService = serviceFactory(trx);
+            return fn(transactionService);
+        } finally {
+            await trx.rollback();
+        }
+    };
+
+    return service;
+}
+
 /** Just for testing purposes */
 export function withFakeTransactional<S>(service: S): WithTransactional<S> {
     const serviceWithFakeTransactional = service as WithTransactional<S>;