site11/services/news-engine-console/backend/app/services/pipeline_service.py

from datetime import datetime
from typing import List, Optional, Dict, Any
from bson import ObjectId
from motor.motor_asyncio import AsyncIOMotorDatabase

from app.models.pipeline import Pipeline, PipelineStats
from app.schemas.pipeline import PipelineCreate, PipelineUpdate, PipelineLog


class PipelineService:
    """Service for managing pipelines"""

    def __init__(self, db: AsyncIOMotorDatabase):
        self.db = db
        self.collection = db.pipelines
        self.logs_collection = db.pipeline_logs

    async def get_pipelines(
        self,
        type: Optional[str] = None,
        status: Optional[str] = None
    ) -> List[Pipeline]:
        """
        Get all pipelines with optional filtering

        Args:
            type: Filter by pipeline type
            status: Filter by status (running/stopped/error)

        Returns:
            List of pipelines
        """
        query = {}
        if type:
            query["type"] = type
        if status:
            query["status"] = status

        cursor = self.collection.find(query).sort("created_at", -1)

        pipelines = []
        async for doc in cursor:
            pipelines.append(Pipeline(**doc))

        return pipelines

    async def get_pipeline_by_id(self, pipeline_id: str) -> Optional[Pipeline]:
        """Get a pipeline by ID"""
        if not ObjectId.is_valid(pipeline_id):
            return None

        doc = await self.collection.find_one({"_id": ObjectId(pipeline_id)})
        if doc:
            return Pipeline(**doc)
        return None

    async def create_pipeline(self, pipeline_data: PipelineCreate) -> Pipeline:
        """Create a new pipeline"""
        pipeline_dict = pipeline_data.model_dump()
        pipeline_dict.update({
            "status": "stopped",
            "stats": PipelineStats().model_dump(),
            "created_at": datetime.utcnow(),
            "updated_at": datetime.utcnow()
        })

        result = await self.collection.insert_one(pipeline_dict)
        pipeline_dict["_id"] = result.inserted_id

        return Pipeline(**pipeline_dict)

    async def update_pipeline(
        self,
        pipeline_id: str,
        update_data: PipelineUpdate
    ) -> Optional[Pipeline]:
        """Update a pipeline"""
        if not ObjectId.is_valid(pipeline_id):
            return None

        update_dict = {
            k: v for k, v in update_data.model_dump().items()
            if v is not None
        }

        if not update_dict:
            return await self.get_pipeline_by_id(pipeline_id)

        update_dict["updated_at"] = datetime.utcnow()

        result = await self.collection.find_one_and_update(
            {"_id": ObjectId(pipeline_id)},
            {"$set": update_dict},
            return_document=True
        )

        if result:
            return Pipeline(**result)
        return None

    async def delete_pipeline(self, pipeline_id: str) -> bool:
        """Delete a pipeline"""
        if not ObjectId.is_valid(pipeline_id):
            return False

        result = await self.collection.delete_one({"_id": ObjectId(pipeline_id)})
        return result.deleted_count > 0

    async def start_pipeline(self, pipeline_id: str) -> Optional[Pipeline]:
        """Start a pipeline"""
        if not ObjectId.is_valid(pipeline_id):
            return None

        pipeline = await self.get_pipeline_by_id(pipeline_id)
        if not pipeline:
            return None

        if pipeline.status == "running":
            return pipeline  # Already running

        # TODO: Actual pipeline start logic would go here
        # For now, just update the status

        result = await self.collection.find_one_and_update(
            {"_id": ObjectId(pipeline_id)},
            {
                "$set": {
                    "status": "running",
                    "last_run": datetime.utcnow(),
                    "updated_at": datetime.utcnow()
                }
            },
            return_document=True
        )

        if result:
            # Log the start event
            await self._add_log(
                pipeline_id,
                "INFO",
                f"Pipeline {pipeline.name} started"
            )
            return Pipeline(**result)
        return None

    async def stop_pipeline(self, pipeline_id: str) -> Optional[Pipeline]:
        """Stop a pipeline"""
        if not ObjectId.is_valid(pipeline_id):
            return None

        pipeline = await self.get_pipeline_by_id(pipeline_id)
        if not pipeline:
            return None

        if pipeline.status == "stopped":
            return pipeline  # Already stopped

        # TODO: Actual pipeline stop logic would go here

        result = await self.collection.find_one_and_update(
            {"_id": ObjectId(pipeline_id)},
            {
                "$set": {
                    "status": "stopped",
                    "updated_at": datetime.utcnow()
                }
            },
            return_document=True
        )

        if result:
            # Log the stop event
            await self._add_log(
                pipeline_id,
                "INFO",
                f"Pipeline {pipeline.name} stopped"
            )
            return Pipeline(**result)
        return None

    async def restart_pipeline(self, pipeline_id: str) -> Optional[Pipeline]:
        """Restart a pipeline"""
        # Stop first
        await self.stop_pipeline(pipeline_id)
        # Then start
        return await self.start_pipeline(pipeline_id)

    async def get_pipeline_stats(self, pipeline_id: str) -> Optional[PipelineStats]:
        """Get statistics for a pipeline"""
        pipeline = await self.get_pipeline_by_id(pipeline_id)
        if not pipeline:
            return None

        return pipeline.stats

    async def update_pipeline_stats(
        self,
        pipeline_id: str,
        success: bool = True,
        duration_seconds: Optional[float] = None
    ) -> Optional[Pipeline]:
        """
        Update pipeline statistics after a run

        Args:
            pipeline_id: Pipeline ID
            success: Whether the run was successful
            duration_seconds: Duration of the run
        """
        if not ObjectId.is_valid(pipeline_id):
            return None

        pipeline = await self.get_pipeline_by_id(pipeline_id)
        if not pipeline:
            return None

        # Update stats
        new_total = pipeline.stats.total_processed + 1
        new_success = pipeline.stats.success_count + (1 if success else 0)
        new_error = pipeline.stats.error_count + (0 if success else 1)

        # Calculate average duration
        if duration_seconds:
            current_avg = pipeline.stats.average_duration_seconds or 0
            current_count = pipeline.stats.total_processed
            new_avg = ((current_avg * current_count) + duration_seconds) / new_total
        else:
            new_avg = pipeline.stats.average_duration_seconds

        result = await self.collection.find_one_and_update(
            {"_id": ObjectId(pipeline_id)},
            {
                "$set": {
                    "stats.total_processed": new_total,
                    "stats.success_count": new_success,
                    "stats.error_count": new_error,
                    "stats.last_run": datetime.utcnow(),
                    "stats.average_duration_seconds": new_avg,
                    "last_run": datetime.utcnow(),
                    "updated_at": datetime.utcnow()
                }
            },
            return_document=True
        )

        if result:
            return Pipeline(**result)
        return None

    async def get_pipeline_logs(
        self,
        pipeline_id: str,
        limit: int = 100,
        level: Optional[str] = None
    ) -> List[PipelineLog]:
        """
        Get logs for a pipeline

        Args:
            pipeline_id: Pipeline ID
            limit: Maximum number of logs to return
            level: Filter by log level (INFO, WARNING, ERROR)

        Returns:
            List of pipeline logs
        """
        if not ObjectId.is_valid(pipeline_id):
            return []

        query = {"pipeline_id": pipeline_id}
        if level:
            query["level"] = level

        cursor = self.logs_collection.find(query).sort("timestamp", -1).limit(limit)

        logs = []
        async for doc in cursor:
            logs.append(PipelineLog(
                timestamp=doc["timestamp"],
                level=doc["level"],
                message=doc["message"],
                details=doc.get("details")
            ))

        return logs

    async def _add_log(
        self,
        pipeline_id: str,
        level: str,
        message: str,
        details: Optional[Dict[str, Any]] = None
    ):
        """Add a log entry for a pipeline"""
        log_entry = {
            "pipeline_id": pipeline_id,
            "timestamp": datetime.utcnow(),
            "level": level,
            "message": message,
            "details": details or {}
        }

        await self.logs_collection.insert_one(log_entry)

    async def update_pipeline_config(
        self,
        pipeline_id: str,
        config: Dict[str, Any]
    ) -> Optional[Pipeline]:
        """Update pipeline configuration"""
        if not ObjectId.is_valid(pipeline_id):
            return None

        result = await self.collection.find_one_and_update(
            {"_id": ObjectId(pipeline_id)},
            {
                "$set": {
                    "config": config,
                    "updated_at": datetime.utcnow()
                }
            },
            return_document=True
        )

        if result:
            await self._add_log(
                pipeline_id,
                "INFO",
                "Pipeline configuration updated"
            )
            return Pipeline(**result)
        return None