cms-SiteForge/app/jobs.py

'''Durable build queue: SQLite single-host or PostgreSQL with DB-owned leases.

Redis is an optional wake-up signal, not an authoritative queue. PostgreSQL
lock/lease semantics require actual multi-worker acceptance on infrastructure.
'''
from __future__ import annotations
import hashlib
import json
import sqlite3
import os
import uuid
from datetime import datetime, timedelta, timezone
from fastapi import HTTPException
from . import db

class JobConflict(ValueError): pass

class JobCapacityError(ValueError): pass

def _limit(name, default, minimum=1, maximum=10000):
    try:return max(minimum,min(maximum,int(os.environ.get(name, default))))
    except ValueError:return default

def iso_after(seconds: int) -> str:
    return (datetime.now(timezone.utc)+timedelta(seconds=seconds)).isoformat(timespec='seconds')

def migrate():
    with db.connect() as c:
        existing={r['name'] for r in c.execute('PRAGMA table_info(build_jobs)')}
        columns={
          'payload_json':"TEXT NOT NULL DEFAULT '{}'",
          'idempotency_key':"TEXT",
          'max_attempts':"INTEGER NOT NULL DEFAULT 3",
          'timeout_seconds':"INTEGER NOT NULL DEFAULT 300",
          'available_at':"TEXT",
          'cancel_requested':"INTEGER NOT NULL DEFAULT 0",
          'worker_id':"TEXT",
          'lease_until':"TEXT",
          'heartbeat_at':"TEXT",
        }
        for column,definition in columns.items():
            if column not in existing:c.execute(f'ALTER TABLE build_jobs ADD COLUMN {column} {definition}')
        c.execute("CREATE UNIQUE INDEX IF NOT EXISTS idx_build_key ON build_jobs(site_id,idempotency_key) WHERE idempotency_key IS NOT NULL")
        c.execute("CREATE INDEX IF NOT EXISTS idx_build_available ON build_jobs(state,available_at,created_at)")

def submit(site_id, payload:dict, actor:str, key:str|None=None) -> dict:
    normalized=json.dumps({'mirror_images':bool(payload.get('mirror_images',False))},sort_keys=True,separators=(',',':'))
    if key is not None:
        if not (8<=len(key)<=128 and all(c.isalnum() or c in '_-:.' for c in key)):
            raise JobConflict('幂等键必须为 8–128 位字母、数字或 _-:.')
    # No client key: de-duplicate pending jobs with identical payload for a site.
    with db.connect() as c:
        c.execute('BEGIN IMMEDIATE')
        if key:
            prior=c.execute('SELECT * FROM build_jobs WHERE site_id=? AND idempotency_key=?',(site_id,key)).fetchone()
            if prior:
                if prior['payload_json']!=normalized:raise JobConflict('同一幂等键不能对应不同构建参数')
                return {'id':prior['id'],'state':prior['state'],'existing':True}
        else:
            prior=c.execute("SELECT id,state FROM build_jobs WHERE site_id=? AND kind='build' AND state IN ('queued','running','retrying') AND payload_json=? ORDER BY created_at LIMIT 1",(site_id,normalized)).fetchone()
            if prior:return {'id':prior['id'],'state':prior['state'],'existing':True}
        # Capacity gates run inside the same immediate write transaction as
        # INSERT so simultaneous callers cannot bypass the limits.
        global_cap=_limit('SITEFORGE_MAX_PENDING_JOBS',100)
        site_cap=_limit('SITEFORGE_MAX_PENDING_JOBS_PER_SITE',5)
        active="state IN ('queued','running','retrying')"
        global_count=c.execute(f'SELECT COUNT(*) FROM build_jobs WHERE {active}').fetchone()[0]
        site_count=c.execute(f'SELECT COUNT(*) FROM build_jobs WHERE {active} AND site_id=?',(site_id,)).fetchone()[0]
        if global_count>=global_cap:raise JobCapacityError('全局待处理任务达到上限')
        if site_count>=site_cap:raise JobCapacityError('该网站待处理任务达到上限')
        job_id=uuid.uuid4().hex[:16]
        max_attempts=max(1,min(6,int(os.getenv('SITEFORGE_BUILD_MAX_ATTEMPTS','3'))))
        timeout=max(10,min(3600,int(os.getenv('SITEFORGE_BUILD_TIMEOUT_SECONDS','300'))))
        c.execute('''INSERT INTO build_jobs(id,site_id,kind,state,created_by,created_at,
            payload_json,idempotency_key,max_attempts,timeout_seconds,available_at)
            VALUES (?,?,?,?,?,?,?,?,?,?,?)''',(job_id,site_id,'build','queued',actor,db.now(),normalized,key,max_attempts,timeout,db.now()))
    from . import redis_signal
    redis_signal.notify()
    return {'id':job_id,'state':'queued','existing':False}

def cancel(site_id,job_id):
    with db.connect() as c:
        c.execute('BEGIN IMMEDIATE')
        row=c.execute('SELECT state FROM build_jobs WHERE site_id=? AND id=?',(site_id,job_id)).fetchone()
        if row is None:raise HTTPException(404,'任务不存在')
        if row['state'] in ('done','failed','cancelled'):
            return {'id':job_id,'state':row['state'],'cancelled':False}
        if row['state'] in ('queued','retrying'):
            c.execute("UPDATE build_jobs SET state='cancelled',cancel_requested=1,finished_at=? WHERE id=?",(db.now(),job_id))
            return {'id':job_id,'state':'cancelled','cancelled':True}
        c.execute('UPDATE build_jobs SET cancel_requested=1 WHERE id=?',(job_id,))
        return {'id':job_id,'state':'cancelling','cancelled':True}

def recover_expired():
    # Claims are recoverable after a worker crashes. Running jobs are never
    # blindly duplicated while their lease is live.
    with db.connect() as c:
        c.execute('BEGIN IMMEDIATE')
        rows=c.execute("SELECT * FROM build_jobs WHERE state='running' AND (lease_until IS NULL OR lease_until < ?)",(db.now(),)).fetchall()
        for row in rows:
            if row['cancel_requested']:
                state='cancelled'
            else:
                # Publishing is an external side effect: it may have completed remotely
                # even though the Worker lost its lease. Never automatically replay it.
                state=('failed' if row['kind'] in ('remote_deploy','remote_rollback')
                       else ('retrying' if row['attempts']<row['max_attempts'] else 'failed'))
            message=('Release lease expired; reconcile remote release before manual retry'
                     if row['kind'] in ('remote_deploy','remote_rollback') else
                     'Worker lease expired; verify partial builds')
            c.execute('UPDATE build_jobs SET state=?,worker_id=NULL,lease_until=NULL,available_at=?,error=? WHERE id=?',
                     (state,db.now(),message,row['id']))
    return len(rows)


def submit_migration(site_id,session_id,actor,key=None):
    # Queue WordPress migration through the shared lease-based Worker.
    from .wp_migration_v81 import _one
    _one(site_id,session_id)
    if key is not None and not (8<=len(key)<=128 and all(c.isalnum() or c in '_-:.' for c in key)):
        raise JobConflict('幂等键必须为 8–128 位字母、数字或 _-:.')
    payload=json.dumps({'session_id':session_id},sort_keys=True,separators=(',',':'))
    with db.connect() as c:
        c.execute('BEGIN IMMEDIATE')
        if key:
            row=c.execute('SELECT * FROM build_jobs WHERE site_id=? AND idempotency_key=?',(site_id,key)).fetchone()
            if row:
                if row['kind']!='wordpress_scan' or row['payload_json']!=payload:
                    raise JobConflict('同一幂等键不能对应不同迁移任务')
                return {'id':row['id'],'state':row['state'],'existing':True}
        row=c.execute("SELECT id,state FROM build_jobs WHERE kind='wordpress_scan' AND site_id=? AND payload_json=? AND state IN ('queued','running','retrying') ORDER BY created_at LIMIT 1",(site_id,payload)).fetchone()
        if row:return {'id':row['id'],'state':row['state'],'existing':True}
        global_cap=_limit('SITEFORGE_MAX_PENDING_JOBS',100)
        site_cap=_limit('SITEFORGE_MAX_PENDING_JOBS_PER_SITE',5)
        if c.execute("SELECT COUNT(*) FROM build_jobs WHERE state IN ('queued','running','retrying')").fetchone()[0]>=global_cap:
            raise JobCapacityError('全局待处理任务达到上限')
        if c.execute("SELECT COUNT(*) FROM build_jobs WHERE site_id=? AND state IN ('queued','running','retrying')",(site_id,)).fetchone()[0]>=site_cap:
            raise JobCapacityError('该网站待处理任务达到上限')
        job_id=uuid.uuid4().hex[:16]
        attempts=max(1,min(6,int(os.getenv('SITEFORGE_WP_SCAN_MAX_ATTEMPTS','3'))))
        timeout=max(60,min(3600,int(os.getenv('SITEFORGE_WP_SCAN_TIMEOUT_SECONDS','900'))))
        c.execute('''INSERT INTO build_jobs(id,site_id,kind,state,created_by,created_at,
          payload_json,idempotency_key,max_attempts,timeout_seconds,available_at)
          VALUES(?,?,?,?,?,?,?,?,?,?,?)''',(job_id,site_id,'wordpress_scan','queued',actor,
           db.now(),payload,key,attempts,timeout,db.now()))
    from . import redis_signal
    redis_signal.notify()
    return {'id':job_id,'state':'queued','existing':False}


def submit_delivery(site_id,kind,payload,actor,key=None):
    # One logical release per job; serialized with build jobs by Worker.
    if kind not in ('remote_deploy','remote_rollback'):
        raise JobConflict('Unknown delivery job kind')
    if key is not None and not (8<=len(key)<=128 and all(c.isalnum() or c in '_-:.' for c in key)):
        raise JobConflict('无效的幂等键')
    normalized=json.dumps(payload,sort_keys=True,separators=(',',':'))
    with db.connect() as c:
        c.execute('BEGIN IMMEDIATE')
        if key:
            row=c.execute('SELECT * FROM build_jobs WHERE site_id=? AND idempotency_key=?',(site_id,key)).fetchone()
            if row:
                if row['kind']!=kind or row['payload_json']!=normalized:raise JobConflict('幂等键存在冲突')
                return {'id':row['id'],'state':row['state'],'existing':True}
        row=c.execute("SELECT id,state FROM build_jobs WHERE site_id=? AND kind=? AND payload_json=? AND state IN ('queued','running','retrying') ORDER BY created_at LIMIT 1",(site_id,kind,normalized)).fetchone()
        if row:return {'id':row['id'],'state':row['state'],'existing':True}
        active="state IN ('queued','running','retrying')"
        if c.execute(f'SELECT COUNT(*) FROM build_jobs WHERE {active}').fetchone()[0]>=_limit('SITEFORGE_MAX_PENDING_JOBS',100):
            raise JobCapacityError('系统任务队列已满')
        if c.execute(f'SELECT COUNT(*) FROM build_jobs WHERE site_id=? AND {active}',(site_id,)).fetchone()[0]>=_limit('SITEFORGE_MAX_PENDING_JOBS_PER_SITE',5):
            raise JobCapacityError('网站任务队列已满')
        job_id=uuid.uuid4().hex[:16]
        c.execute('''INSERT INTO build_jobs(id,site_id,kind,state,created_by,created_at,payload_json,
          idempotency_key,max_attempts,timeout_seconds,available_at)
          VALUES(?,?,?,?,?,?,?,?,?,?,?)''',
          (job_id,site_id,kind,'queued',actor,db.now(),normalized,key,3,900,db.now()))
    from . import redis_signal
    redis_signal.notify()
    return {'id':job_id,'state':'queued','existing':False}