Make it a wrapper

src/lib/db/transaction.ts

@@ -1,6 +1,6 @@
 import type { Knex } from 'knex';
 import type { IUnleashConfig } from '../types/index.ts';
-import { transactionContext } from '../util/transactionContext.js';
+import { TransactionContext } from '../util/transactionContext.js';
 
 export type KnexTransaction = Knex.Transaction;
 
@@ -42,6 +42,12 @@ export type WithTransactional<S> = S & {
     transactional: <R>(fn: (service: S) => R) => Promise<R>;
 };
 
+export type WithTrackedTransactional<S> = S & {
+    trackedTransactional: <R>(
+        fn: (transactionContext: TransactionContext) => R,
+    ) => Promise<R>;
+};
+
 export type WithRollbackTransaction<S> = S & {
     rollbackTransaction: <R>(fn: (service: S) => R) => Promise<R>;
 };
@@ -76,12 +82,29 @@ export function withTransactional<S>(
 ): WithTransactional<S> {
     const service = serviceFactory(db) as WithTransactional<S>;
 
-    service.transactional = async <R>(fn: (service: S) => R): Promise<R> => {
+    service.transactional = async <R>(fn: (service: S) => R) =>
+        // Maybe: inTransaction(db, async (trx: Knex.Transaction) => fn(serviceFactory(trx)));
+        // this assumes that the caller didn't start a transaction already and opens a new one.
+        db.transaction(async (trx: Knex.Transaction) => {
+            const transactionalService = serviceFactory(trx);
+            return fn(transactionalService);
+        });
+
+    return service;
+}
+
+export function withContextualTransactional<S>(
+    serviceFactory: (db: Knex) => S,
+    db: Knex,
+): WithTrackedTransactional<S> {
+    const service = serviceFactory(db) as WithTrackedTransactional<S>;
+
+    service.trackedTransactional = async <R>(
+        fn: (transactionContext: TransactionContext) => R,
+    ): Promise<R> => {
         return db.transaction(async (trx) => {
-            return transactionContext.run(async () => {
+            const transactionContext = new TransactionContext(trx);
-                const transactionalService = serviceFactory(trx);
+            return fn(transactionContext);
-                return fn(transactionalService);
-            });
         });
     };
 

src/lib/util/transactionContext.test.ts

@@ -1,191 +1,108 @@
-import { transactionContext } from './transactionContext.js';
+import {
+    TransactionContext,
+    type OperationContext,
+} from './transactionContext.js';
+import { vi } from 'vitest';
 
-describe('transactionContext', () => {
+describe('TransactionContext', () => {
-    describe('run', () => {
+    let mockTransaction: any;
-        it('should execute callback with transaction context', async () => {
-            const result = await transactionContext.run(async () => {
-                return 'callback-result';
-            });
 
-            expect(result).toBe('callback-result');
+    beforeEach(() => {
-        });
+        mockTransaction = {
-
+            select: vi.fn(),
-        it('should make transaction context available inside the callback', async () => {
+            insert: vi.fn(),
-            await transactionContext.run(async () => {
+            update: vi.fn(),
-                expect(transactionContext.getOperationType()).toBe(
+            delete: vi.fn(),
-                    'transaction',
+            where: vi.fn(),
-                );
+            commit: vi.fn(),
-                expect(transactionContext.getOperationId()).toBeDefined();
+            rollback: vi.fn(),
-                expect(typeof transactionContext.getOperationId()).toBe(
+            isTransaction: true,
-                    'number',
+        };
-                );
-            });
-        });
-
-        it('should generate unique numeric transaction IDs', async () => {
-            const ids: (string | number | undefined)[] = [];
-
-            await transactionContext.run(async () => {
-                ids.push(transactionContext.getOperationId());
-            });
-
-            await transactionContext.run(async () => {
-                ids.push(transactionContext.getOperationId());
-            });
-
-            expect(ids).toHaveLength(2);
-            expect(ids[0]).not.toBe(ids[1]);
-            expect(typeof ids[0]).toBe('number');
-            expect(typeof ids[1]).toBe('number');
-        });
-
-        it('should handle rejected promises', async () => {
-            const error = new Error('Test error');
-
-            await expect(
-                transactionContext.run(async () => {
-                    throw error;
-                }),
-            ).rejects.toThrow('Test error');
-        });
-
-        it('should handle nested transaction contexts', async () => {
-            let outerOperationId: string | number | undefined;
-            let innerOperationId: string | number | undefined;
-
-            await transactionContext.run(async () => {
-                outerOperationId = transactionContext.getOperationId();
-                expect(transactionContext.getOperationType()).toBe(
-                    'transaction',
-                );
-
-                await transactionContext.run(async () => {
-                    innerOperationId = transactionContext.getOperationId();
-                    expect(transactionContext.getOperationType()).toBe(
-                        'transaction',
-                    );
-                });
-
-                expect(transactionContext.getOperationId()).toBe(
-                    outerOperationId,
-                );
-            });
-
-            expect(outerOperationId).toBeDefined();
-            expect(innerOperationId).toBeDefined();
-            expect(outerOperationId).not.toBe(innerOperationId);
-        });
     });
 
-    describe('getOperation', () => {
+    describe('constructor', () => {
-        it('should return undefined when called outside of transaction context', () => {
+        it('should create transaction context with default operation', () => {
-            expect(transactionContext.getOperation()).toBeUndefined();
+            const txContext = new TransactionContext(mockTransaction);
+
+            expect(txContext.operationContext.type).toBe('transaction');
+            expect(txContext.operationContext.id).toBeDefined();
+            expect(typeof txContext.operationContext.id).toBe('number');
+            expect(txContext.transaction).toBe(mockTransaction);
         });
 
-        it('should return the operation context when called inside context', async () => {
+        it('should create transaction context with custom operation', () => {
-            await transactionContext.run(async () => {
+            const customOperation: OperationContext = {
-                const operation = transactionContext.getOperation();
+                type: 'change-request',
-                expect(operation).toBeDefined();
+                id: 42,
-                expect(operation?.type).toBe('transaction');
+            };
-                expect(operation?.id).toBeDefined();
+
-                expect(typeof operation?.id).toBe('number');
+            const txContext = new TransactionContext(
+                mockTransaction,
+                customOperation,
+            );
+
+            expect(txContext.operationContext.type).toBe('change-request');
+            expect(txContext.operationContext.id).toBe(42);
+            expect(txContext.transaction).toBe(mockTransaction);
+        });
+
+        it('should create transaction context with partial operation context', () => {
+            const txContext = new TransactionContext(mockTransaction, {
+                type: 'change-request',
             });
-        });
-    });
 
-    describe('getOperationType', () => {
+            expect(txContext.operationContext.type).toBe('change-request');
-        it('should return undefined when called outside of transaction context', () => {
+            expect(txContext.operationContext.id).toBeDefined();
-            expect(transactionContext.getOperationType()).toBeUndefined();
+            expect(typeof txContext.operationContext.id).toBe('number');
+            expect(txContext.transaction).toBe(mockTransaction);
         });
 
-        it('should return "transaction" when called inside context', async () => {
+        it('should generate unique IDs for different contexts', () => {
-            await transactionContext.run(async () => {
+            const txContext1 = new TransactionContext(mockTransaction);
-                expect(transactionContext.getOperationType()).toBe(
+            const txContext2 = new TransactionContext(mockTransaction);
-                    'transaction',
-                );
-            });
-        });
-    });
 
-    describe('getOperationId', () => {
+            expect(txContext1.operationContext.id).not.toBe(
-        it('should return undefined when called outside of transaction context', () => {
+                txContext2.operationContext.id,
-            expect(transactionContext.getOperationId()).toBeUndefined();
-        });
-
-        it('should return numeric ID when called inside context', async () => {
-            await transactionContext.run(async () => {
-                const id = transactionContext.getOperationId();
-                expect(id).toBeDefined();
-                expect(typeof id).toBe('number');
-            });
-        });
-    });
-
-    describe('setOperation', () => {
-        it('should throw error when called outside of transaction context', () => {
-            expect(() => {
-                transactionContext.setOperation({
-                    type: 'change-request',
-                    id: 123,
-                });
-            }).toThrow(
-                'No active transaction context found when setting operation',
             );
         });
+    });
 
-        it('should set operation context in active context', async () => {
+    describe('operationContext', () => {
-            await transactionContext.run(async () => {
+        it('should allow updating operation context', () => {
-                expect(transactionContext.getOperationType()).toBe(
+            const txContext = new TransactionContext(mockTransaction);
-                    'transaction',
-                );
 
-                transactionContext.setOperation({
+            expect(txContext.operationContext.type).toBe('transaction');
-                    type: 'change-request',
-                    id: 456,
-                });
 
-                expect(transactionContext.getOperationType()).toBe(
+            txContext.operationContext = {
-                    'change-request',
+                type: 'change-request',
-                );
+                id: 123,
-                expect(transactionContext.getOperationId()).toBe(456);
+            };
-            });
+
+            expect(txContext.operationContext.type).toBe('change-request');
+            expect(txContext.operationContext.id).toBe(123);
         });
 
-        it('should update existing operation context', async () => {
+        it('should allow partial updates to operation context', () => {
-            await transactionContext.run(async () => {
+            const txContext = new TransactionContext(mockTransaction);
-                expect(transactionContext.getOperationType()).toBe(
+            const originalId = txContext.operationContext.id;
-                    'transaction',
-                );
-                const originalId = transactionContext.getOperationId();
 
-                transactionContext.setOperation({
+            Object.assign(txContext.operationContext, {
-                    type: 'change-request',
+                type: 'change-request',
-                    id: 789,
-                });
-
-                expect(transactionContext.getOperationType()).toBe(
-                    'change-request',
-                );
-                expect(transactionContext.getOperationId()).toBe(789);
-                expect(transactionContext.getOperationId()).not.toBe(
-                    originalId,
-                );
             });
+
+            expect(txContext.operationContext.type).toBe('change-request');
+            expect(txContext.operationContext.id).toBe(originalId);
         });
+    });
 
-        it('should not affect transaction context after callback completes', async () => {
+    describe('transaction access', () => {
-            await transactionContext.run(async () => {
+        it('should provide access to the underlying transaction', () => {
-                transactionContext.setOperation({
+            const txContext = new TransactionContext(mockTransaction);
-                    type: 'change-request',
-                    id: 999,
-                });
-                expect(transactionContext.getOperationType()).toBe(
-                    'change-request',
-                );
-                expect(transactionContext.getOperationId()).toBe(999);
-            });
 
-            expect(transactionContext.getOperation()).toBeUndefined();
+            expect(txContext.transaction).toBe(mockTransaction);
+            expect(txContext.transaction.select).toBe(mockTransaction.select);
+            expect(txContext.transaction.insert).toBe(mockTransaction.insert);
+            expect(txContext.transaction.commit).toBe(mockTransaction.commit);
         });
     });
 });

src/lib/util/transactionContext.ts

@@ -1,4 +1,4 @@
-import { AsyncLocalStorage } from 'async_hooks';
+import type { Knex } from 'knex';
 
 export interface OperationContext {
     type: 'change-request' | 'transaction';
@@ -11,38 +11,18 @@ function generateNumericTransactionId(): number {
     return timestamp * 1000 + random;
 }
 
-const storage = new AsyncLocalStorage<OperationContext>();
+export class TransactionContext {
+    public readonly transaction: Knex.Transaction;
+    public operationContext: OperationContext;
 
-export const transactionContext = {
+    constructor(
-    run<T>(callback: () => Promise<T>): Promise<T> {
+        transaction: Knex.Transaction,
-        const data: OperationContext = {
+        operationContext?: Partial<OperationContext>,
-            type: 'transaction',
+    ) {
-            id: generateNumericTransactionId(),
+        this.transaction = transaction;
+        this.operationContext = {
+            type: operationContext?.type || 'transaction',
+            id: operationContext?.id || generateNumericTransactionId(),
         };
-        return storage.run(data, callback) as Promise<T>;
+    }
-    },
+}
-
-    getOperation(): OperationContext | undefined {
-        return storage.getStore();
-    },
-
-    getOperationType(): OperationContext['type'] | undefined {
-        return storage.getStore()?.type;
-    },
-
-    getOperationId(): number | undefined {
-        return storage.getStore()?.id;
-    },
-
-    setOperation(operation: OperationContext): void {
-        const store = storage.getStore();
-        if (store) {
-            store.id = operation.id;
-            store.type = operation.type;
-        } else {
-            throw new Error(
-                'No active transaction context found when setting operation',
-            );
-        }
-    },
-};