job_manager.py

import psycopg2
from psycopg2.extras import RealDictCursor, Json
from datetime import datetime
import json


class JobManager:
    def __init__(self, host='localhost', database='job_registry',
                 user='postgres', password='labpassword'):
        """Initialize database connection."""
        self.conn = psycopg2.connect(
            host=host,
            database=database,
            user=user,
            password=password
        )
        self.conn.autocommit = False
        self.cursor = self.conn.cursor(cursor_factory=RealDictCursor)

    def create_job(self, job_name: str) -> int:
        """
        Create a new job entry with 'pending' status.

        Args:
            job_name: Name of the job to create

        Returns:
            job_id: ID of the created job
        """
        try:
            self.cursor.execute(
                """
                INSERT INTO jobs (job_name, status, created_at)
                VALUES (%s, %s, %s)
                RETURNING job_id;
                """,
                (job_name, 'pending', datetime.now())
            )
            job_id = self.cursor.fetchone()['job_id']
            self.conn.commit()
            return job_id
        except Exception:
            self.conn.rollback()
            raise

    def start_job(self, job_id: int) -> bool:
        """
        Mark job as running and set started_at timestamp.

        Args:
            job_id: ID of the job to start

        Returns:
            bool: True if successful, False otherwise
        """
        try:
            self.cursor.execute(
                """
                UPDATE jobs
                SET status = %s,
                    started_at = %s
                WHERE job_id = %s;
                """,
                ('running', datetime.now(), job_id)
            )
            success = self.cursor.rowcount > 0
            self.conn.commit()
            return success
        except Exception:
            self.conn.rollback()
            return False

    def complete_job(self, job_id: int, result_data: dict) -> bool:
        """
        Mark job as completed and store results.

        Args:
            job_id: ID of the job to complete
            result_data: Dictionary containing job results

        Returns:
            bool: True if successful, False otherwise
        """
        try:
            self.cursor.execute(
                """
                UPDATE jobs
                SET status = %s,
                    completed_at = %s,
                    result_data = %s
                WHERE job_id = %s;
                """,
                ('completed', datetime.now(), Json(result_data), job_id)
            )
            success = self.cursor.rowcount > 0
            self.conn.commit()
            return success
        except Exception:
            self.conn.rollback()
            return False

    def fail_job(self, job_id: int, error_message: str) -> bool:
        """
        Mark job as failed and store error message.

        Args:
            job_id: ID of the job that failed
            error_message: Error description

        Returns:
            bool: True if successful, False otherwise
        """
        try:
            self.cursor.execute(
                """
                UPDATE jobs
                SET status = %s,
                    completed_at = %s,
                    error_message = %s
                WHERE job_id = %s;
                """,
                ('failed', datetime.now(), error_message, job_id)
            )
            success = self.cursor.rowcount > 0
            self.conn.commit()
            return success
        except Exception:
            self.conn.rollback()
            return False

    def get_job_status(self, job_id: int) -> dict:
        """
        Retrieve current job status and metadata.

        Args:
            job_id: ID of the job to query

        Returns:
            dict: Job metadata including status, timestamps, and results
        """
        self.cursor.execute(
            """
            SELECT job_id, job_name, status, created_at, started_at, completed_at,
                   result_data, error_message
            FROM jobs
            WHERE job_id = %s;
            """,
            (job_id,)
        )
        result = self.cursor.fetchone()
        return dict(result) if result else None

    def get_jobs_by_status(self, status: str) -> list:
        """
        Retrieve all jobs with specified status.

        Args:
            status: Job status to filter by

        Returns:
            list: List of job dictionaries
        """
        self.cursor.execute(
            """
            SELECT job_id, job_name, status, created_at, started_at, completed_at,
                   result_data, error_message
            FROM jobs
            WHERE status = %s
            ORDER BY job_id;
            """,
            (status,)
        )
        results = self.cursor.fetchall()
        return [dict(row) for row in results]

    def get_job_duration(self, job_id: int) -> float:
        """
        Calculate job execution duration in seconds.

        Args:
            job_id: ID of the job

        Returns:
            float: Duration in seconds, or None if not completed
        """
        self.cursor.execute(
            """
            SELECT started_at, completed_at
            FROM jobs
            WHERE job_id = %s;
            """,
            (job_id,)
        )
        result = self.cursor.fetchone()

        if not result:
            return None

        started_at = result['started_at']
        completed_at = result['completed_at']

        if started_at is None or completed_at is None:
            return None

        duration = (completed_at - started_at).total_seconds()
        return duration

    def close(self):
        """Close database connection."""
        if hasattr(self, 'cursor') and self.cursor:
            self.cursor.close()
        if hasattr(self, 'conn') and self.conn:
            self.conn.close()