feat(jobs): persist async batch job state in PostgreSQL
- Add jobs table to database schema (job_id, status, progress, result_json, etc.) - Add DatabaseClient methods: create_job, update_job, get_job, list_jobs - Add mark_stale_jobs_failed() called at startup to handle interrupted jobs - Refactor _run_batch_job and job endpoints to read/write from PostgreSQL - Remove in-memory _jobs dict; job state now survives API restarts - Update init_database.py to list all tables in output Closes leeworks-agents/SPARC#8 Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -171,6 +171,26 @@ class DatabaseClient:
|
||||
ON serp_queries(query_hash)
|
||||
""")
|
||||
|
||||
# Create jobs table for persisting async batch job state
|
||||
cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS jobs (
|
||||
job_id VARCHAR(128) PRIMARY KEY,
|
||||
status VARCHAR(20) NOT NULL DEFAULT 'pending',
|
||||
progress INTEGER NOT NULL DEFAULT 0,
|
||||
total_companies INTEGER NOT NULL DEFAULT 0,
|
||||
completed_companies INTEGER NOT NULL DEFAULT 0,
|
||||
result_json JSONB,
|
||||
error TEXT,
|
||||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
||||
)
|
||||
""")
|
||||
|
||||
cursor.execute("""
|
||||
CREATE INDEX IF NOT EXISTS idx_jobs_status
|
||||
ON jobs(status)
|
||||
""")
|
||||
|
||||
self.conn.commit()
|
||||
|
||||
@staticmethod
|
||||
@@ -462,6 +482,131 @@ class DatabaseClient:
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
# Job Persistence Methods
|
||||
|
||||
def create_job(
|
||||
self,
|
||||
job_id: str,
|
||||
total_companies: int,
|
||||
) -> Dict:
|
||||
"""Create a new job record.
|
||||
|
||||
Args:
|
||||
job_id: Unique job identifier
|
||||
total_companies: Number of companies in the batch
|
||||
|
||||
Returns:
|
||||
Job dict
|
||||
"""
|
||||
with self.get_conn() as conn:
|
||||
with conn.cursor(cursor_factory=RealDictCursor) as cursor:
|
||||
cursor.execute(
|
||||
"""
|
||||
INSERT INTO jobs (job_id, status, progress, total_companies, completed_companies)
|
||||
VALUES (%s, 'pending', 0, %s, 0)
|
||||
RETURNING *
|
||||
""",
|
||||
(job_id, total_companies),
|
||||
)
|
||||
job = cursor.fetchone()
|
||||
conn.commit()
|
||||
return dict(job)
|
||||
|
||||
def update_job(
|
||||
self,
|
||||
job_id: str,
|
||||
status: Optional[str] = None,
|
||||
progress: Optional[int] = None,
|
||||
completed_companies: Optional[int] = None,
|
||||
result_json: Optional[str] = None,
|
||||
error: Optional[str] = None,
|
||||
) -> Optional[Dict]:
|
||||
"""Update a job's state.
|
||||
|
||||
Only non-None fields are updated.
|
||||
"""
|
||||
updates = []
|
||||
params = []
|
||||
if status is not None:
|
||||
updates.append("status = %s")
|
||||
params.append(status)
|
||||
if progress is not None:
|
||||
updates.append("progress = %s")
|
||||
params.append(progress)
|
||||
if completed_companies is not None:
|
||||
updates.append("completed_companies = %s")
|
||||
params.append(completed_companies)
|
||||
if result_json is not None:
|
||||
updates.append("result_json = %s")
|
||||
params.append(result_json)
|
||||
if error is not None:
|
||||
updates.append("error = %s")
|
||||
params.append(error)
|
||||
|
||||
if not updates:
|
||||
return self.get_job(job_id)
|
||||
|
||||
updates.append("updated_at = CURRENT_TIMESTAMP")
|
||||
params.append(job_id)
|
||||
|
||||
with self.get_conn() as conn:
|
||||
with conn.cursor(cursor_factory=RealDictCursor) as cursor:
|
||||
cursor.execute(
|
||||
f"UPDATE jobs SET {', '.join(updates)} WHERE job_id = %s RETURNING *",
|
||||
params,
|
||||
)
|
||||
job = cursor.fetchone()
|
||||
conn.commit()
|
||||
return dict(job) if job else None
|
||||
|
||||
def get_job(self, job_id: str) -> Optional[Dict]:
|
||||
"""Get a job by ID."""
|
||||
with self.get_conn() as conn:
|
||||
with conn.cursor(cursor_factory=RealDictCursor) as cursor:
|
||||
cursor.execute("SELECT * FROM jobs WHERE job_id = %s", (job_id,))
|
||||
job = cursor.fetchone()
|
||||
return dict(job) if job else None
|
||||
|
||||
def list_jobs(
|
||||
self,
|
||||
status: Optional[str] = None,
|
||||
limit: int = 10,
|
||||
) -> List[Dict]:
|
||||
"""List jobs, optionally filtered by status."""
|
||||
query = "SELECT * FROM jobs"
|
||||
params: list = []
|
||||
if status:
|
||||
query += " WHERE status = %s"
|
||||
params.append(status)
|
||||
query += " ORDER BY created_at DESC LIMIT %s"
|
||||
params.append(limit)
|
||||
|
||||
with self.get_conn() as conn:
|
||||
with conn.cursor(cursor_factory=RealDictCursor) as cursor:
|
||||
cursor.execute(query, params)
|
||||
return [dict(row) for row in cursor.fetchall()]
|
||||
|
||||
def mark_stale_jobs_failed(self) -> int:
|
||||
"""Mark any jobs in 'running' or 'pending' state as 'failed'.
|
||||
|
||||
Called at startup to clean up jobs that were interrupted by a restart.
|
||||
|
||||
Returns:
|
||||
Number of jobs marked as failed.
|
||||
"""
|
||||
with self.get_conn() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
cursor.execute(
|
||||
"""
|
||||
UPDATE jobs SET status = 'failed', error = 'Interrupted by server restart',
|
||||
updated_at = CURRENT_TIMESTAMP
|
||||
WHERE status IN ('running', 'pending')
|
||||
"""
|
||||
)
|
||||
count = cursor.rowcount
|
||||
conn.commit()
|
||||
return count
|
||||
|
||||
# User Authentication Methods
|
||||
|
||||
@staticmethod
|
||||
|
||||
Reference in New Issue
Block a user