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 = (db: Knex) => S; export type ServiceFactory = ( config: IUnleashConfig, ) => DeferredServiceFactory; + export type WithTransactional = S & { transactional: (fn: (service: S) => R) => Promise; }; +export type WithRollback = S & { + rollback: (fn: (service: S) => R) => Promise; +}; + /** * @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( return service; } +export function withRollback( + serviceFactory: (db: Knex) => S, + db: Knex, +): WithRollback { + const service = serviceFactory(db) as WithRollback; + + service.rollback = async (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(service: S): WithTransactional { const serviceWithFakeTransactional = service as WithTransactional;