@axiom-lattice/pg-stores 1.0.1 → 1.0.3

package/dist/index.mjs

@@ -573,11 +573,556 @@ var PostgreSQLAssistantStore = class {
     };
   }
 };
+
+// src/stores/PostgreSQLScheduleStorage.ts
+import { Pool as Pool3 } from "pg";
+import {
+  ScheduledTaskStatus
+} from "@axiom-lattice/protocols";
+
+// src/migrations/schedule_migrations.ts
+var createScheduledTasksTable = {
+  name: "create_scheduled_tasks_table",
+  version: 2,
+  up: async (client) => {
+    await client.query(`
+      CREATE TABLE IF NOT EXISTS lattice_scheduled_tasks (
+        -- Primary key
+        task_id VARCHAR(255) PRIMARY KEY,
+        
+        -- Task type (maps to handler)
+        task_type VARCHAR(255) NOT NULL,
+        
+        -- Payload (JSON-serializable data)
+        payload JSONB NOT NULL DEFAULT '{}',
+        
+        -- Context fields for querying
+        assistant_id VARCHAR(255),
+        thread_id VARCHAR(255),
+        
+        -- Execution type: 'once' or 'cron'
+        execution_type VARCHAR(50) NOT NULL,
+        
+        -- For ONCE type
+        execute_at TIMESTAMPTZ,
+        delay_ms BIGINT,
+        
+        -- For CRON type
+        cron_expression VARCHAR(255),
+        timezone VARCHAR(100),
+        next_run_at TIMESTAMPTZ,
+        last_run_at TIMESTAMPTZ,
+        
+        -- Status
+        status VARCHAR(50) NOT NULL DEFAULT 'pending',
+        
+        -- Execution tracking
+        run_count INTEGER NOT NULL DEFAULT 0,
+        max_runs INTEGER,
+        
+        -- Error handling
+        retry_count INTEGER NOT NULL DEFAULT 0,
+        max_retries INTEGER NOT NULL DEFAULT 0,
+        last_error TEXT,
+        
+        -- Timestamps
+        created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
+        updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
+        expires_at TIMESTAMPTZ,
+        
+        -- Metadata
+        metadata JSONB
+      );
+
+      -- Index for querying active tasks
+      CREATE INDEX IF NOT EXISTS idx_scheduled_tasks_status 
+        ON lattice_scheduled_tasks (status);
+      
+      -- Index for querying by task type
+      CREATE INDEX IF NOT EXISTS idx_scheduled_tasks_task_type 
+        ON lattice_scheduled_tasks (task_type);
+      
+      -- Index for querying by execution type
+      CREATE INDEX IF NOT EXISTS idx_scheduled_tasks_execution_type 
+        ON lattice_scheduled_tasks (execution_type);
+      
+      -- Index for querying by assistant ID
+      CREATE INDEX IF NOT EXISTS idx_scheduled_tasks_assistant_id 
+        ON lattice_scheduled_tasks (assistant_id)
+        WHERE assistant_id IS NOT NULL;
+      
+      -- Index for querying by thread ID
+      CREATE INDEX IF NOT EXISTS idx_scheduled_tasks_thread_id 
+        ON lattice_scheduled_tasks (thread_id)
+        WHERE thread_id IS NOT NULL;
+      
+      -- Index for querying pending tasks by next run time
+      CREATE INDEX IF NOT EXISTS idx_scheduled_tasks_next_run 
+        ON lattice_scheduled_tasks (next_run_at) 
+        WHERE status = 'pending';
+      
+      -- Index for querying pending one-time tasks by execute time
+      CREATE INDEX IF NOT EXISTS idx_scheduled_tasks_execute_at 
+        ON lattice_scheduled_tasks (execute_at) 
+        WHERE status = 'pending' AND execution_type = 'once';
+      
+      -- Index for cleanup queries (old completed/cancelled tasks)
+      CREATE INDEX IF NOT EXISTS idx_scheduled_tasks_cleanup 
+        ON lattice_scheduled_tasks (updated_at) 
+        WHERE status IN ('completed', 'cancelled', 'failed');
+    `);
+  },
+  down: async (client) => {
+    await client.query(`
+      DROP INDEX IF EXISTS idx_scheduled_tasks_cleanup;
+      DROP INDEX IF EXISTS idx_scheduled_tasks_execute_at;
+      DROP INDEX IF EXISTS idx_scheduled_tasks_next_run;
+      DROP INDEX IF EXISTS idx_scheduled_tasks_thread_id;
+      DROP INDEX IF EXISTS idx_scheduled_tasks_assistant_id;
+      DROP INDEX IF EXISTS idx_scheduled_tasks_execution_type;
+      DROP INDEX IF EXISTS idx_scheduled_tasks_task_type;
+      DROP INDEX IF EXISTS idx_scheduled_tasks_status;
+      DROP TABLE IF EXISTS lattice_scheduled_tasks;
+    `);
+  }
+};
+
+// src/stores/PostgreSQLScheduleStorage.ts
+var PostgreSQLScheduleStorage = class {
+  constructor(options) {
+    this.initialized = false;
+    if (typeof options.poolConfig === "string") {
+      this.pool = new Pool3({ connectionString: options.poolConfig });
+    } else {
+      this.pool = new Pool3(options.poolConfig);
+    }
+    this.migrationManager = new MigrationManager(this.pool);
+    this.migrationManager.register(createScheduledTasksTable);
+    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() {
+    if (this.pool) {
+      await this.pool.end();
+    }
+  }
+  /**
+   * Initialize the store and run migrations
+   */
+  async initialize() {
+    if (this.initialized) {
+      return;
+    }
+    await this.migrationManager.migrate();
+    this.initialized = true;
+  }
+  /**
+   * Ensure store is initialized
+   */
+  async ensureInitialized() {
+    if (!this.initialized) {
+      await this.initialize();
+    }
+  }
+  /**
+   * Save a new task
+   */
+  async save(task) {
+    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) {
+    await this.ensureInitialized();
+    const result = await this.pool.query(
+      `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, updates) {
+    await this.ensureInitialized();
+    const setClauses = [];
+    const values = [];
+    let paramIndex = 1;
+    if (updates.taskType !== void 0) {
+      setClauses.push(`task_type = $${paramIndex++}`);
+      values.push(updates.taskType);
+    }
+    if (updates.payload !== void 0) {
+      setClauses.push(`payload = $${paramIndex++}`);
+      values.push(JSON.stringify(updates.payload));
+    }
+    if (updates.assistantId !== void 0) {
+      setClauses.push(`assistant_id = $${paramIndex++}`);
+      values.push(updates.assistantId);
+    }
+    if (updates.threadId !== void 0) {
+      setClauses.push(`thread_id = $${paramIndex++}`);
+      values.push(updates.threadId);
+    }
+    if (updates.executionType !== void 0) {
+      setClauses.push(`execution_type = $${paramIndex++}`);
+      values.push(updates.executionType);
+    }
+    if (updates.executeAt !== void 0) {
+      setClauses.push(`execute_at = $${paramIndex++}`);
+      values.push(updates.executeAt ? new Date(updates.executeAt) : null);
+    }
+    if (updates.delayMs !== void 0) {
+      setClauses.push(`delay_ms = $${paramIndex++}`);
+      values.push(updates.delayMs);
+    }
+    if (updates.cronExpression !== void 0) {
+      setClauses.push(`cron_expression = $${paramIndex++}`);
+      values.push(updates.cronExpression);
+    }
+    if (updates.timezone !== void 0) {
+      setClauses.push(`timezone = $${paramIndex++}`);
+      values.push(updates.timezone);
+    }
+    if (updates.nextRunAt !== void 0) {
+      setClauses.push(`next_run_at = $${paramIndex++}`);
+      values.push(updates.nextRunAt ? new Date(updates.nextRunAt) : null);
+    }
+    if (updates.lastRunAt !== void 0) {
+      setClauses.push(`last_run_at = $${paramIndex++}`);
+      values.push(updates.lastRunAt ? new Date(updates.lastRunAt) : null);
+    }
+    if (updates.status !== void 0) {
+      setClauses.push(`status = $${paramIndex++}`);
+      values.push(updates.status);
+    }
+    if (updates.runCount !== void 0) {
+      setClauses.push(`run_count = $${paramIndex++}`);
+      values.push(updates.runCount);
+    }
+    if (updates.maxRuns !== void 0) {
+      setClauses.push(`max_runs = $${paramIndex++}`);
+      values.push(updates.maxRuns);
+    }
+    if (updates.retryCount !== void 0) {
+      setClauses.push(`retry_count = $${paramIndex++}`);
+      values.push(updates.retryCount);
+    }
+    if (updates.maxRetries !== void 0) {
+      setClauses.push(`max_retries = $${paramIndex++}`);
+      values.push(updates.maxRetries);
+    }
+    if (updates.lastError !== void 0) {
+      setClauses.push(`last_error = $${paramIndex++}`);
+      values.push(updates.lastError);
+    }
+    if (updates.expiresAt !== void 0) {
+      setClauses.push(`expires_at = $${paramIndex++}`);
+      values.push(updates.expiresAt ? new Date(updates.expiresAt) : null);
+    }
+    if (updates.metadata !== void 0) {
+      setClauses.push(`metadata = $${paramIndex++}`);
+      values.push(updates.metadata ? JSON.stringify(updates.metadata) : null);
+    }
+    setClauses.push(`updated_at = $${paramIndex++}`);
+    values.push(/* @__PURE__ */ new Date());
+    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) {
+    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() {
+    await this.ensureInitialized();
+    const result = await this.pool.query(
+      `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) {
+    await this.ensureInitialized();
+    const result = await this.pool.query(
+      `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) {
+    await this.ensureInitialized();
+    const result = await this.pool.query(
+      `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) {
+    await this.ensureInitialized();
+    const result = await this.pool.query(
+      `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) {
+    await this.ensureInitialized();
+    const result = await this.pool.query(
+      `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) {
+    await this.ensureInitialized();
+    const result = await this.pool.query(
+      `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) {
+    await this.ensureInitialized();
+    const whereClauses = [];
+    const values = [];
+    let paramIndex = 1;
+    if (filters?.status !== void 0) {
+      whereClauses.push(`status = $${paramIndex++}`);
+      values.push(filters.status);
+    }
+    if (filters?.executionType !== void 0) {
+      whereClauses.push(`execution_type = $${paramIndex++}`);
+      values.push(filters.executionType);
+    }
+    if (filters?.taskType !== void 0) {
+      whereClauses.push(`task_type = $${paramIndex++}`);
+      values.push(filters.taskType);
+    }
+    if (filters?.assistantId !== void 0) {
+      whereClauses.push(`assistant_id = $${paramIndex++}`);
+      values.push(filters.assistantId);
+    }
+    if (filters?.threadId !== void 0) {
+      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 !== void 0) {
+      query += ` LIMIT $${paramIndex++}`;
+      values.push(filters.limit);
+    }
+    if (filters?.offset !== void 0) {
+      query += ` OFFSET $${paramIndex++}`;
+      values.push(filters.offset);
+    }
+    const result = await this.pool.query(query, values);
+    return result.rows.map((row) => this.mapRowToTask(row));
+  }
+  /**
+   * Count tasks with optional filters
+   */
+  async countTasks(filters) {
+    await this.ensureInitialized();
+    const whereClauses = [];
+    const values = [];
+    let paramIndex = 1;
+    if (filters?.status !== void 0) {
+      whereClauses.push(`status = $${paramIndex++}`);
+      values.push(filters.status);
+    }
+    if (filters?.executionType !== void 0) {
+      whereClauses.push(`execution_type = $${paramIndex++}`);
+      values.push(filters.executionType);
+    }
+    if (filters?.taskType !== void 0) {
+      whereClauses.push(`task_type = $${paramIndex++}`);
+      values.push(filters.taskType);
+    }
+    if (filters?.assistantId !== void 0) {
+      whereClauses.push(`assistant_id = $${paramIndex++}`);
+      values.push(filters.assistantId);
+    }
+    if (filters?.threadId !== void 0) {
+      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(query, values);
+    return parseInt(result.rows[0].count, 10);
+  }
+  /**
+   * Delete completed/cancelled/failed tasks older than specified time
+   */
+  async deleteOldTasks(olderThanMs) {
+    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
+   */
+  mapRowToTask(row) {
+    return {
+      taskId: row.task_id,
+      taskType: row.task_type,
+      payload: typeof row.payload === "string" ? JSON.parse(row.payload) : row.payload || {},
+      assistantId: row.assistant_id ?? void 0,
+      threadId: row.thread_id ?? void 0,
+      executionType: row.execution_type,
+      executeAt: row.execute_at ? row.execute_at.getTime() : void 0,
+      delayMs: row.delay_ms ?? void 0,
+      cronExpression: row.cron_expression ?? void 0,
+      timezone: row.timezone ?? void 0,
+      nextRunAt: row.next_run_at ? row.next_run_at.getTime() : void 0,
+      lastRunAt: row.last_run_at ? row.last_run_at.getTime() : void 0,
+      status: row.status,
+      runCount: row.run_count,
+      maxRuns: row.max_runs ?? void 0,
+      retryCount: row.retry_count,
+      maxRetries: row.max_retries,
+      lastError: row.last_error ?? void 0,
+      createdAt: row.created_at.getTime(),
+      updatedAt: row.updated_at.getTime(),
+      expiresAt: row.expires_at ? row.expires_at.getTime() : void 0,
+      metadata: row.metadata === null ? void 0 : typeof row.metadata === "string" ? JSON.parse(row.metadata) : row.metadata
+    };
+  }
+};
 export {
   MigrationManager,
   PostgreSQLAssistantStore,
+  PostgreSQLScheduleStorage,
   PostgreSQLThreadStore,
   createAssistantsTable,
+  createScheduledTasksTable,
   createThreadsTable
 };
 //# sourceMappingURL=index.mjs.map