import { pgTable, uuid, varchar, text, integer, jsonb, timestamp, pgEnum, index, check, type AnyPgColumn, } from 'drizzle-orm/pg-core'; import { sql } from 'drizzle-orm'; import { projects } from './projects'; import { flows } from './flows'; import { apiKeys } from './apiKeys'; // Enum for generation status export const generationStatusEnum = pgEnum('generation_status', [ 'pending', 'processing', 'success', 'failed', ]); // Type for referenced images JSONB export type ReferencedImage = { imageId: string; alias: string; }; export const generations = pgTable( 'generations', { id: uuid('id').primaryKey().defaultRandom(), // Relations projectId: uuid('project_id') .notNull() .references(() => projects.id, { onDelete: 'cascade' }), flowId: uuid('flow_id').references(() => flows.id, { onDelete: 'set null' }), pendingFlowId: text('pending_flow_id'), // Temporary UUID for lazy flow pattern apiKeyId: uuid('api_key_id').references(() => apiKeys.id, { onDelete: 'set null' }), // Status status: generationStatusEnum('status').notNull().default('pending'), // Prompts (Section 2.1: Reversed semantics) // prompt: The prompt that was ACTUALLY USED for generation (enhanced OR original) // originalPrompt: User's ORIGINAL input, only stored if autoEnhance was used prompt: text('prompt').notNull(), // Prompt used for generation originalPrompt: text('original_prompt'), // User's original (nullable, only if enhanced) // Generation parameters aspectRatio: varchar('aspect_ratio', { length: 10 }), width: integer('width'), height: integer('height'), // AI Model modelName: varchar('model_name', { length: 100 }).notNull().default('gemini-flash-image-001'), modelVersion: varchar('model_version', { length: 50 }), // Result outputImageId: uuid('output_image_id').references( (): AnyPgColumn => { const { images } = require('./images'); return images.id; }, { onDelete: 'set null' }, ), // Referenced images used in generation // Format: [{ "imageId": "uuid", "alias": "@product" }, ...] referencedImages: jsonb('referenced_images').$type(), // Error handling errorMessage: text('error_message'), errorCode: varchar('error_code', { length: 50 }), retryCount: integer('retry_count').notNull().default(0), // Metrics processingTimeMs: integer('processing_time_ms'), cost: integer('cost'), // In cents (USD) // Request context requestId: uuid('request_id'), userAgent: text('user_agent'), ipAddress: text('ip_address'), // Metadata meta: jsonb('meta').$type>().notNull().default({}), // Audit createdAt: timestamp('created_at').notNull().defaultNow(), updatedAt: timestamp('updated_at') .notNull() .defaultNow() .$onUpdate(() => new Date()), }, (table) => ({ // CHECK constraints statusSuccessCheck: check( 'status_success_check', sql`(${table.status} = 'success' AND ${table.outputImageId} IS NOT NULL) OR (${table.status} != 'success')`, ), statusFailedCheck: check( 'status_failed_check', sql`(${table.status} = 'failed' AND ${table.errorMessage} IS NOT NULL) OR (${table.status} != 'failed')`, ), retryCountCheck: check('retry_count_check', sql`${table.retryCount} >= 0`), processingTimeCheck: check( 'processing_time_check', sql`${table.processingTimeMs} IS NULL OR ${table.processingTimeMs} >= 0`, ), costCheck: check('cost_check', sql`${table.cost} IS NULL OR ${table.cost} >= 0`), // Indexes // Index for querying generations by project and status projectStatusIdx: index('idx_generations_project_status').on( table.projectId, table.status, table.createdAt.desc(), ), // Index for flow-scoped generations (partial index) flowIdx: index('idx_generations_flow') .on(table.flowId, table.createdAt.desc()) .where(sql`${table.flowId} IS NOT NULL`), // Index for pending flow-scoped generations (partial index) pendingFlowIdx: index('idx_generations_pending_flow') .on(table.pendingFlowId, table.createdAt.desc()) .where(sql`${table.pendingFlowId} IS NOT NULL`), // Index for output image lookup outputIdx: index('idx_generations_output').on(table.outputImageId), // Index for request correlation requestIdx: index('idx_generations_request').on(table.requestId), // Index for API key audit trail apiKeyIdx: index('idx_generations_api_key').on(table.apiKeyId), }), ); export type Generation = typeof generations.$inferSelect; export type NewGeneration = typeof generations.$inferInsert;