@fragno-dev/db 0.1.15 → 0.2.1

package/src/db-fragment-definition-builder.ts

@@ -1,9 +1,13 @@
 import type { AnySchema } from "./schema/create";
-import type { AbstractQuery } from "./query/query";
+import type { SimpleQueryInterface } from "./query/simple-query-interface";
 import type { DatabaseAdapter } from "./adapters/adapters";
-import type { IUnitOfWork } from "./query/unit-of-work";
-import { TypedUnitOfWork, UnitOfWork } from "./query/unit-of-work";
-import type { RequestThisContext, FragnoPublicConfig } from "@fragno-dev/core";
+import type { IUnitOfWork } from "./query/unit-of-work/unit-of-work";
+import { TypedUnitOfWork, UnitOfWork } from "./query/unit-of-work/unit-of-work";
+import type {
+  RequestThisContext,
+  FragnoPublicConfig,
+  AnyFragnoInstantiatedFragment,
+} from "@fragno-dev/core";
 import {
   FragmentDefinitionBuilder,
   type FragmentDefinition,
@@ -11,9 +15,22 @@ import {
 } from "@fragno-dev/core";
 import {
   executeRestrictedUnitOfWork,
+  executeTxArray,
+  executeTxCallbacks,
+  executeServiceTx,
   type AwaitedPromisesInObject,
   type ExecuteRestrictedUnitOfWorkOptions,
-} from "./query/execute-unit-of-work";
+  type HandlerTxCallbacks,
+  type ServiceTxCallbacks,
+} from "./query/unit-of-work/execute-unit-of-work";
+import {
+  prepareHookMutations,
+  processHooks,
+  type HooksMap,
+  type HookFn,
+  type HookContext,
+} from "./hooks/hooks";
+import type { InternalFragmentInstance } from "./fragments/internal-fragment";
 
 /**
  * Extended FragnoPublicConfig that includes a database adapter.
@@ -32,7 +49,7 @@ export type ImplicitDatabaseDependencies<TSchema extends AnySchema> = {
   /**
    * Database query engine for the fragment's schema.
    */
-  db: AbstractQuery<TSchema>;
+  db: SimpleQueryInterface<TSchema>;
   /**
    * The schema definition for this fragment.
    */
@@ -50,19 +67,43 @@ export type ImplicitDatabaseDependencies<TSchema extends AnySchema> = {
 /**
  * Service context for database fragments - provides restricted UOW access without execute methods.
  */
-export type DatabaseServiceContext = RequestThisContext & {
+export type DatabaseServiceContext<THooks extends HooksMap> = RequestThisContext & {
   /**
    * Get a typed, restricted Unit of Work for the given schema.
    * @param schema - Schema to get a typed view for
    * @returns TypedUnitOfWork (restricted version without execute methods)
    */
-  uow<TSchema extends AnySchema>(schema: TSchema): TypedUnitOfWork<TSchema, [], unknown>;
+  uow<TSchema extends AnySchema>(schema: TSchema): TypedUnitOfWork<TSchema, [], unknown, THooks>;
+
+  /**
+   * Execute a transaction with two-phase callbacks (retrieve + mutate).
+   *
+   * @param schema - Schema to use for the transaction
+   * @param callbacks - Object containing retrieve and mutate callbacks
+   * @returns Promise resolving to the mutation result with promises awaited 1 level deep
+   *
+   * @example
+   * ```ts
+   * return this.tx(schema, {
+   *   retrieve: (uow) => uow.findFirst("users", ...),
+   *   mutate: async (uow, [user]) => {
+   *     await validateUser(user);
+   *     uow.update("users", user.id, ...).check();
+   *     return { ok: true };
+   *   }
+   * });
+   * ```
+   */
+  tx<TSchema extends AnySchema, TRetrievalResults extends unknown[], TMutationResult = void>(
+    schema: TSchema,
+    callbacks: ServiceTxCallbacks<TSchema, TRetrievalResults, TMutationResult, THooks>,
+  ): Promise<AwaitedPromisesInObject<TMutationResult>>;
 };
 
 /**
  * Handler context for database fragments - provides UOW execution with automatic retry support.
  */
-export type DatabaseHandlerContext = RequestThisContext & {
+export type DatabaseHandlerContext<THooks extends HooksMap = {}> = RequestThisContext & {
   /**
    * Execute a Unit of Work with explicit phase control and automatic retry support.
    * This method provides an API where users call forSchema to create a schema-specific
@@ -96,7 +137,10 @@ export type DatabaseHandlerContext = RequestThisContext & {
    */
   uow<TResult>(
     callback: (context: {
-      forSchema: <S extends AnySchema>(schema: S) => TypedUnitOfWork<S, [], unknown>;
+      forSchema: <TSchema extends AnySchema, H extends HooksMap = THooks>(
+        schema: TSchema,
+        hooks?: H,
+      ) => TypedUnitOfWork<TSchema, [], unknown, H>;
       executeRetrieve: () => Promise<void>;
       executeMutate: () => Promise<void>;
       nonce: string;
@@ -104,6 +148,49 @@ export type DatabaseHandlerContext = RequestThisContext & {
     }) => Promise<TResult> | TResult,
     options?: Omit<ExecuteRestrictedUnitOfWorkOptions, "createUnitOfWork">,
   ): Promise<AwaitedPromisesInObject<TResult>>;
+
+  /**
+   * Execute a transaction with automatic retry support.
+   * Provides two overloads: array syntax (common case) and callback syntax (advanced).
+   *
+   * Array syntax - pass a factory function that returns service calls:
+   * ```ts
+   * const [transfer, notify] = await this.tx(
+   *   () => [
+   *     services.transfer({ from, to, amount }),
+   *     services.notify({ userId, message })
+   *   ]
+   * );
+   * ```
+   *
+   * Callback syntax - for handlers that need direct UOW access:
+   * ```ts
+   * const result = await this.tx({
+   *   retrieve: ({ forSchema }) => {
+   *     const uow = forSchema(schema);
+   *     uow.find("users", ...);
+   *     return services.transfer({ ... });
+   *   },
+   *   mutate: ({ forSchema }, transferPromise) => {
+   *     const uow = forSchema(schema);
+   *     uow.create("auditLog", { ... });
+   *     return transferPromise;
+   *   }
+   * });
+   * ```
+   *
+   * Note: Handler callbacks are synchronous only to prevent accidentally awaiting services in wrong place.
+   */
+  // Overload 1: Array syntax (common case) - accepts a factory
+  tx<T extends readonly unknown[]>(
+    servicesFactory: () => readonly [...{ [K in keyof T]: Promise<T[K]> }],
+    options?: Omit<ExecuteRestrictedUnitOfWorkOptions, "createUnitOfWork">,
+  ): Promise<{ [K in keyof T]: T[K] }>;
+  // Overload 2: Callback syntax (advanced, sync callbacks only)
+  tx<TRetrieveResult, TMutationResult>(
+    callbacks: HandlerTxCallbacks<TRetrieveResult, TMutationResult, THooks>,
+    options?: Omit<ExecuteRestrictedUnitOfWorkOptions, "createUnitOfWork">,
+  ): Promise<AwaitedPromisesInObject<TMutationResult>>;
 };
 
 /**
@@ -117,7 +204,7 @@ export type DatabaseFragmentContext<TSchema extends AnySchema> = {
   /**
    * ORM query engine for the fragment's schema.
    */
-  db: AbstractQuery<TSchema>;
+  db: SimpleQueryInterface<TSchema>;
 };
 
 /**
@@ -163,8 +250,10 @@ export class DatabaseFragmentDefinitionBuilder<
   TServices,
   TServiceDependencies,
   TPrivateServices,
+  THooks extends HooksMap = {},
   TServiceThisContext extends RequestThisContext = DatabaseHandlerContext,
   THandlerThisContext extends RequestThisContext = DatabaseHandlerContext,
+  TLinkedFragments extends Record<string, AnyFragnoInstantiatedFragment> = {},
 > {
   // Store the base builder - we'll replace its storage and context setup when building
   #baseBuilder: FragmentDefinitionBuilder<
@@ -177,10 +266,12 @@ export class DatabaseFragmentDefinitionBuilder<
     TPrivateServices,
     TServiceThisContext,
     THandlerThisContext,
-    DatabaseRequestStorage
+    DatabaseRequestStorage,
+    TLinkedFragments
   >;
   #schema: TSchema;
   #namespace: string;
+  #hooksFactory?: (context: { config: TConfig; options: FragnoPublicConfigWithDatabase }) => THooks;
 
   constructor(
     baseBuilder: FragmentDefinitionBuilder<
@@ -193,14 +284,20 @@ export class DatabaseFragmentDefinitionBuilder<
       TPrivateServices,
       TServiceThisContext,
       THandlerThisContext,
-      DatabaseRequestStorage
+      DatabaseRequestStorage,
+      TLinkedFragments
     >,
     schema: TSchema,
     namespace?: string,
+    hooksFactory?: (context: {
+      config: TConfig;
+      options: FragnoPublicConfigWithDatabase;
+    }) => THooks,
   ) {
     this.#baseBuilder = baseBuilder;
     this.#schema = schema;
-    this.#namespace = namespace ?? baseBuilder.name + "-db";
+    this.#namespace = namespace ?? baseBuilder.name;
+    this.#hooksFactory = hooksFactory;
   }
 
   /**
@@ -211,7 +308,7 @@ export class DatabaseFragmentDefinitionBuilder<
     fn: (context: {
       config: TConfig;
       options: FragnoPublicConfigWithDatabase;
-      db: AbstractQuery<TSchema>;
+      db: SimpleQueryInterface<TSchema>;
       databaseAdapter: DatabaseAdapter<any>; // eslint-disable-line @typescript-eslint/no-explicit-any
     }) => TNewDeps,
   ): DatabaseFragmentDefinitionBuilder<
@@ -222,8 +319,10 @@ export class DatabaseFragmentDefinitionBuilder<
     {},
     TServiceDependencies,
     {},
+    THooks,
     TServiceThisContext,
-    THandlerThisContext
+    THandlerThisContext,
+    TLinkedFragments
   > {
     // Wrap user function to inject DB context
     const wrappedFn = (context: { config: TConfig; options: FragnoPublicConfigWithDatabase }) => {
@@ -255,8 +354,12 @@ export class DatabaseFragmentDefinitionBuilder<
     // Create new base builder with wrapped function
     const newBaseBuilder = this.#baseBuilder.withDependencies(wrappedFn);
 
-    // Return new database builder with updated base
-    return new DatabaseFragmentDefinitionBuilder(newBaseBuilder, this.#schema, this.#namespace);
+    return new DatabaseFragmentDefinitionBuilder(
+      newBaseBuilder,
+      this.#schema,
+      this.#namespace,
+      this.#hooksFactory,
+    );
   }
 
   providesBaseService<TNewService>(
@@ -277,12 +380,19 @@ export class DatabaseFragmentDefinitionBuilder<
     TServices,
     TServiceDependencies,
     TPrivateServices,
+    THooks,
     TServiceThisContext,
-    THandlerThisContext
+    THandlerThisContext,
+    TLinkedFragments
   > {
     const newBaseBuilder = this.#baseBuilder.providesBaseService<TNewService>(fn);
 
-    return new DatabaseFragmentDefinitionBuilder(newBaseBuilder, this.#schema, this.#namespace);
+    return new DatabaseFragmentDefinitionBuilder(
+      newBaseBuilder,
+      this.#schema,
+      this.#namespace,
+      this.#hooksFactory,
+    );
   }
 
   providesService<TServiceName extends string, TService>(
@@ -304,15 +414,147 @@ export class DatabaseFragmentDefinitionBuilder<
     TServices & { [K in TServiceName]: TService },
     TServiceDependencies,
     TPrivateServices,
+    THooks,
     TServiceThisContext,
-    THandlerThisContext
+    THandlerThisContext,
+    TLinkedFragments
   > {
     const newBaseBuilder = this.#baseBuilder.providesService<TServiceName, TService>(
       serviceName,
       fn,
     );
 
-    return new DatabaseFragmentDefinitionBuilder(newBaseBuilder, this.#schema, this.#namespace);
+    return new DatabaseFragmentDefinitionBuilder(
+      newBaseBuilder,
+      this.#schema,
+      this.#namespace,
+      this.#hooksFactory,
+    );
+  }
+
+  /**
+   * Provide a private service that is only accessible to the fragment author.
+   * Private services are NOT exposed on the fragment instance, but can be used
+   * when defining other services (baseServices, namedServices, and other privateServices).
+   * Private services are instantiated in order, so earlier private services are available
+   * to later ones.
+   */
+  providesPrivateService<TServiceName extends string, TService>(
+    serviceName: TServiceName,
+    fn: ServiceConstructorFn<
+      TConfig,
+      FragnoPublicConfigWithDatabase,
+      TDeps,
+      TServiceDependencies,
+      TPrivateServices,
+      TService,
+      TServiceThisContext
+    >,
+  ): DatabaseFragmentDefinitionBuilder<
+    TSchema,
+    TConfig,
+    TDeps,
+    TBaseServices,
+    TServices,
+    TServiceDependencies,
+    TPrivateServices & { [K in TServiceName]: TService },
+    THooks,
+    TServiceThisContext,
+    THandlerThisContext,
+    TLinkedFragments
+  > {
+    const newBaseBuilder = this.#baseBuilder.providesPrivateService<TServiceName, TService>(
+      serviceName,
+      fn,
+    );
+
+    return new DatabaseFragmentDefinitionBuilder(
+      newBaseBuilder,
+      this.#schema,
+      this.#namespace,
+      this.#hooksFactory,
+    );
+  }
+
+  /**
+   * Define durable hooks for this fragment.
+   * Hooks are automatically persisted and retried on failure.
+   *
+   * @param fn - Function that receives defineHook helper and returns a hooks map
+   * @returns Builder with hooks type set
+   *
+   * @example
+   * ```ts
+   * .provideHooks(({ defineHook, config }) => ({
+   *   onSubscribe: defineHook(async function (payload: { email: string }) {
+   *     // 'this' context available (HookServiceContext with nonce)
+   *     await config.onSubscribe?.(payload.email);
+   *   }),
+   * }))
+   * ```
+   */
+  provideHooks<TNewHooks extends HooksMap>(
+    fn: (context: {
+      config: TConfig;
+      options: FragnoPublicConfigWithDatabase;
+      defineHook: <TPayload>(
+        hook: (this: HookContext, payload: TPayload) => void | Promise<void>,
+      ) => HookFn<TPayload>;
+    }) => TNewHooks,
+  ): DatabaseFragmentDefinitionBuilder<
+    TSchema,
+    TConfig,
+    TDeps,
+    TBaseServices,
+    TServices,
+    TServiceDependencies,
+    TPrivateServices,
+    TNewHooks,
+    DatabaseServiceContext<TNewHooks>,
+    THandlerThisContext,
+    TLinkedFragments
+  > {
+    const defineHook = <TPayload>(
+      hook: (this: HookContext, payload: TPayload) => void | Promise<void>,
+    ): HookFn<TPayload> => {
+      return hook;
+    };
+
+    // Store the hooks factory - it will be called in build() with config/options
+    const hooksFactory = (context: {
+      config: TConfig;
+      options: FragnoPublicConfigWithDatabase;
+    }) => {
+      return fn({
+        config: context.config,
+        options: context.options,
+        defineHook,
+      });
+    };
+
+    // Create new builder with hooks factory stored
+    // Cast is safe: we're only changing THooks and TServiceThisContext type parameters
+    const newBuilder = new DatabaseFragmentDefinitionBuilder(
+      this.#baseBuilder,
+      this.#schema,
+      this.#namespace,
+    ) as unknown as DatabaseFragmentDefinitionBuilder<
+      TSchema,
+      TConfig,
+      TDeps,
+      TBaseServices,
+      TServices,
+      TServiceDependencies,
+      TPrivateServices,
+      TNewHooks,
+      DatabaseServiceContext<TNewHooks>,
+      THandlerThisContext,
+      TLinkedFragments
+    >;
+
+    newBuilder.#hooksFactory = hooksFactory;
+
+    return newBuilder;
   }
 
   /**
@@ -329,12 +571,19 @@ export class DatabaseFragmentDefinitionBuilder<
     TServices,
     TServiceDependencies & { [K in TServiceName]: TService },
     TPrivateServices,
+    THooks,
     TServiceThisContext,
-    THandlerThisContext
+    THandlerThisContext,
+    TLinkedFragments
   > {
     const newBaseBuilder = this.#baseBuilder.usesService<TServiceName, TService>(serviceName);
 
-    return new DatabaseFragmentDefinitionBuilder(newBaseBuilder, this.#schema, this.#namespace);
+    return new DatabaseFragmentDefinitionBuilder(
+      newBaseBuilder,
+      this.#schema,
+      this.#namespace,
+      this.#hooksFactory,
+    );
   }
 
   /**
@@ -351,14 +600,21 @@ export class DatabaseFragmentDefinitionBuilder<
     TServices,
     TServiceDependencies & { [K in TServiceName]: TService | undefined },
     TPrivateServices,
+    THooks,
     TServiceThisContext,
-    THandlerThisContext
+    THandlerThisContext,
+    TLinkedFragments
   > {
     const newBaseBuilder = this.#baseBuilder.usesOptionalService<TServiceName, TService>(
       serviceName,
     );
 
-    return new DatabaseFragmentDefinitionBuilder(newBaseBuilder, this.#schema, this.#namespace);
+    return new DatabaseFragmentDefinitionBuilder(
+      newBaseBuilder,
+      this.#schema,
+      this.#namespace,
+      this.#hooksFactory,
+    );
   }
 
   /**
@@ -374,18 +630,38 @@ export class DatabaseFragmentDefinitionBuilder<
     TServices,
     TServiceDependencies,
     TPrivateServices,
-    DatabaseServiceContext,
-    DatabaseHandlerContext,
-    DatabaseRequestStorage
+    DatabaseServiceContext<THooks>,
+    DatabaseHandlerContext<THooks>,
+    DatabaseRequestStorage,
+    TLinkedFragments
   > {
+    const baseDef = this.#baseBuilder.build();
+
     // Ensure dependencies callback always exists for database fragments
     // If no user dependencies were defined, create a minimal one that only adds implicit deps
     const dependencies = (context: {
       config: TConfig;
       options: FragnoPublicConfigWithDatabase;
     }): TDeps => {
-      const baseDef = this.#baseBuilder.build();
-      const userDeps = baseDef.dependencies?.(context);
+      // In dry run mode, allow user deps to fail gracefully.
+      // This is critical for database fragments because the CLI needs access to the schema
+      // even when user dependencies (like API clients) can't be initialized.
+      // Without this, if user deps fail, we'd lose the implicit database dependencies
+      // (including schema), and the CLI couldn't extract schema information.
+      let userDeps;
+      try {
+        userDeps = baseDef.dependencies?.(context);
+      } catch (error) {
+        if (process.env["FRAGNO_INIT_DRY_RUN"] === "true") {
+          console.warn(
+            "Warning: Failed to initialize user dependencies in dry run mode:",
+            error instanceof Error ? error.message : String(error),
+          );
+          userDeps = {};
+        } else {
+          throw error;
+        }
+      }
 
       const { db } = createDatabaseContext(context.options, this.#schema, this.#namespace);
 
@@ -423,15 +699,29 @@ export class DatabaseFragmentDefinitionBuilder<
       },
     );
 
-    // Services get restricted context (no execute methods), handlers get execution context
+    // Get the internal fragment factory from linked fragments (added by withDatabase)
+    // Cast is safe: withDatabase() guarantees this fragment exists and has the correct type
+    const internalFragmentFactory = baseDef.linkedFragments?.["_fragno_internal"] as (context: {
+      config: TConfig;
+      options: FragnoPublicConfigWithDatabase;
+    }) => InternalFragmentInstance;
+
     const builderWithContext = builderWithStorage.withThisContext<
-      DatabaseServiceContext,
-      DatabaseHandlerContext
-    >(({ storage }) => {
-      // Service context - forSchema method to get restricted typed UOW
+      DatabaseServiceContext<THooks>,
+      DatabaseHandlerContext<THooks>
+    >(({ storage, config, options }) => {
+      // Create hooks config if hooks factory is defined
+      const hooksConfig = this.#hooksFactory
+        ? {
+            hooks: this.#hooksFactory({ config, options }),
+            namespace: this.#namespace,
+            internalFragment: internalFragmentFactory({ config, options }),
+          }
+        : undefined;
+
       function forSchema<TSchema extends AnySchema>(
         schema: TSchema,
-      ): TypedUnitOfWork<TSchema, [], unknown> {
+      ): TypedUnitOfWork<TSchema, [], unknown, THooks> {
         const uow = storage.getStore()?.uow;
         if (!uow) {
           throw new Error(
@@ -439,24 +729,44 @@ export class DatabaseFragmentDefinitionBuilder<
           );
         }
 
-        // Return typed view of restricted UOW
-        return uow.restrict().forSchema(schema);
+        return uow.restrict().forSchema<TSchema, THooks>(schema);
       }
 
-      const serviceContext: DatabaseServiceContext = {
+      const serviceTx = async <
+        TSchema extends AnySchema,
+        TRetrievalResults extends unknown[],
+        TMutationResult = void,
+      >(
+        schema: TSchema,
+        callbacks: ServiceTxCallbacks<TSchema, TRetrievalResults, TMutationResult, THooks>,
+      ): Promise<AwaitedPromisesInObject<TMutationResult>> => {
+        const uow = storage.getStore()?.uow;
+        if (!uow) {
+          throw new Error(
+            "No UnitOfWork in context. Service must be called within a route handler OR using `withUnitOfWork`.",
+          );
+        }
+
+        return executeServiceTx(schema, callbacks, uow);
+      };
+
+      const serviceContext: DatabaseServiceContext<THooks> = {
         uow: forSchema,
+        tx: serviceTx,
       };
 
-      // Handler context - only executeRestrictedUnitOfWork
       async function uow<TResult>(
         callback: (context: {
-          forSchema: <S extends AnySchema>(schema: S) => TypedUnitOfWork<S, [], unknown>;
+          forSchema: <S extends AnySchema, H extends HooksMap = THooks>(
+            schema: S,
+            hooks?: H,
+          ) => TypedUnitOfWork<S, [], unknown, H>;
           executeRetrieve: () => Promise<void>;
           executeMutate: () => Promise<void>;
           nonce: string;
           currentAttempt: number;
         }) => Promise<TResult> | TResult,
-        options?: Omit<ExecuteRestrictedUnitOfWorkOptions, "createUnitOfWork">,
+        execOptions?: Omit<ExecuteRestrictedUnitOfWorkOptions, "createUnitOfWork">,
       ): Promise<AwaitedPromisesInObject<TResult>> {
         const currentStorage = storage.getStore();
         if (!currentStorage) {
@@ -465,30 +775,131 @@ export class DatabaseFragmentDefinitionBuilder<
           );
         }
 
-        // Wrap callback to ensure it always returns a Promise
         const wrappedCallback = async (context: {
-          forSchema: <S extends AnySchema>(schema: S) => TypedUnitOfWork<S, [], unknown>;
+          forSchema: <S extends AnySchema, H extends HooksMap = THooks>(
+            schema: S,
+            hooks?: H,
+          ) => TypedUnitOfWork<S, [], unknown, H>;
           executeRetrieve: () => Promise<void>;
           executeMutate: () => Promise<void>;
           nonce: string;
           currentAttempt: number;
         }): Promise<TResult> => {
-          return await callback(context);
+          return callback(context);
         };
 
-        // Use the UOW from storage - reset it before each attempt for retry support
-        // Cast is safe because IUnitOfWork is actually implemented by UnitOfWork
+        const userOnBeforeMutate = execOptions?.onBeforeMutate;
+        const userOnSuccess = execOptions?.onSuccess;
+
         return executeRestrictedUnitOfWork(wrappedCallback, {
-          ...options,
+          ...execOptions,
           createUnitOfWork: () => {
             currentStorage.uow.reset();
+            if (hooksConfig) {
+              currentStorage.uow.registerSchema(
+                hooksConfig.internalFragment.$internal.deps.schema,
+                hooksConfig.internalFragment.$internal.deps.namespace,
+              );
+            }
+            // Safe cast: currentStorage.uow is always a UnitOfWork instance
             return currentStorage.uow as UnitOfWork;
           },
+          onBeforeMutate: (uow) => {
+            if (hooksConfig) {
+              prepareHookMutations(uow, hooksConfig);
+            }
+            if (userOnBeforeMutate) {
+              userOnBeforeMutate(uow);
+            }
+          },
+          onSuccess: async (uow) => {
+            if (hooksConfig) {
+              await processHooks(hooksConfig);
+            }
+            if (userOnSuccess) {
+              await userOnSuccess(uow);
+            }
+          },
         });
       }
 
-      const handlerContext: DatabaseHandlerContext = {
+      // Handler tx method with two overloads
+      function handlerTx<T extends readonly unknown[]>(
+        servicesFactory: () => readonly [...{ [K in keyof T]: Promise<T[K]> }],
+        execOptions?: Omit<ExecuteRestrictedUnitOfWorkOptions, "createUnitOfWork">,
+      ): Promise<{ [K in keyof T]: T[K] }>;
+      function handlerTx<TRetrieveResult, TMutationResult>(
+        callbacks: HandlerTxCallbacks<TRetrieveResult, TMutationResult, THooks>,
+        execOptions?: Omit<ExecuteRestrictedUnitOfWorkOptions, "createUnitOfWork">,
+      ): Promise<AwaitedPromisesInObject<TMutationResult>>;
+      async function handlerTx<T extends readonly unknown[], TRetrieveResult, TMutationResult>(
+        factoryOrCallbacks:
+          | (() => readonly [...{ [K in keyof T]: Promise<T[K]> }])
+          | HandlerTxCallbacks<TRetrieveResult, TMutationResult, THooks>,
+        execOptions?: Omit<ExecuteRestrictedUnitOfWorkOptions, "createUnitOfWork">,
+      ): Promise<{ [K in keyof T]: T[K] } | AwaitedPromisesInObject<TMutationResult>> {
+        const currentStorage = storage.getStore();
+        if (!currentStorage) {
+          throw new Error(
+            "No storage in context. Handler must be called within a request context.",
+          );
+        }
+
+        const userOnBeforeMutate = execOptions?.onBeforeMutate;
+        const userOnSuccess = execOptions?.onSuccess;
+
+        const createUow = () => {
+          currentStorage.uow.reset();
+
+          if (hooksConfig) {
+            currentStorage.uow.registerSchema(
+              hooksConfig.internalFragment.$internal.deps.schema,
+              hooksConfig.internalFragment.$internal.deps.namespace,
+            );
+          }
+          // Safe cast: currentStorage.uow is always a UnitOfWork instance
+          return currentStorage.uow as UnitOfWork;
+        };
+
+        const options: ExecuteRestrictedUnitOfWorkOptions = {
+          ...execOptions,
+          createUnitOfWork: createUow,
+          onBeforeMutate: (uow) => {
+            if (hooksConfig) {
+              prepareHookMutations(uow, hooksConfig);
+            }
+            if (userOnBeforeMutate) {
+              userOnBeforeMutate(uow);
+            }
+          },
+          onSuccess: async (uow) => {
+            if (hooksConfig) {
+              await processHooks(hooksConfig);
+            }
+            if (userOnSuccess) {
+              await userOnSuccess(uow);
+            }
+          },
+        };
+
+        // Check if it's a function (array syntax factory) or callbacks object (callback syntax)
+        if (typeof factoryOrCallbacks === "function") {
+          // Array syntax - factoryOrCallbacks is a factory function
+          return executeTxArray(
+            factoryOrCallbacks as () => readonly [...{ [K in keyof T]: Promise<T[K]> }],
+            options,
+          ) as Promise<{ [K in keyof T]: T[K] }>;
+        } else {
+          return executeTxCallbacks(
+            factoryOrCallbacks as HandlerTxCallbacks<TRetrieveResult, TMutationResult, THooks>,
+            options,
+          );
+        }
+      }
+
+      const handlerContext: DatabaseHandlerContext<THooks> = {
         uow,
+        tx: handlerTx,
       };
 
       return { serviceContext, handlerContext };