@geekmidas/constructs 0.0.12 → 0.0.13

package/src/functions/FunctionBuilder.ts

@@ -1,4 +1,4 @@
-import type { AuditStorage } from '@geekmidas/audit';
+import type { AuditStorage, AuditableAction } from '@geekmidas/audit';
 import type { EventPublisher } from '@geekmidas/events';
 import type { Logger } from '@geekmidas/logger';
 import { ConsoleLogger } from '@geekmidas/logger/console';
@@ -21,6 +21,12 @@ export class FunctionBuilder<
   TEventPublisherServiceName extends string = string,
   TAuditStorage extends AuditStorage | undefined = undefined,
   TAuditStorageServiceName extends string = string,
+  TDatabase = undefined,
+  TDatabaseServiceName extends string = string,
+  TAuditAction extends AuditableAction<string, unknown> = AuditableAction<
+    string,
+    unknown
+  >,
 > extends BaseFunctionBuilder<
   TInput,
   OutSchema,
@@ -29,7 +35,9 @@ export class FunctionBuilder<
   TEventPublisher,
   TEventPublisherServiceName,
   TAuditStorage,
-  TAuditStorageServiceName
+  TAuditStorageServiceName,
+  TDatabase,
+  TDatabaseServiceName
 > {
   protected _memorySize?: number;
 
@@ -57,7 +65,10 @@ export class FunctionBuilder<
     TEventPublisher,
     TEventPublisherServiceName,
     TAuditStorage,
-    TAuditStorageServiceName
+    TAuditStorageServiceName,
+    TDatabase,
+    TDatabaseServiceName,
+    TAuditAction
   > {
     this.outputSchema = schema as unknown as OutSchema;
 
@@ -69,7 +80,10 @@ export class FunctionBuilder<
       TEventPublisher,
       TEventPublisherServiceName,
       TAuditStorage,
-      TAuditStorageServiceName
+      TAuditStorageServiceName,
+      TDatabase,
+      TDatabaseServiceName,
+      TAuditAction
     >;
   }
 
@@ -83,7 +97,10 @@ export class FunctionBuilder<
     TEventPublisher,
     TEventPublisherServiceName,
     TAuditStorage,
-    TAuditStorageServiceName
+    TAuditStorageServiceName,
+    TDatabase,
+    TDatabaseServiceName,
+    TAuditAction
   > {
     this.inputSchema = schema as unknown as TInput;
 
@@ -95,7 +112,10 @@ export class FunctionBuilder<
       TEventPublisher,
       TEventPublisherServiceName,
       TAuditStorage,
-      TAuditStorageServiceName
+      TAuditStorageServiceName,
+      TDatabase,
+      TDatabaseServiceName,
+      TAuditAction
     >;
   }
 
@@ -109,7 +129,10 @@ export class FunctionBuilder<
     TEventPublisher,
     TEventPublisherServiceName,
     TAuditStorage,
-    TAuditStorageServiceName
+    TAuditStorageServiceName,
+    TDatabase,
+    TDatabaseServiceName,
+    TAuditAction
   > {
     this._services = uniqBy(
       [...this._services, ...services],
@@ -124,7 +147,10 @@ export class FunctionBuilder<
       TEventPublisher,
       TEventPublisherServiceName,
       TAuditStorage,
-      TAuditStorageServiceName
+      TAuditStorageServiceName,
+      TDatabase,
+      TDatabaseServiceName,
+      TAuditAction
     >;
   }
 
@@ -138,7 +164,10 @@ export class FunctionBuilder<
     TEventPublisher,
     TEventPublisherServiceName,
     TAuditStorage,
-    TAuditStorageServiceName
+    TAuditStorageServiceName,
+    TDatabase,
+    TDatabaseServiceName,
+    TAuditAction
   > {
     this._logger = logger as unknown as TLogger;
 
@@ -150,7 +179,10 @@ export class FunctionBuilder<
       TEventPublisher,
       TEventPublisherServiceName,
       TAuditStorage,
-      TAuditStorageServiceName
+      TAuditStorageServiceName,
+      TDatabase,
+      TDatabaseServiceName,
+      TAuditAction
     >;
   }
 
@@ -164,7 +196,10 @@ export class FunctionBuilder<
     T,
     TName,
     TAuditStorage,
-    TAuditStorageServiceName
+    TAuditStorageServiceName,
+    TDatabase,
+    TDatabaseServiceName,
+    TAuditAction
   > {
     this._publisher = publisher as unknown as Service<
       TEventPublisherServiceName,
@@ -179,7 +214,10 @@ export class FunctionBuilder<
       T,
       TName,
       TAuditStorage,
-      TAuditStorageServiceName
+      TAuditStorageServiceName,
+      TDatabase,
+      TDatabaseServiceName,
+      TAuditAction
     >;
   }
 
@@ -193,7 +231,10 @@ export class FunctionBuilder<
     TEventPublisher,
     TEventPublisherServiceName,
     T,
-    TName
+    TName,
+    TDatabase,
+    TDatabaseServiceName,
+    TAuditAction
   > {
     this._auditorStorage = storage as unknown as Service<
       TAuditStorageServiceName,
@@ -208,22 +249,115 @@ export class FunctionBuilder<
       TEventPublisher,
       TEventPublisherServiceName,
       T,
-      TName
+      TName,
+      TDatabase,
+      TDatabaseServiceName,
+      TAuditAction
+    >;
+  }
+
+  /**
+   * Set the audit action types for this function.
+   * This provides type-safety for the auditor in the handler context.
+   */
+  actions<T extends AuditableAction<string, unknown>>(): FunctionBuilder<
+    TInput,
+    OutSchema,
+    TServices,
+    TLogger,
+    TEventPublisher,
+    TEventPublisherServiceName,
+    TAuditStorage,
+    TAuditStorageServiceName,
+    TDatabase,
+    TDatabaseServiceName,
+    T
+  > {
+    return this as unknown as FunctionBuilder<
+      TInput,
+      OutSchema,
+      TServices,
+      TLogger,
+      TEventPublisher,
+      TEventPublisherServiceName,
+      TAuditStorage,
+      TAuditStorageServiceName,
+      TDatabase,
+      TDatabaseServiceName,
+      T
+    >;
+  }
+
+  /**
+   * Set the database service for this function.
+   * The database will be available in the handler context as `db`.
+   */
+  database<T, TName extends string>(
+    service: Service<TName, T>,
+  ): FunctionBuilder<
+    TInput,
+    OutSchema,
+    TServices,
+    TLogger,
+    TEventPublisher,
+    TEventPublisherServiceName,
+    TAuditStorage,
+    TAuditStorageServiceName,
+    T,
+    TName,
+    TAuditAction
+  > {
+    this._databaseService = service as unknown as Service<
+      TDatabaseServiceName,
+      TDatabase
+    >;
+
+    return this as unknown as FunctionBuilder<
+      TInput,
+      OutSchema,
+      TServices,
+      TLogger,
+      TEventPublisher,
+      TEventPublisherServiceName,
+      TAuditStorage,
+      TAuditStorageServiceName,
+      T,
+      TName,
+      TAuditAction
     >;
   }
 
   handle(
-    fn: FunctionHandler<TInput, TServices, TLogger, OutSchema>,
+    fn: FunctionHandler<
+      TInput,
+      TServices,
+      TLogger,
+      OutSchema,
+      TDatabase,
+      TAuditStorage,
+      TAuditAction
+    >,
   ): Function<
     TInput,
     TServices,
     TLogger,
     OutSchema,
-    FunctionHandler<TInput, TServices, TLogger, OutSchema>,
     TEventPublisher,
     TEventPublisherServiceName,
     TAuditStorage,
-    TAuditStorageServiceName
+    TAuditStorageServiceName,
+    TDatabase,
+    TDatabaseServiceName,
+    TAuditAction,
+    FunctionHandler<
+      TInput,
+      TServices,
+      TLogger,
+      OutSchema,
+      TDatabase,
+      TAuditStorage,
+      TAuditAction
+    >
   > {
     const func = new Function(
       fn,
@@ -237,6 +371,7 @@ export class FunctionBuilder<
       this._events,
       this._memorySize,
       this._auditorStorage,
+      this._databaseService,
     );
 
     // Reset builder state after creating the function to prevent pollution
@@ -245,6 +380,7 @@ export class FunctionBuilder<
     this._events = [];
     this._publisher = undefined;
     this._auditorStorage = undefined;
+    this._databaseService = undefined;
     this.inputSchema = undefined;
     this.outputSchema = undefined;
     this._timeout = undefined;

package/src/functions/FunctionExecutionWrapper.ts

@@ -1,3 +1,6 @@
+import type { AuditStorage, AuditableAction, Auditor } from '@geekmidas/audit';
+import { DefaultAuditor } from '@geekmidas/audit';
+import { withAuditableTransaction } from '@geekmidas/audit/kysely';
 import type { EnvironmentParser } from '@geekmidas/envkit';
 import type { EventPublisher } from '@geekmidas/events';
 import type { Logger } from '@geekmidas/logger';
@@ -23,6 +26,14 @@ export abstract class FunctionExecutionWrapper<
   TLogger extends Logger = Logger,
   TEventPublisher extends EventPublisher<any> | undefined = undefined,
   TEventPublisherServiceName extends string = string,
+  TAuditStorage extends AuditStorage | undefined = undefined,
+  TAuditStorageServiceName extends string = string,
+  TDatabase = undefined,
+  TDatabaseServiceName extends string = string,
+  TAuditAction extends AuditableAction<string, unknown> = AuditableAction<
+    string,
+    unknown
+  >,
 > {
   constructor(
     protected envParser: EnvironmentParser<{}>,
@@ -31,9 +42,22 @@ export abstract class FunctionExecutionWrapper<
       TServices,
       TLogger,
       TOutSchema,
-      FunctionHandler<TInput, TServices, TLogger, TOutSchema>,
       TEventPublisher,
-      TEventPublisherServiceName
+      TEventPublisherServiceName,
+      TAuditStorage,
+      TAuditStorageServiceName,
+      TDatabase,
+      TDatabaseServiceName,
+      TAuditAction,
+      FunctionHandler<
+        TInput,
+        TServices,
+        TLogger,
+        TOutSchema,
+        TDatabase,
+        TAuditStorage,
+        TAuditAction
+      >
     >,
   ) {}
 
@@ -56,6 +80,113 @@ export abstract class FunctionExecutionWrapper<
     return this.serviceDiscovery.register(this.fn.services);
   }
 
+  async getDatabase(): Promise<TDatabase | undefined> {
+    if (!this.fn.databaseService) {
+      return undefined;
+    }
+
+    const services = await this.serviceDiscovery.register([
+      this.fn.databaseService,
+    ]);
+
+    return services[
+      this.fn.databaseService.serviceName as keyof typeof services
+    ] as TDatabase;
+  }
+
+  /**
+   * Get the audit storage service if configured.
+   * Returns undefined if no auditor storage is configured.
+   */
+  async getAuditStorage(): Promise<TAuditStorage | undefined> {
+    if (!this.fn.auditorStorageService) {
+      return undefined;
+    }
+
+    const services = await this.serviceDiscovery.register([
+      this.fn.auditorStorageService,
+    ]);
+
+    return services[
+      this.fn.auditorStorageService.serviceName as keyof typeof services
+    ] as TAuditStorage;
+  }
+
+  /**
+   * Create an auditor instance for the function.
+   * Returns undefined if no auditor storage is configured.
+   */
+  async createAuditor(): Promise<Auditor<TAuditAction> | undefined> {
+    const storage = await this.getAuditStorage();
+    if (!storage) {
+      return undefined;
+    }
+
+    return new DefaultAuditor<TAuditAction>({
+      actor: { id: 'system', type: 'system' },
+      storage: storage as AuditStorage,
+      metadata: {
+        function: this.fn.type,
+      },
+    });
+  }
+
+  /**
+   * Execute handler with audit transaction support.
+   * If the audit storage has a database (via getDatabase()), wraps execution
+   * in a transaction so audits are atomic with handler's database operations.
+   *
+   * @param handler - The handler function to execute (receives auditor and db)
+   * @returns The handler result
+   */
+  async executeWithAudit<T>(
+    handler: (ctx: {
+      auditor?: Auditor<TAuditAction>;
+      db?: TDatabase;
+    }) => Promise<T>,
+  ): Promise<T> {
+    const auditor = await this.createAuditor();
+    const storage = await this.getAuditStorage();
+
+    // No audit context - just run handler with regular db
+    if (!auditor || !storage) {
+      const db = await this.getDatabase();
+      return handler({ db });
+    }
+
+    // Check if storage has a database and db service names match
+    const storageDb = (storage as AuditStorage).getDatabase?.();
+    const databaseServiceName = this.fn.databaseService?.serviceName;
+    const auditDbServiceName = this.fn.auditorStorageService?.serviceName;
+
+    // If the audit storage has a database and we're using the same database service
+    // (or the audit storage provides the database), use transactional execution
+    if (storageDb && databaseServiceName && auditDbServiceName) {
+      return withAuditableTransaction(
+        storageDb as any,
+        auditor as any,
+        async (trx) => {
+          // Use transaction as db
+          const response = await handler({
+            auditor,
+            db: trx as TDatabase,
+          });
+          // Audits are flushed by withAuditableTransaction before commit
+          return response;
+        },
+      );
+    }
+
+    // No database on storage or service names don't match - run handler and flush audits after
+    const db = await this.getDatabase();
+    const response = await handler({ auditor, db });
+
+    // Flush audits (no transaction)
+    await auditor.flush();
+
+    return response;
+  }
+
   async getFunctionInput<TEvent>(
     event: TEvent,
   ): Promise<InferComposableStandardSchema<TInput>> {

package/src/functions/TestFunctionAdaptor.ts

@@ -1,3 +1,5 @@
+import type { AuditStorage, AuditableAction, Auditor } from '@geekmidas/audit';
+import { DefaultAuditor } from '@geekmidas/audit';
 import { EnvironmentParser } from '@geekmidas/envkit';
 import type { EventPublisher } from '@geekmidas/events';
 import type { Logger } from '@geekmidas/logger';
@@ -17,12 +19,20 @@ import type { Function } from './Function';
 import { FunctionBuilder } from './FunctionBuilder';
 
 export class TestFunctionAdaptor<
-  TInput extends StandardSchemaV1 | undefined = undefined,
+  TInput extends ComposableStandardSchema | undefined = undefined,
   TOutSchema extends StandardSchemaV1 | undefined = undefined,
   TServices extends Service[] = [],
   TLogger extends Logger = Logger,
   TEventPublisher extends EventPublisher<any> | undefined = undefined,
   TEventPublisherServiceName extends string = string,
+  TAuditStorage extends AuditStorage | undefined = undefined,
+  TAuditStorageServiceName extends string = string,
+  TDatabase = undefined,
+  TDatabaseServiceName extends string = string,
+  TAuditAction extends AuditableAction<string, unknown> = AuditableAction<
+    string,
+    unknown
+  >,
 > {
   static getDefaultServiceDiscovery<
     TInput extends ComposableStandardSchema | undefined = undefined,
@@ -31,15 +41,28 @@ export class TestFunctionAdaptor<
     TLogger extends Logger = Logger,
     TEventPublisher extends EventPublisher<any> | undefined = undefined,
     TEventPublisherServiceName extends string = string,
+    TAuditStorage extends AuditStorage | undefined = undefined,
+    TAuditStorageServiceName extends string = string,
+    TDatabase = undefined,
+    TDatabaseServiceName extends string = string,
+    TAuditAction extends AuditableAction<string, unknown> = AuditableAction<
+      string,
+      unknown
+    >,
   >(
     fn: Function<
       TInput,
       TServices,
       TLogger,
       TOutSchema,
-      any,
       TEventPublisher,
-      TEventPublisherServiceName
+      TEventPublisherServiceName,
+      TAuditStorage,
+      TAuditStorageServiceName,
+      TDatabase,
+      TDatabaseServiceName,
+      TAuditAction,
+      any
     >,
   ) {
     return ServiceDiscovery.getInstance(fn.logger, new EnvironmentParser({}));
@@ -51,9 +74,14 @@ export class TestFunctionAdaptor<
       TServices,
       TLogger,
       TOutSchema,
-      any,
       TEventPublisher,
-      TEventPublisherServiceName
+      TEventPublisherServiceName,
+      TAuditStorage,
+      TAuditStorageServiceName,
+      TDatabase,
+      TDatabaseServiceName,
+      TAuditAction,
+      any
     >,
     private serviceDiscovery: ServiceDiscovery<
       any,
@@ -66,7 +94,9 @@ export class TestFunctionAdaptor<
       TInput,
       TServices,
       TEventPublisher,
-      TEventPublisherServiceName
+      TEventPublisherServiceName,
+      TDatabase,
+      TAuditAction
     >,
   ): Promise<InferStandardSchema<TOutSchema>> {
     // Parse input if schema is provided
@@ -89,16 +119,65 @@ export class TestFunctionAdaptor<
       services = await this.serviceDiscovery.register(this.fn.services);
     }
 
+    // Resolve database (use provided db or register from function)
+    let db: TDatabase | undefined;
+    if ('db' in ctx && ctx.db !== undefined) {
+      db = ctx.db;
+    } else if (this.fn.databaseService) {
+      const dbServices = await this.serviceDiscovery.register([
+        this.fn.databaseService,
+      ]);
+      db = dbServices[
+        this.fn.databaseService.serviceName as keyof typeof dbServices
+      ] as TDatabase;
+    }
+
+    // Resolve auditor (use provided auditor or create from function)
+    let auditor: Auditor<TAuditAction> | undefined;
+    if ('auditor' in ctx && ctx.auditor !== undefined) {
+      auditor = ctx.auditor;
+    } else if (this.fn.auditorStorageService) {
+      const auditServices = await this.serviceDiscovery.register([
+        this.fn.auditorStorageService,
+      ]);
+      const storage = auditServices[
+        this.fn.auditorStorageService.serviceName as keyof typeof auditServices
+      ] as AuditStorage;
+
+      auditor = new DefaultAuditor<TAuditAction>({
+        actor: { id: 'system', type: 'system' },
+        storage,
+        metadata: {
+          function: this.fn.type,
+          test: true,
+        },
+      });
+    }
+
     // Execute the function
     const response = await this.fn['fn']({
       input: parsedInput,
       services,
       logger,
-    });
+      db,
+      auditor,
+    } as any);
 
     // Parse output if schema is provided
     const output = await this.fn.parseOutput(response);
 
+    // Flush audits if any were recorded
+    if (auditor) {
+      const records = auditor.getRecords();
+      if (records.length > 0) {
+        logger.debug(
+          { auditCount: records.length },
+          'Flushing function audits',
+        );
+        await auditor.flush();
+      }
+    }
+
     // Register publisher service if provided in context
 
     await publishEvents(
@@ -114,12 +193,19 @@ export class TestFunctionAdaptor<
 }
 
 export type TestFunctionRequest<
-  TInput extends StandardSchemaV1 | undefined = undefined,
+  TInput extends ComposableStandardSchema | undefined = undefined,
   TServices extends Service[] = [],
   TEventPublisher extends EventPublisher<any> | undefined = undefined,
   TEventPublisherServiceName extends string = string,
+  TDatabase = undefined,
+  TAuditAction extends AuditableAction<string, unknown> = AuditableAction<
+    string,
+    unknown
+  >,
 > = {
   input: InferComposableStandardSchema<TInput>;
   services: ServiceRecord<TServices>;
   publisher?: Service<TEventPublisherServiceName, TEventPublisher>;
+  db?: TDatabase;
+  auditor?: Auditor<TAuditAction>;
 } & InferComposableStandardSchema<{ input: TInput }>;

package/src/functions/__tests__/AWSLambdaFunctionAdaptor.spec.ts

@@ -373,4 +373,86 @@ describe('AWSLambdaFunction', () => {
   });
 
   // Skipping timeout tests for now as middleware order needs to be adjusted
+
+  describe('database', () => {
+    // Mock database service
+    class MockDatabase {
+      async query(sql: string) {
+        return [{ id: '1', name: 'Test User' }];
+      }
+    }
+
+    class DatabaseService implements Service<'database', MockDatabase> {
+      serviceName = 'database' as const;
+      db = new MockDatabase();
+
+      async register() {
+        return this.db;
+      }
+    }
+
+    it('should inject database into function context', async () => {
+      const databaseService = new DatabaseService();
+      const handler = vi.fn(async ({ db }) => {
+        const result = await db.query('SELECT * FROM users');
+        return { users: result };
+      });
+
+      const outputSchema = z.object({
+        users: z.array(z.object({ id: z.string(), name: z.string() })),
+      });
+
+      const fn = new Function(
+        handler,
+        undefined,
+        undefined,
+        undefined,
+        outputSchema,
+        [],
+        logger,
+        undefined,
+        [],
+        undefined,
+        undefined,
+        databaseService,
+      );
+
+      const adaptor = new AWSLambdaFunction(envParser, fn);
+      const lambdaHandler = adaptor.handler;
+
+      const result = await lambdaHandler({}, createMockContext(), vi.fn());
+
+      expect(result).toEqual({ users: [{ id: '1', name: 'Test User' }] });
+      expect(handler).toHaveBeenCalledWith(
+        expect.objectContaining({
+          db: expect.any(MockDatabase),
+        }),
+      );
+    });
+
+    it('should have db as undefined when no database service is configured', async () => {
+      const handler = vi.fn(async ({ db }) => {
+        return { hasDb: db !== undefined };
+      });
+
+      const outputSchema = z.object({ hasDb: z.boolean() });
+
+      const fn = new Function(
+        handler,
+        undefined,
+        undefined,
+        undefined,
+        outputSchema,
+        [],
+        logger,
+      );
+
+      const adaptor = new AWSLambdaFunction(envParser, fn);
+      const lambdaHandler = adaptor.handler;
+
+      const result = await lambdaHandler({}, createMockContext(), vi.fn());
+
+      expect(result).toEqual({ hasDb: false });
+    });
+  });
 });