@fragno-dev/db 0.2.0 → 0.2.1

package/src/query/unit-of-work/execute-unit-of-work.ts

@@ -309,13 +309,11 @@ export async function executeUnitOfWork<
         return { success: false, reason: "conflict" };
       }
 
-      // Wait before retrying
       const delayMs = retryPolicy.getDelayMs(attempt);
       if (delayMs > 0) {
         await new Promise((resolve) => setTimeout(resolve, delayMs));
       }
 
-      // Increment attempt counter for next iteration
       attempt++;
     } catch (error) {
       // An error was thrown during execution
@@ -356,6 +354,57 @@ export interface ExecuteRestrictedUnitOfWorkOptions {
   onSuccess?: (uow: IUnitOfWork) => Promise<void>;
 }
 
+/**
+ * Context provided to handler tx callbacks
+ */
+export interface TxPhaseContext<THooks extends HooksMap> {
+  /**
+   * Get a typed Unit of Work for the given schema
+   */
+  forSchema: <S extends AnySchema, H extends HooksMap = THooks>(
+    schema: S,
+    hooks?: H,
+  ) => TypedUnitOfWork<S, [], unknown, H>;
+}
+
+/**
+ * Handler callbacks for tx() - SYNCHRONOUS ONLY (no Promise return allowed)
+ * This prevents accidentally awaiting services in the wrong place
+ */
+export interface HandlerTxCallbacks<TRetrieveResult, TMutationResult, THooks extends HooksMap> {
+  /**
+   * Retrieval phase callback - schedules retrievals and optionally calls services
+   * Must be synchronous - cannot await promises
+   */
+  retrieve?: (context: TxPhaseContext<THooks>) => TRetrieveResult;
+  /**
+   * Mutation phase callback - receives retrieve result, schedules mutations
+   * Must be synchronous - cannot await promises (but may return a promise to be awaited)
+   */
+  mutate?: (context: TxPhaseContext<THooks>, retrieveResult: TRetrieveResult) => TMutationResult;
+}
+
+export interface ServiceTxCallbacks<
+  TSchema extends AnySchema,
+  TRetrievalResults extends unknown[],
+  TMutationResult,
+  THooks extends HooksMap,
+> {
+  /**
+   * Retrieval phase callback - schedules retrievals, returns typed UOW
+   */
+  retrieve?: (
+    uow: TypedUnitOfWork<TSchema, [], unknown, THooks>,
+  ) => TypedUnitOfWork<TSchema, TRetrievalResults, unknown, THooks>;
+  /**
+   * Mutation phase callback - receives retrieval results, schedules mutations and hooks
+   */
+  mutate?: (
+    uow: TypedUnitOfWork<TSchema, TRetrievalResults, unknown, THooks>,
+    results: TRetrievalResults,
+  ) => TMutationResult | Promise<TMutationResult>;
+}
+
 /**
  * Execute a Unit of Work with explicit phase control and automatic retry support.
  *
@@ -491,14 +540,273 @@ export async function executeRestrictedUnitOfWork<TResult, THooks extends HooksM
         });
       }
 
-      // Wait before retrying
       const delayMs = retryPolicy.getDelayMs(attempt);
       if (delayMs > 0) {
         await new Promise((resolve) => setTimeout(resolve, delayMs));
       }
 
-      // Increment attempt counter for next iteration
       attempt++;
     }
   }
 }
+
+/**
+ * Execute a transaction with array syntax (handler context).
+ * Takes a factory function that creates an array of service promises, enabling proper retry support.
+ *
+ * @param servicesFactory - Function that creates an array of service promises
+ * @param options - Configuration including UOW factory, retry policy, and abort signal
+ * @returns Promise resolving to array of awaited service results
+ *
+ * @example
+ * ```ts
+ * const [result1, result2] = await executeTxArray(
+ *   () => [
+ *     executeServiceTx(schema, callbacks1, uow),
+ *     executeServiceTx(schema, callbacks2, uow)
+ *   ],
+ *   { createUnitOfWork }
+ * );
+ * ```
+ */
+export async function executeTxArray<T extends readonly unknown[]>(
+  servicesFactory: () => readonly [...{ [K in keyof T]: Promise<T[K]> }],
+  options: ExecuteRestrictedUnitOfWorkOptions,
+): Promise<{ [K in keyof T]: T[K] }> {
+  const retryPolicy =
+    options.retryPolicy ??
+    new ExponentialBackoffRetryPolicy({
+      maxRetries: 5,
+      initialDelayMs: 10,
+      maxDelayMs: 100,
+    });
+  const signal = options.signal;
+  let attempt = 0;
+
+  while (true) {
+    // Check if aborted before starting attempt
+    if (signal?.aborted) {
+      throw new Error("Unit of Work execution aborted");
+    }
+
+    try {
+      // Create a fresh UOW for this attempt
+      const baseUow = options.createUnitOfWork();
+
+      // Call factory to create fresh service promises for this attempt
+      const services = servicesFactory();
+
+      await baseUow.executeRetrieve();
+
+      if (options.onBeforeMutate) {
+        options.onBeforeMutate(baseUow);
+      }
+
+      const result = await baseUow.executeMutations();
+      if (!result.success) {
+        throw new ConcurrencyConflictError();
+      }
+
+      if (options.onSuccess) {
+        await options.onSuccess(baseUow);
+      }
+
+      // Now await all service promises - they should all resolve now that mutations executed
+      const results = await Promise.all(services);
+      return results as { [K in keyof T]: T[K] };
+    } catch (error) {
+      if (signal?.aborted) {
+        throw new Error("Unit of Work execution aborted");
+      }
+
+      // Only retry concurrency conflicts, not other errors
+      if (!(error instanceof ConcurrencyConflictError)) {
+        throw error;
+      }
+
+      if (!retryPolicy.shouldRetry(attempt, error, signal)) {
+        if (signal?.aborted) {
+          throw new Error("Unit of Work execution aborted");
+        }
+        throw new Error("Unit of Work execution failed: optimistic concurrency conflict", {
+          cause: error,
+        });
+      }
+
+      const delayMs = retryPolicy.getDelayMs(attempt);
+      if (delayMs > 0) {
+        await new Promise((resolve) => setTimeout(resolve, delayMs));
+      }
+
+      attempt++;
+    }
+  }
+}
+
+/**
+ * Execute a transaction with callback syntax (handler context).
+ * Callbacks are synchronous only to prevent accidentally awaiting services in wrong place.
+ *
+ * @param callbacks - Object containing retrieve and mutate callbacks
+ * @param options - Configuration including UOW factory, retry policy, and abort signal
+ * @returns Promise resolving to the mutation result with promises awaited 1 level deep
+ */
+export async function executeTxCallbacks<
+  TRetrieveResult,
+  TMutationResult,
+  THooks extends HooksMap = {},
+>(
+  callbacks: HandlerTxCallbacks<TRetrieveResult, TMutationResult, THooks>,
+  options: ExecuteRestrictedUnitOfWorkOptions,
+): Promise<AwaitedPromisesInObject<TMutationResult>> {
+  const retryPolicy =
+    options.retryPolicy ??
+    new ExponentialBackoffRetryPolicy({
+      maxRetries: 5,
+      initialDelayMs: 10,
+      maxDelayMs: 100,
+    });
+  const signal = options.signal;
+  let attempt = 0;
+
+  while (true) {
+    // Check if aborted before starting attempt
+    if (signal?.aborted) {
+      throw new Error("Unit of Work execution aborted");
+    }
+
+    try {
+      // Create a fresh UOW for this attempt
+      const baseUow = options.createUnitOfWork();
+
+      const context: TxPhaseContext<THooks> = {
+        forSchema: <S extends AnySchema, H extends HooksMap = THooks>(schema: S, hooks?: H) => {
+          return baseUow.forSchema(schema, hooks);
+        },
+      };
+
+      let retrieveResult: TRetrieveResult;
+      if (callbacks.retrieve) {
+        retrieveResult = callbacks.retrieve(context);
+      } else {
+        retrieveResult = undefined as TRetrieveResult;
+      }
+
+      await baseUow.executeRetrieve();
+
+      let mutationResult: TMutationResult;
+      if (callbacks.mutate) {
+        mutationResult = callbacks.mutate(context, retrieveResult);
+      } else {
+        mutationResult = retrieveResult as unknown as TMutationResult;
+      }
+
+      const awaitedMutationResult = await awaitPromisesInObject(mutationResult);
+
+      if (options.onBeforeMutate) {
+        options.onBeforeMutate(baseUow);
+      }
+
+      const result = await baseUow.executeMutations();
+      if (!result.success) {
+        throw new ConcurrencyConflictError();
+      }
+
+      if (options.onSuccess) {
+        await options.onSuccess(baseUow);
+      }
+
+      return awaitedMutationResult;
+    } catch (error) {
+      if (signal?.aborted) {
+        throw new Error("Unit of Work execution aborted");
+      }
+
+      // Only retry concurrency conflicts, not other errors
+      if (!(error instanceof ConcurrencyConflictError)) {
+        throw error;
+      }
+
+      if (!retryPolicy.shouldRetry(attempt, error, signal)) {
+        if (signal?.aborted) {
+          throw new Error("Unit of Work execution aborted");
+        }
+        throw new Error("Unit of Work execution failed: optimistic concurrency conflict", {
+          cause: error,
+        });
+      }
+
+      const delayMs = retryPolicy.getDelayMs(attempt);
+      if (delayMs > 0) {
+        await new Promise((resolve) => setTimeout(resolve, delayMs));
+      }
+
+      attempt++;
+    }
+  }
+}
+
+/**
+ * Execute a transaction for service context.
+ * Service callbacks can be async for ergonomic async work.
+ *
+ * @param schema - Schema to use for the transaction
+ * @param callbacks - Object containing retrieve and mutate callbacks
+ * @param baseUow - Base Unit of Work (restricted) to use
+ * @returns Promise resolving to the mutation result with promises awaited 1 level deep
+ */
+export async function executeServiceTx<
+  TSchema extends AnySchema,
+  TRetrievalResults extends unknown[],
+  TMutationResult,
+  THooks extends HooksMap,
+>(
+  schema: TSchema,
+  callbacks: ServiceTxCallbacks<TSchema, TRetrievalResults, TMutationResult, THooks>,
+  baseUow: IUnitOfWork,
+): Promise<AwaitedPromisesInObject<TMutationResult>> {
+  const typedUow = baseUow.restrict({ readyFor: "none" }).forSchema<TSchema, THooks>(schema);
+
+  let retrievalUow: TypedUnitOfWork<TSchema, TRetrievalResults, unknown, THooks>;
+  try {
+    if (callbacks.retrieve) {
+      retrievalUow = callbacks.retrieve(typedUow);
+    } else {
+      // Safe cast: when there's no retrieve callback, TRetrievalResults should be []
+      retrievalUow = typedUow as unknown as TypedUnitOfWork<
+        TSchema,
+        TRetrievalResults,
+        unknown,
+        THooks
+      >;
+    }
+  } catch (error) {
+    typedUow.signalReadyForRetrieval();
+    typedUow.signalReadyForMutation();
+    throw error;
+  }
+
+  typedUow.signalReadyForRetrieval();
+
+  // Safe cast: retrievalPhase returns the correct type based on the UOW's type parameters
+  const results = (await retrievalUow.retrievalPhase) as TRetrievalResults;
+
+  let mutationResult: TMutationResult;
+  try {
+    if (callbacks.mutate) {
+      mutationResult = await callbacks.mutate(retrievalUow, results);
+    } else {
+      // Safe cast: when there's no mutate callback, TMutationResult should be void
+      mutationResult = undefined as TMutationResult;
+    }
+  } catch (error) {
+    typedUow.signalReadyForMutation();
+    throw error;
+  }
+
+  typedUow.signalReadyForMutation();
+
+  await retrievalUow.mutationPhase;
+
+  return await awaitPromisesInObject(mutationResult);
+}

package/src/query/unit-of-work/unit-of-work-coordinator.test.ts

@@ -769,7 +769,7 @@ describe("UOW Coordinator - Parent-Child Execution", () => {
     expect(child1.nonce).not.toBe(child2.nonce);
   });
 
-  it.skip("should not cause unhandled rejection when service method awaits retrievalPhase and executeRetrieve fails", async () => {
+  it("should not cause unhandled rejection when service method awaits retrievalPhase and executeRetrieve fails", async () => {
     const testSchema = schema((s) =>
       s.addTable("settings", (t) =>
         t
@@ -828,4 +828,252 @@ describe("UOW Coordinator - Parent-Child Execution", () => {
 
     expect(await deferred.promise).toContain('relation "settings" does not exist');
   });
+
+  it("should allow child UOW to call getCreatedIds() after parent executes mutations", async () => {
+    const testSchema = schema((s) =>
+      s.addTable("products", (t) =>
+        t.addColumn("id", idColumn()).addColumn("name", "string").addColumn("price", "integer"),
+      ),
+    );
+
+    const executor = createMockExecutor();
+    const parentUow = createUnitOfWork(createCompiler(), executor, createMockDecoder());
+
+    // Service method that creates a product using a child UOW and returns the child
+    const createProduct = (name: string, price: number) => {
+      const childUow = parentUow.restrict();
+      const productId = childUow.forSchema(testSchema).create("products", { name, price });
+      // Return both the child UOW and the product ID reference
+      return { childUow, productId };
+    };
+
+    // Call service to create a product via child UOW
+    const { childUow, productId } = createProduct("Widget", 1999);
+
+    // Parent executes mutations
+    await parentUow.executeMutations();
+
+    // Child should be able to call getCreatedIds() after parent has executed
+    // This tests that child checks parent's state, not its own stale state
+    const createdIds = childUow.getCreatedIds();
+
+    expect(createdIds).toHaveLength(1);
+    expect(createdIds[0].externalId).toBe(productId.externalId);
+  });
+
+  it("should preserve internal IDs in child UOW when using two-phase pattern with mutationPhase await", async () => {
+    const testSchema = schema((s) =>
+      s.addTable("orders", (t) =>
+        t
+          .addColumn("id", idColumn())
+          .addColumn("customerId", "string")
+          .addColumn("total", "integer"),
+      ),
+    );
+
+    const executor = createMockExecutor();
+    const parentUow = createUnitOfWork(createCompiler(), executor, createMockDecoder());
+
+    // Service method that uses two-phase pattern (common with hooks/async operations)
+    // This simulates a service that creates a record and needs to return the internal ID
+    const createOrder = async (customerId: string, total: number) => {
+      const childUow = parentUow.restrict();
+      const typedUow = childUow.forSchema(testSchema);
+
+      const orderId = typedUow.create("orders", { customerId, total });
+
+      // Service awaits mutationPhase to coordinate with parent execution
+      await childUow.mutationPhase;
+
+      // After mutationPhase resolves, service should be able to get internal IDs
+      const createdIds = childUow.getCreatedIds();
+      const foundId = createdIds.find((id) => id.externalId === orderId.externalId);
+
+      return {
+        externalId: orderId.externalId,
+        internalId: foundId?.internalId,
+      };
+    };
+
+    // Handler orchestrates the service call and mutation execution
+    const handler = async () => {
+      const orderPromise = createOrder("customer-123", 9999);
+
+      // Execute mutations - this should resolve the service's mutationPhase await
+      await parentUow.executeMutations();
+
+      // Now the service can complete and return the result with internal ID
+      return await orderPromise;
+    };
+
+    const result = await handler();
+
+    // The key assertion: internal ID should be defined (not undefined)
+    // This tests that child UOW preserves shared reference to parent's createdInternalIds array
+    expect(result.internalId).toBeDefined();
+    expect(typeof result.internalId).toBe("bigint");
+    expect(result.internalId).toBeGreaterThan(0n);
+  });
+
+  it("should fail when handler executes mutations before service finishes scheduling them (anti-pattern)", async () => {
+    const testSchema = schema((s) =>
+      s.addTable("totp_secret", (t) =>
+        t
+          .addColumn("id", idColumn())
+          .addColumn("userId", "string")
+          .addColumn("secret", "string")
+          .addColumn("backupCodes", "string")
+          .createIndex("idx_totp_user", ["userId"]),
+      ),
+    );
+
+    // Create executor that returns empty results (no existing record)
+    const customExecutor: UOWExecutor<CompiledQuery, unknown> = {
+      executeRetrievalPhase: async (queries: CompiledQuery[]) => {
+        // Return empty array for each query (no existing records)
+        return queries.map(() => []);
+      },
+      executeMutationPhase: async (mutations: CompiledMutation<CompiledQuery>[]) => {
+        return {
+          success: true,
+          createdInternalIds: mutations.map(() => BigInt(Math.floor(Math.random() * 1000))),
+        };
+      },
+    };
+
+    const parentUow = createUnitOfWork(createCompiler(), customExecutor, createMockDecoder());
+
+    // Service that has async work AFTER awaiting retrievalPhase but BEFORE scheduling mutations
+    // This is a problematic pattern that can lead to race conditions
+    const enableTotp = async (userId: string) => {
+      const childUow = parentUow.restrict();
+      const typedUow = childUow
+        .forSchema(testSchema)
+        .findFirst("totp_secret", (b) =>
+          b.whereIndex("idx_totp_user", (eb) => eb("userId", "=", userId)),
+        );
+
+      // Service awaits retrieval phase
+      const [existing] = await typedUow.retrievalPhase;
+
+      if (existing) {
+        throw new Error("TOTP already enabled");
+      }
+
+      // Simulate async work (like hashing backup codes) that yields control
+      await new Promise((resolve) => setTimeout(resolve, 10));
+
+      // By the time we get here, if the handler has already called executeMutate(),
+      // the UOW will be in "executed" state and this will fail
+      typedUow.create("totp_secret", {
+        userId,
+        secret: "TESTSECRET123",
+        backupCodes: "[]",
+      });
+
+      await typedUow.mutationPhase;
+
+      return { success: true };
+    };
+
+    // ANTI-PATTERN: Handler executes both phases immediately without awaiting service
+    const badHandler = async () => {
+      // Call service - returns promise immediately
+      const resultPromise = enableTotp("user-123");
+
+      // Execute retrieval phase
+      await parentUow.executeRetrieve();
+
+      // Execute mutation phase BEFORE service has finished scheduling mutations
+      // This is the bug - service is still running async work and hasn't scheduled mutations yet
+      await parentUow.executeMutations();
+
+      // Now await service promise - but it's too late, UOW is already in "executed" state
+      return await resultPromise;
+    };
+
+    // This should throw "Cannot add mutation operation in executed state"
+    await expect(badHandler()).rejects.toThrow("Cannot add mutation operation in executed state");
+  });
+
+  it("should succeed when handler awaits service promise between phases (correct pattern)", async () => {
+    const testSchema = schema((s) =>
+      s.addTable("totp_secret", (t) =>
+        t
+          .addColumn("id", idColumn())
+          .addColumn("userId", "string")
+          .addColumn("secret", "string")
+          .addColumn("backupCodes", "string")
+          .createIndex("idx_totp_user", ["userId"]),
+      ),
+    );
+
+    // Create executor that returns empty results (no existing record)
+    const customExecutor: UOWExecutor<CompiledQuery, unknown> = {
+      executeRetrievalPhase: async (queries: CompiledQuery[]) => {
+        return queries.map(() => []);
+      },
+      executeMutationPhase: async (mutations: CompiledMutation<CompiledQuery>[]) => {
+        return {
+          success: true,
+          createdInternalIds: mutations.map(() => BigInt(Math.floor(Math.random() * 1000))),
+        };
+      },
+    };
+
+    const parentUow = createUnitOfWork(createCompiler(), customExecutor, createMockDecoder());
+
+    // Same service as before - has async work between retrieval and mutation scheduling
+    const enableTotp = async (userId: string) => {
+      const childUow = parentUow.restrict();
+      const typedUow = childUow
+        .forSchema(testSchema)
+        .findFirst("totp_secret", (b) =>
+          b.whereIndex("idx_totp_user", (eb) => eb("userId", "=", userId)),
+        );
+
+      const [existing] = await typedUow.retrievalPhase;
+
+      if (existing) {
+        throw new Error("TOTP already enabled");
+      }
+
+      // Simulate async work (like hashing backup codes)
+      await new Promise((resolve) => setTimeout(resolve, 10));
+
+      // Schedule mutation - this will work because handler waits for service to finish
+      typedUow.create("totp_secret", {
+        userId,
+        secret: "TESTSECRET123",
+        backupCodes: "[]",
+      });
+
+      return { success: true };
+    };
+
+    // CORRECT PATTERN: Handler awaits service promise between phase executions
+    const goodHandler = async () => {
+      // Call service first - it schedules retrieval operations synchronously
+      const resultPromise = enableTotp("user-123");
+
+      // Execute retrieval phase - service can now continue past its retrieval await
+      await parentUow.executeRetrieve();
+
+      // Wait for service to finish - it will schedule mutations
+      const result = await resultPromise;
+
+      // Execute mutations that service scheduled
+      await parentUow.executeMutations();
+
+      return result;
+    };
+
+    // This should succeed without errors
+    const result = await goodHandler();
+    expect(result.success).toBe(true);
+
+    // Verify operations were registered
+    expect(parentUow.getRetrievalOperations()).toHaveLength(1);
+    expect(parentUow.getMutationOperations()).toHaveLength(1);
+  });
 });