worker_service.py

import requests
import time
import json
import subprocess
import os
import shutil
from datetime import datetime

API_BASE_URL = 'http://localhost:5000'
POLL_INTERVAL = 5  # seconds


class JobWorker:
    """Worker service that fetches and executes jobs"""

    def __init__(self, api_url):
        self.api_url = api_url
        self.running = True

    def fetch_pending_jobs(self):
        """
        Fetch jobs with 'pending' status from the API

        Returns:
        list: List of pending jobs
        """
        try:
            response = requests.get(f'{self.api_url}/jobs?status=pending', timeout=10)
            response.raise_for_status()
            jobs = response.json()
            if isinstance(jobs, list):
                return jobs
            return []
        except requests.RequestException as e:
            print(f"[{datetime.now()}] Error fetching pending jobs: {e}")
            return []
        except ValueError as e:
            print(f"[{datetime.now()}] Error parsing API response: {e}")
            return []

    def update_job_status(self, job_id, status, result=None):
        """
        Update job status via API

        Args:
        job_id: Job ID to update
        status: New status (processing, completed, failed)
        result: Optional result data
        """
        payload = {
            'status': status,
            'result': result
        }

        try:
            response = requests.put(
                f'{self.api_url}/jobs/{job_id}/status',
                json=payload,
                timeout=10
            )
            response.raise_for_status()
            print(f"[{datetime.now()}] Updated job {job_id} status to '{status}'")
            return response.json()
        except requests.RequestException as e:
            print(f"[{datetime.now()}] Failed to update status for job {job_id}: {e}")
            return None

    def execute_backup_job(self, payload):
        """
        Execute a backup job

        Args:
        payload: JSON string with backup parameters

        Returns:
        dict: Result with success status and message
        """
        data = json.loads(payload)
        source_path = data.get('path')
        destination_dir = data.get('destination')

        if not source_path or not destination_dir:
            raise ValueError("Backup payload must include 'path' and 'destination'")

        if not os.path.exists(source_path):
            raise FileNotFoundError(f"Source path does not exist: {source_path}")

        os.makedirs(destination_dir, exist_ok=True)

        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')

        if os.path.isdir(source_path):
            archive_name = os.path.join(destination_dir, f"backup_{timestamp}")
            result = subprocess.run(
                ['tar', '-czf', f'{archive_name}.tar.gz', '-C', os.path.dirname(source_path), os.path.basename(source_path)],
                capture_output=True,
                text=True
            )
            if result.returncode != 0:
                raise RuntimeError(f"Backup failed: {result.stderr.strip()}")

            return {
                'success': True,
                'operation': 'directory_backup',
                'source': source_path,
                'backup_file': f'{archive_name}.tar.gz',
                'message': f'Backup archive created successfully for directory {source_path}'
            }

        if os.path.isfile(source_path):
            backup_file = os.path.join(destination_dir, f"{os.path.basename(source_path)}_{timestamp}")
            shutil.copy2(source_path, backup_file)
            return {
                'success': True,
                'operation': 'file_backup',
                'source': source_path,
                'backup_file': backup_file,
                'message': f'File copied successfully from {source_path} to {backup_file}'
            }

        raise ValueError(f"Unsupported source path type: {source_path}")

    def execute_cleanup_job(self, payload):
        """
        Execute a cleanup job

        Args:
        payload: JSON string with cleanup parameters

        Returns:
        dict: Result with files cleaned and message
        """
        data = json.loads(payload)
        target_directory = data.get('directory')
        days = int(data.get('days', 0))

        if not target_directory:
            raise ValueError("Cleanup payload must include 'directory'")

        if not os.path.exists(target_directory):
            raise FileNotFoundError(f"Directory does not exist: {target_directory}")

        cutoff_time = time.time() - (days * 86400)

        files_found = []
        files_removed = []

        for root, dirs, files in os.walk(target_directory):
            for filename in files:
                file_path = os.path.join(root, filename)
                try:
                    if os.path.getmtime(file_path) < cutoff_time:
                        files_found.append(file_path)
                        os.remove(file_path)
                        files_removed.append(file_path)
                except Exception as e:
                    print(f"[{datetime.now()}] Could not process file {file_path}: {e}")

        return {
            'success': True,
            'directory': target_directory,
            'days_threshold': days,
            'files_found': len(files_found),
            'files_removed': len(files_removed),
            'removed_files': files_removed,
            'message': f'Cleanup completed in {target_directory}'
        }

    def execute_report_job(self, payload):
        """
        Execute a report generation job

        Args:
        payload: JSON string with report parameters

        Returns:
        dict: Result with report data
        """
        data = json.loads(payload)
        report_type = data.get('type', 'system')
        output_format = data.get('format', 'json')

        if report_type != 'system':
            raise ValueError(f"Unsupported report type: {report_type}")

        disk_usage = subprocess.run(['df', '-h'], capture_output=True, text=True)
        memory_usage = subprocess.run(['free', '-m'], capture_output=True, text=True)
        uptime_info = subprocess.run(['uptime'], capture_output=True, text=True)

        if disk_usage.returncode != 0:
            raise RuntimeError(f"df command failed: {disk_usage.stderr.strip()}")

        if memory_usage.returncode != 0:
            raise RuntimeError(f"free command failed: {memory_usage.stderr.strip()}")

        if uptime_info.returncode != 0:
            raise RuntimeError(f"uptime command failed: {uptime_info.stderr.strip()}")

        report = {
            'success': True,
            'report_type': report_type,
            'format': output_format,
            'generated_at': datetime.now().isoformat(),
            'disk_usage': disk_usage.stdout,
            'memory_usage': memory_usage.stdout,
            'uptime': uptime_info.stdout.strip(),
            'message': 'System report generated successfully'
        }

        return report

    def process_job(self, job):
        """
        Process a single job based on its type

        Args:
        job: Job dictionary from API
        """
        job_id = job['id']
        job_type = job['job_type']
        payload = job['payload']

        print(f"[{datetime.now()}] Processing job {job_id} ({job_type})")

        self.update_job_status(job_id, 'processing')

        try:
            if job_type == 'backup':
                result = self.execute_backup_job(payload)
            elif job_type == 'cleanup':
                result = self.execute_cleanup_job(payload)
            elif job_type == 'report':
                result = self.execute_report_job(payload)
            else:
                raise ValueError(f"Unknown job type: {job_type}")

            self.update_job_status(job_id, 'completed', json.dumps(result))
            print(f"[{datetime.now()}] Job {job_id} completed successfully")

        except Exception as e:
            error_result = {'error': str(e)}
            self.update_job_status(job_id, 'failed', json.dumps(error_result))
            print(f"[{datetime.now()}] Job {job_id} failed: {e}")

    def run(self):
        """Main worker loop"""
        print(f"Worker service started. Polling every {POLL_INTERVAL} seconds...")

        while self.running:
            try:
                jobs = self.fetch_pending_jobs()

                if jobs:
                    print(f"Found {len(jobs)} pending job(s)")
                    for job in jobs:
                        self.process_job(job)
                else:
                    print(f"[{datetime.now()}] No pending jobs")

                time.sleep(POLL_INTERVAL)

            except KeyboardInterrupt:
                print("\nShutting down worker service...")
                self.running = False
            except Exception as e:
                print(f"Error in worker loop: {e}")
                time.sleep(POLL_INTERVAL)


if __name__ == '__main__':
    worker = JobWorker(API_BASE_URL)
    worker.run()