244 lines
9.2 KiB
Python
244 lines
9.2 KiB
Python
import uuid
|
|
import logging
|
|
from typing import Dict, Optional, List
|
|
from datetime import datetime, timedelta
|
|
import threading
|
|
import json
|
|
from sqlalchemy.orm import Session
|
|
from sqlalchemy import desc
|
|
|
|
from ..services.postgresql_manager import PostgreSQLManager
|
|
from ..utils.config import config
|
|
from ..models.job import Job
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class JobManager:
|
|
_instance = None
|
|
|
|
def __new__(cls):
|
|
if cls._instance is None:
|
|
cls._instance = super(JobManager, cls).__new__(cls)
|
|
cls._instance.db_manager = PostgreSQLManager(config.DATABASE_URL)
|
|
# We still keep active_jobs in memory for simple locking/status
|
|
# But the detailed state and history will be DB backed
|
|
cls._instance.job_lock = threading.Lock()
|
|
return cls._instance
|
|
|
|
def _get_db(self):
|
|
return self.db_manager.get_db_session()
|
|
|
|
def run_serialized(self, job_id: str, func, *args, **kwargs):
|
|
"""Run a function with a global lock to ensure serial execution."""
|
|
if self.should_cancel(job_id):
|
|
self.update_job(job_id, status="cancelled", message="Cancelled before start")
|
|
return
|
|
|
|
self.update_job(job_id, message="Queued (Waiting for lock)...")
|
|
|
|
with self.job_lock:
|
|
if self.should_cancel(job_id):
|
|
self.update_job(job_id, status="cancelled", message="Cancelled while queued")
|
|
return
|
|
|
|
self.update_job(job_id, message="Starting...")
|
|
try:
|
|
func(job_id, *args, **kwargs)
|
|
except Exception as e:
|
|
logger.error(f"Error in serialized job {job_id}: {e}")
|
|
self.fail_job(job_id, str(e))
|
|
|
|
def request_pause(self, job_id: str) -> bool:
|
|
with self._get_db() as db:
|
|
job = db.query(Job).filter(Job.id == job_id).first()
|
|
if job and job.status == 'running':
|
|
job.paused = True
|
|
job.status = 'paused'
|
|
job.message = "Paused..."
|
|
db.commit()
|
|
return True
|
|
return False
|
|
|
|
def resume_job(self, job_id: str) -> bool:
|
|
with self._get_db() as db:
|
|
job = db.query(Job).filter(Job.id == job_id).first()
|
|
if job and job.paused:
|
|
job.paused = False
|
|
job.status = 'running'
|
|
job.message = "Resuming..."
|
|
db.commit()
|
|
return True
|
|
return False
|
|
|
|
def should_pause(self, job_id: str) -> bool:
|
|
with self._get_db() as db:
|
|
job = db.query(Job).filter(Job.id == job_id).first()
|
|
return job.paused if job else False
|
|
|
|
def create_job(self, operation: str) -> str:
|
|
job_id = str(uuid.uuid4())
|
|
new_job = Job(
|
|
id=job_id,
|
|
operation=operation,
|
|
status="running",
|
|
start_time=datetime.now(),
|
|
progress=0,
|
|
message="Starting..."
|
|
)
|
|
|
|
with self._get_db() as db:
|
|
db.add(new_job)
|
|
db.commit()
|
|
|
|
logger.info(f"Created job {job_id} for {operation}")
|
|
return job_id
|
|
|
|
def _cleanup_jobs(self):
|
|
"""Delete jobs older than 30 days."""
|
|
try:
|
|
with self._get_db() as db:
|
|
cutoff = datetime.now() - timedelta(days=30)
|
|
db.query(Job).filter(Job.start_time < cutoff).delete()
|
|
db.commit()
|
|
except Exception as e:
|
|
logger.error(f"Error cleaning up jobs: {e}")
|
|
|
|
def get_job_history(self, limit: int = 10, offset: int = 0) -> Dict:
|
|
# self._cleanup_jobs() # Optional: Run periodically, running on every fetch is okay-ish but maybe expensive?
|
|
# Let's run it async or less frequently? For now, run it here is safe.
|
|
self._cleanup_jobs()
|
|
|
|
with self._get_db() as db:
|
|
# Sort desc by start_time
|
|
query = db.query(Job).order_by(desc(Job.start_time))
|
|
total = query.count()
|
|
jobs = query.offset(offset).limit(limit).all()
|
|
|
|
# Convert to dict
|
|
items = []
|
|
for j in jobs:
|
|
items.append({
|
|
"id": j.id,
|
|
"operation": j.operation,
|
|
"status": j.status,
|
|
"start_time": j.start_time.isoformat() if j.start_time else None,
|
|
"end_time": j.end_time.isoformat() if j.end_time else None,
|
|
"completed_at": j.end_time.isoformat() if j.end_time else None, # Compatibility
|
|
"duration_s": round((j.end_time - j.start_time).total_seconds(), 2) if j.end_time and j.start_time else None,
|
|
"progress": j.progress,
|
|
"message": j.message,
|
|
"result": j.result
|
|
})
|
|
|
|
return {"total": total, "items": items}
|
|
|
|
def get_job(self, job_id: str) -> Optional[Dict]:
|
|
with self._get_db() as db:
|
|
j = db.query(Job).filter(Job.id == job_id).first()
|
|
if j:
|
|
return {
|
|
"id": j.id,
|
|
"operation": j.operation,
|
|
"status": j.status,
|
|
"start_time": j.start_time,
|
|
"progress": j.progress,
|
|
"message": j.message,
|
|
"paused": j.paused,
|
|
"cancel_requested": j.cancel_requested
|
|
}
|
|
return None
|
|
|
|
def get_active_jobs(self) -> List[Dict]:
|
|
with self._get_db() as db:
|
|
active_jobs = db.query(Job).filter(Job.status.in_(['running', 'queued', 'paused'])).all()
|
|
return [{
|
|
"id": j.id,
|
|
"operation": j.operation,
|
|
"status": j.status,
|
|
"start_time": j.start_time,
|
|
"progress": j.progress,
|
|
"message": j.message,
|
|
"paused": j.paused,
|
|
"cancel_requested": j.cancel_requested
|
|
} for j in active_jobs]
|
|
|
|
def update_job(self, job_id: str, status: str = None, progress: int = None, message: str = None):
|
|
with self._get_db() as db:
|
|
job = db.query(Job).filter(Job.id == job_id).first()
|
|
if job:
|
|
if status:
|
|
job.status = status
|
|
if status in ["completed", "failed", "cancelled"] and not job.end_time:
|
|
job.end_time = datetime.now()
|
|
if progress is not None:
|
|
job.progress = progress
|
|
if message:
|
|
job.message = message
|
|
db.commit()
|
|
|
|
def request_cancel(self, job_id: str) -> bool:
|
|
with self._get_db() as db:
|
|
job = db.query(Job).filter(Job.id == job_id).first()
|
|
if job and job.status in ['running', 'queued', 'paused']:
|
|
# If paused, it's effectively stopped, so cancel immediately
|
|
if job.status == 'paused':
|
|
job.status = 'cancelled'
|
|
job.message = "Cancelled (while paused)"
|
|
job.end_time = datetime.now()
|
|
else:
|
|
job.cancel_requested = True
|
|
job.message = "Cancelling..."
|
|
|
|
db.commit()
|
|
logger.info(f"Cancellation requested for job {job_id}")
|
|
return True
|
|
return False
|
|
|
|
def should_cancel(self, job_id: str) -> bool:
|
|
with self._get_db() as db:
|
|
job = db.query(Job).filter(Job.id == job_id).first()
|
|
return job.cancel_requested if job else False
|
|
|
|
def complete_job(self, job_id: str, result: Dict = None):
|
|
with self._get_db() as db:
|
|
job = db.query(Job).filter(Job.id == job_id).first()
|
|
if job:
|
|
job.status = "completed"
|
|
job.progress = 100
|
|
job.message = "Completed"
|
|
job.end_time = datetime.now()
|
|
if result:
|
|
job.result = result
|
|
db.commit()
|
|
|
|
def fail_job(self, job_id: str, error: str):
|
|
with self._get_db() as db:
|
|
job = db.query(Job).filter(Job.id == job_id).first()
|
|
if job:
|
|
job.status = "failed"
|
|
job.message = error
|
|
job.end_time = datetime.now()
|
|
db.commit()
|
|
|
|
def force_fail_job(self, job_id: str):
|
|
"""
|
|
Forcefully mark a job as failed in the database.
|
|
This does not guarantee the underlying thread stops immediately,
|
|
but it releases the UI state.
|
|
"""
|
|
with self._get_db() as db:
|
|
job = db.query(Job).filter(Job.id == job_id).first()
|
|
if job:
|
|
# We update status regardless of current state if user wants to force it
|
|
prev_status = job.status
|
|
job.status = "failed"
|
|
job.message = f"Forcefully killed by user (was {prev_status})"
|
|
job.end_time = datetime.now()
|
|
job.cancel_requested = True # Hint to thread if it's still alive
|
|
db.commit()
|
|
logger.warning(f"Job {job_id} was forcefully killed by user.")
|
|
return True
|
|
return False
|
|
|
|
job_manager = JobManager()
|