@axiom-lattice/pg-stores 1.0.2 → 1.0.4

package/src/stores/PostgreSQLScheduleStorage.ts

@@ -0,0 +1,591 @@
+/**
+ * PostgreSQL implementation of ScheduleStorage
+ *
+ * Provides persistent storage for scheduled tasks
+ * Data survives service restarts
+ */
+
+import { Pool, PoolConfig } from "pg";
+import {
+  ScheduleStorage,
+  ScheduledTaskDefinition,
+  ScheduledTaskStatus,
+  ScheduleExecutionType,
+} from "@axiom-lattice/protocols";
+import { MigrationManager } from "../migrations/migration";
+import { createScheduledTasksTable } from "../migrations/schedule_migrations";
+
+/**
+ * PostgreSQL ScheduleStorage options
+ */
+export interface PostgreSQLScheduleStorageOptions {
+  /**
+   * PostgreSQL connection pool configuration
+   * Can be a connection string or PoolConfig object
+   */
+  poolConfig: string | PoolConfig;
+
+  /**
+   * Whether to run migrations automatically on initialization
+   * @default true
+   */
+  autoMigrate?: boolean;
+}
+
+/**
+ * Database row type for scheduled tasks
+ */
+interface ScheduledTaskRow {
+  task_id: string;
+  task_type: string;
+  payload: string | Record<string, any>;
+  assistant_id: string | null;
+  thread_id: string | null;
+  execution_type: string;
+  execute_at: Date | null;
+  delay_ms: number | null;
+  cron_expression: string | null;
+  timezone: string | null;
+  next_run_at: Date | null;
+  last_run_at: Date | null;
+  status: string;
+  run_count: number;
+  max_runs: number | null;
+  retry_count: number;
+  max_retries: number;
+  last_error: string | null;
+  created_at: Date;
+  updated_at: Date;
+  expires_at: Date | null;
+  metadata: string | Record<string, any> | null;
+}
+
+/**
+ * PostgreSQL implementation of ScheduleStorage
+ */
+export class PostgreSQLScheduleStorage implements ScheduleStorage {
+  private pool: Pool;
+  private migrationManager: MigrationManager;
+  private initialized: boolean = false;
+
+  constructor(options: PostgreSQLScheduleStorageOptions) {
+    // Create Pool from config
+    if (typeof options.poolConfig === "string") {
+      this.pool = new Pool({ connectionString: options.poolConfig });
+    } else {
+      this.pool = new Pool(options.poolConfig);
+    }
+
+    this.migrationManager = new MigrationManager(this.pool);
+    this.migrationManager.register(createScheduledTasksTable);
+
+    // Auto-migrate by default
+    if (options.autoMigrate !== false) {
+      this.initialize().catch((error) => {
+        console.error("Failed to initialize PostgreSQLScheduleStorage:", error);
+        throw error;
+      });
+    }
+  }
+
+  /**
+   * Dispose resources and close the connection pool
+   */
+  async dispose(): Promise<void> {
+    if (this.pool) {
+      await this.pool.end();
+    }
+  }
+
+  /**
+   * Initialize the store and run migrations
+   */
+  async initialize(): Promise<void> {
+    if (this.initialized) {
+      return;
+    }
+
+    await this.migrationManager.migrate();
+    this.initialized = true;
+  }
+
+  /**
+   * Ensure store is initialized
+   */
+  private async ensureInitialized(): Promise<void> {
+    if (!this.initialized) {
+      await this.initialize();
+    }
+  }
+
+  /**
+   * Save a new task
+   */
+  async save(task: ScheduledTaskDefinition): Promise<void> {
+    await this.ensureInitialized();
+
+    await this.pool.query(
+      `
+      INSERT INTO lattice_scheduled_tasks (
+        task_id, task_type, payload, assistant_id, thread_id, execution_type,
+        execute_at, delay_ms, cron_expression, timezone,
+        next_run_at, last_run_at, status, run_count, max_runs,
+        retry_count, max_retries, last_error,
+        created_at, updated_at, expires_at, metadata
+      ) VALUES (
+        $1, $2, $3, $4, $5, $6,
+        $7, $8, $9, $10,
+        $11, $12, $13, $14, $15,
+        $16, $17, $18,
+        $19, $20, $21, $22
+      )
+      ON CONFLICT (task_id) DO UPDATE SET
+        task_type = EXCLUDED.task_type,
+        payload = EXCLUDED.payload,
+        assistant_id = EXCLUDED.assistant_id,
+        thread_id = EXCLUDED.thread_id,
+        execution_type = EXCLUDED.execution_type,
+        execute_at = EXCLUDED.execute_at,
+        delay_ms = EXCLUDED.delay_ms,
+        cron_expression = EXCLUDED.cron_expression,
+        timezone = EXCLUDED.timezone,
+        next_run_at = EXCLUDED.next_run_at,
+        last_run_at = EXCLUDED.last_run_at,
+        status = EXCLUDED.status,
+        run_count = EXCLUDED.run_count,
+        max_runs = EXCLUDED.max_runs,
+        retry_count = EXCLUDED.retry_count,
+        max_retries = EXCLUDED.max_retries,
+        last_error = EXCLUDED.last_error,
+        updated_at = EXCLUDED.updated_at,
+        expires_at = EXCLUDED.expires_at,
+        metadata = EXCLUDED.metadata
+      `,
+      [
+        task.taskId,
+        task.taskType,
+        JSON.stringify(task.payload),
+        task.assistantId ?? null,
+        task.threadId ?? null,
+        task.executionType,
+        task.executeAt ? new Date(task.executeAt) : null,
+        task.delayMs ?? null,
+        task.cronExpression ?? null,
+        task.timezone ?? null,
+        task.nextRunAt ? new Date(task.nextRunAt) : null,
+        task.lastRunAt ? new Date(task.lastRunAt) : null,
+        task.status,
+        task.runCount,
+        task.maxRuns ?? null,
+        task.retryCount,
+        task.maxRetries,
+        task.lastError ?? null,
+        new Date(task.createdAt),
+        new Date(task.updatedAt),
+        task.expiresAt ? new Date(task.expiresAt) : null,
+        task.metadata ? JSON.stringify(task.metadata) : null,
+      ]
+    );
+  }
+
+  /**
+   * Get task by ID
+   */
+  async get(taskId: string): Promise<ScheduledTaskDefinition | null> {
+    await this.ensureInitialized();
+
+    const result = await this.pool.query<ScheduledTaskRow>(
+      `SELECT * FROM lattice_scheduled_tasks WHERE task_id = $1`,
+      [taskId]
+    );
+
+    if (result.rows.length === 0) {
+      return null;
+    }
+
+    return this.mapRowToTask(result.rows[0]);
+  }
+
+  /**
+   * Update task
+   */
+  async update(
+    taskId: string,
+    updates: Partial<ScheduledTaskDefinition>
+  ): Promise<void> {
+    await this.ensureInitialized();
+
+    const setClauses: string[] = [];
+    const values: any[] = [];
+    let paramIndex = 1;
+
+    // Build dynamic UPDATE query
+    if (updates.taskType !== undefined) {
+      setClauses.push(`task_type = $${paramIndex++}`);
+      values.push(updates.taskType);
+    }
+    if (updates.payload !== undefined) {
+      setClauses.push(`payload = $${paramIndex++}`);
+      values.push(JSON.stringify(updates.payload));
+    }
+    if (updates.assistantId !== undefined) {
+      setClauses.push(`assistant_id = $${paramIndex++}`);
+      values.push(updates.assistantId);
+    }
+    if (updates.threadId !== undefined) {
+      setClauses.push(`thread_id = $${paramIndex++}`);
+      values.push(updates.threadId);
+    }
+    if (updates.executionType !== undefined) {
+      setClauses.push(`execution_type = $${paramIndex++}`);
+      values.push(updates.executionType);
+    }
+    if (updates.executeAt !== undefined) {
+      setClauses.push(`execute_at = $${paramIndex++}`);
+      values.push(updates.executeAt ? new Date(updates.executeAt) : null);
+    }
+    if (updates.delayMs !== undefined) {
+      setClauses.push(`delay_ms = $${paramIndex++}`);
+      values.push(updates.delayMs);
+    }
+    if (updates.cronExpression !== undefined) {
+      setClauses.push(`cron_expression = $${paramIndex++}`);
+      values.push(updates.cronExpression);
+    }
+    if (updates.timezone !== undefined) {
+      setClauses.push(`timezone = $${paramIndex++}`);
+      values.push(updates.timezone);
+    }
+    if (updates.nextRunAt !== undefined) {
+      setClauses.push(`next_run_at = $${paramIndex++}`);
+      values.push(updates.nextRunAt ? new Date(updates.nextRunAt) : null);
+    }
+    if (updates.lastRunAt !== undefined) {
+      setClauses.push(`last_run_at = $${paramIndex++}`);
+      values.push(updates.lastRunAt ? new Date(updates.lastRunAt) : null);
+    }
+    if (updates.status !== undefined) {
+      setClauses.push(`status = $${paramIndex++}`);
+      values.push(updates.status);
+    }
+    if (updates.runCount !== undefined) {
+      setClauses.push(`run_count = $${paramIndex++}`);
+      values.push(updates.runCount);
+    }
+    if (updates.maxRuns !== undefined) {
+      setClauses.push(`max_runs = $${paramIndex++}`);
+      values.push(updates.maxRuns);
+    }
+    if (updates.retryCount !== undefined) {
+      setClauses.push(`retry_count = $${paramIndex++}`);
+      values.push(updates.retryCount);
+    }
+    if (updates.maxRetries !== undefined) {
+      setClauses.push(`max_retries = $${paramIndex++}`);
+      values.push(updates.maxRetries);
+    }
+    if (updates.lastError !== undefined) {
+      setClauses.push(`last_error = $${paramIndex++}`);
+      values.push(updates.lastError);
+    }
+    if (updates.expiresAt !== undefined) {
+      setClauses.push(`expires_at = $${paramIndex++}`);
+      values.push(updates.expiresAt ? new Date(updates.expiresAt) : null);
+    }
+    if (updates.metadata !== undefined) {
+      setClauses.push(`metadata = $${paramIndex++}`);
+      values.push(updates.metadata ? JSON.stringify(updates.metadata) : null);
+    }
+
+    // Always update updated_at
+    setClauses.push(`updated_at = $${paramIndex++}`);
+    values.push(new Date());
+
+    // Add taskId as the last parameter
+    values.push(taskId);
+
+    if (setClauses.length === 0) {
+      return;
+    }
+
+    await this.pool.query(
+      `UPDATE lattice_scheduled_tasks SET ${setClauses.join(
+        ", "
+      )} WHERE task_id = $${paramIndex}`,
+      values
+    );
+  }
+
+  /**
+   * Delete task
+   */
+  async delete(taskId: string): Promise<void> {
+    await this.ensureInitialized();
+
+    await this.pool.query(
+      `DELETE FROM lattice_scheduled_tasks WHERE task_id = $1`,
+      [taskId]
+    );
+  }
+
+  /**
+   * Get all active tasks (pending or paused)
+   */
+  async getActiveTasks(): Promise<ScheduledTaskDefinition[]> {
+    await this.ensureInitialized();
+
+    const result = await this.pool.query<ScheduledTaskRow>(
+      `SELECT * FROM lattice_scheduled_tasks WHERE status IN ($1, $2) ORDER BY created_at ASC`,
+      [ScheduledTaskStatus.PENDING, ScheduledTaskStatus.PAUSED]
+    );
+
+    return result.rows.map((row) => this.mapRowToTask(row));
+  }
+
+  /**
+   * Get tasks by type
+   */
+  async getTasksByType(taskType: string): Promise<ScheduledTaskDefinition[]> {
+    await this.ensureInitialized();
+
+    const result = await this.pool.query<ScheduledTaskRow>(
+      `SELECT * FROM lattice_scheduled_tasks WHERE task_type = $1 ORDER BY created_at DESC`,
+      [taskType]
+    );
+
+    return result.rows.map((row) => this.mapRowToTask(row));
+  }
+
+  /**
+   * Get tasks by status
+   */
+  async getTasksByStatus(
+    status: ScheduledTaskStatus
+  ): Promise<ScheduledTaskDefinition[]> {
+    await this.ensureInitialized();
+
+    const result = await this.pool.query<ScheduledTaskRow>(
+      `SELECT * FROM lattice_scheduled_tasks WHERE status = $1 ORDER BY created_at DESC`,
+      [status]
+    );
+
+    return result.rows.map((row) => this.mapRowToTask(row));
+  }
+
+  /**
+   * Get tasks by execution type
+   */
+  async getTasksByExecutionType(
+    executionType: ScheduleExecutionType
+  ): Promise<ScheduledTaskDefinition[]> {
+    await this.ensureInitialized();
+
+    const result = await this.pool.query<ScheduledTaskRow>(
+      `SELECT * FROM lattice_scheduled_tasks WHERE execution_type = $1 ORDER BY created_at DESC`,
+      [executionType]
+    );
+
+    return result.rows.map((row) => this.mapRowToTask(row));
+  }
+
+  /**
+   * Get tasks by assistant ID
+   */
+  async getTasksByAssistantId(
+    assistantId: string
+  ): Promise<ScheduledTaskDefinition[]> {
+    await this.ensureInitialized();
+
+    const result = await this.pool.query<ScheduledTaskRow>(
+      `SELECT * FROM lattice_scheduled_tasks WHERE assistant_id = $1 ORDER BY created_at DESC`,
+      [assistantId]
+    );
+
+    return result.rows.map((row) => this.mapRowToTask(row));
+  }
+
+  /**
+   * Get tasks by thread ID
+   */
+  async getTasksByThreadId(
+    threadId: string
+  ): Promise<ScheduledTaskDefinition[]> {
+    await this.ensureInitialized();
+
+    const result = await this.pool.query<ScheduledTaskRow>(
+      `SELECT * FROM lattice_scheduled_tasks WHERE thread_id = $1 ORDER BY created_at DESC`,
+      [threadId]
+    );
+
+    return result.rows.map((row) => this.mapRowToTask(row));
+  }
+
+  /**
+   * Get all tasks with optional filters
+   */
+  async getAllTasks(filters?: {
+    status?: ScheduledTaskStatus;
+    executionType?: ScheduleExecutionType;
+    taskType?: string;
+    assistantId?: string;
+    threadId?: string;
+    limit?: number;
+    offset?: number;
+  }): Promise<ScheduledTaskDefinition[]> {
+    await this.ensureInitialized();
+
+    const whereClauses: string[] = [];
+    const values: any[] = [];
+    let paramIndex = 1;
+
+    if (filters?.status !== undefined) {
+      whereClauses.push(`status = $${paramIndex++}`);
+      values.push(filters.status);
+    }
+    if (filters?.executionType !== undefined) {
+      whereClauses.push(`execution_type = $${paramIndex++}`);
+      values.push(filters.executionType);
+    }
+    if (filters?.taskType !== undefined) {
+      whereClauses.push(`task_type = $${paramIndex++}`);
+      values.push(filters.taskType);
+    }
+    if (filters?.assistantId !== undefined) {
+      whereClauses.push(`assistant_id = $${paramIndex++}`);
+      values.push(filters.assistantId);
+    }
+    if (filters?.threadId !== undefined) {
+      whereClauses.push(`thread_id = $${paramIndex++}`);
+      values.push(filters.threadId);
+    }
+
+    let query = `SELECT * FROM lattice_scheduled_tasks`;
+    if (whereClauses.length > 0) {
+      query += ` WHERE ${whereClauses.join(" AND ")}`;
+    }
+    query += ` ORDER BY created_at DESC`;
+
+    if (filters?.limit !== undefined) {
+      query += ` LIMIT $${paramIndex++}`;
+      values.push(filters.limit);
+    }
+    if (filters?.offset !== undefined) {
+      query += ` OFFSET $${paramIndex++}`;
+      values.push(filters.offset);
+    }
+
+    const result = await this.pool.query<ScheduledTaskRow>(query, values);
+
+    return result.rows.map((row) => this.mapRowToTask(row));
+  }
+
+  /**
+   * Count tasks with optional filters
+   */
+  async countTasks(filters?: {
+    status?: ScheduledTaskStatus;
+    executionType?: ScheduleExecutionType;
+    taskType?: string;
+    assistantId?: string;
+    threadId?: string;
+  }): Promise<number> {
+    await this.ensureInitialized();
+
+    const whereClauses: string[] = [];
+    const values: any[] = [];
+    let paramIndex = 1;
+
+    if (filters?.status !== undefined) {
+      whereClauses.push(`status = $${paramIndex++}`);
+      values.push(filters.status);
+    }
+    if (filters?.executionType !== undefined) {
+      whereClauses.push(`execution_type = $${paramIndex++}`);
+      values.push(filters.executionType);
+    }
+    if (filters?.taskType !== undefined) {
+      whereClauses.push(`task_type = $${paramIndex++}`);
+      values.push(filters.taskType);
+    }
+    if (filters?.assistantId !== undefined) {
+      whereClauses.push(`assistant_id = $${paramIndex++}`);
+      values.push(filters.assistantId);
+    }
+    if (filters?.threadId !== undefined) {
+      whereClauses.push(`thread_id = $${paramIndex++}`);
+      values.push(filters.threadId);
+    }
+
+    let query = `SELECT COUNT(*) as count FROM lattice_scheduled_tasks`;
+    if (whereClauses.length > 0) {
+      query += ` WHERE ${whereClauses.join(" AND ")}`;
+    }
+
+    const result = await this.pool.query<{ count: string }>(query, values);
+
+    return parseInt(result.rows[0].count, 10);
+  }
+
+  /**
+   * Delete completed/cancelled/failed tasks older than specified time
+   */
+  async deleteOldTasks(olderThanMs: number): Promise<number> {
+    await this.ensureInitialized();
+
+    const cutoff = new Date(Date.now() - olderThanMs);
+
+    const result = await this.pool.query(
+      `
+      DELETE FROM lattice_scheduled_tasks 
+      WHERE status IN ($1, $2, $3) 
+      AND updated_at < $4
+      `,
+      [
+        ScheduledTaskStatus.COMPLETED,
+        ScheduledTaskStatus.CANCELLED,
+        ScheduledTaskStatus.FAILED,
+        cutoff,
+      ]
+    );
+
+    return result.rowCount ?? 0;
+  }
+
+  /**
+   * Map database row to ScheduledTaskDefinition
+   */
+  private mapRowToTask(row: ScheduledTaskRow): ScheduledTaskDefinition {
+    return {
+      taskId: row.task_id,
+      taskType: row.task_type,
+      payload:
+        typeof row.payload === "string"
+          ? JSON.parse(row.payload)
+          : row.payload || {},
+      assistantId: row.assistant_id ?? undefined,
+      threadId: row.thread_id ?? undefined,
+      executionType: row.execution_type as ScheduleExecutionType,
+      executeAt: row.execute_at ? row.execute_at.getTime() : undefined,
+      delayMs: row.delay_ms ?? undefined,
+      cronExpression: row.cron_expression ?? undefined,
+      timezone: row.timezone ?? undefined,
+      nextRunAt: row.next_run_at ? row.next_run_at.getTime() : undefined,
+      lastRunAt: row.last_run_at ? row.last_run_at.getTime() : undefined,
+      status: row.status as ScheduledTaskStatus,
+      runCount: row.run_count,
+      maxRuns: row.max_runs ?? undefined,
+      retryCount: row.retry_count,
+      maxRetries: row.max_retries,
+      lastError: row.last_error ?? undefined,
+      createdAt: row.created_at.getTime(),
+      updatedAt: row.updated_at.getTime(),
+      expiresAt: row.expires_at ? row.expires_at.getTime() : undefined,
+      metadata:
+        row.metadata === null
+          ? undefined
+          : typeof row.metadata === "string"
+          ? JSON.parse(row.metadata)
+          : row.metadata,
+    };
+  }
+}