banatie-service/packages/database/src/schema/generations.ts

149 lines
4.7 KiB
TypeScript

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<ReferencedImage[]>(),
// 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<Record<string, unknown>>().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;