pulse-engine 0.2.0__py3-none-any.whl

pulse_engine/extractor/schemas.py

@@ -0,0 +1,93 @@
+"""Pydantic request/response schemas for the Jobs API."""
+
+from datetime import datetime
+from enum import Enum
+from typing import Any, Literal
+
+from pydantic import BaseModel, Field
+
+
+class JobStatus(str, Enum):
+    pending = "pending"
+    running = "running"
+    completed = "completed"
+    failed = "failed"
+    cancelled = "cancelled"
+
+
+class JobPriority(str, Enum):
+    low = "low"
+    normal = "normal"
+    high = "high"
+
+
+class CreateJobRequest(BaseModel):
+    product: str
+    stage: Literal["extractor", "processor", "storage"] = "extractor"
+    orchestrator: str = "prefect"
+    compute: str = "ecs"
+    chain: bool = True
+    job_id: str | None = None  # reuse existing job for isolation runs
+    config: dict[str, Any] = Field(default_factory=dict)
+    # Retained for backwards compatibility:
+    job_type: str | None = None
+    priority: JobPriority = JobPriority.normal
+    orchestrator_run_id: str | None = None
+    parameters: dict[str, Any] = Field(default_factory=dict)
+    callback_url: str | None = None
+
+
+class CreateJobResponse(BaseModel):
+    job_id: str
+    job_type: str
+    product: str
+    stage: str
+    status: str
+    created_at: datetime
+
+
+class StatusUpdateRequest(BaseModel):
+    status: JobStatus
+    stage: str | None = None
+    result_summary: dict[str, Any] | None = None
+    error: str | None = None
+
+
+class StatusUpdateResponse(BaseModel):
+    job_id: str
+    status: str
+    updated_at: datetime
+
+
+class JobStageResponse(BaseModel):
+    stage: str
+    status: str
+    prefect_flow_run_id: str | None = None
+    started_at: datetime | None = None
+    completed_at: datetime | None = None
+    error: str | None = None
+
+
+class JobResponse(BaseModel):
+    job_id: str
+    job_type: str
+    product: str
+    tenant_id: str
+    status: str
+    priority: str
+    parameters: dict[str, Any]
+    orchestrator_run_id: str | None = None
+    callback_url: str | None = None
+    created_at: datetime
+    started_at: datetime | None = None
+    completed_at: datetime | None = None
+    result_summary: dict[str, Any] | None = None
+    error: str | None = None
+    stages: list[JobStageResponse] = Field(default_factory=list)
+
+
+class JobListResponse(BaseModel):
+    total: int
+    jobs: list[JobResponse]
+    limit: int
+    offset: int

pulse_engine/extractor/service.py

@@ -0,0 +1,431 @@
+"""Job service — business logic for job management."""
+
+from __future__ import annotations
+
+from datetime import UTC, datetime
+from typing import TYPE_CHECKING, Any
+
+import httpx
+import structlog
+
+from pulse_engine.config import Settings
+from pulse_engine.core.exceptions import NotFoundError, TooManyRequestsError
+from pulse_engine.extractor.schemas import (
+    CreateJobRequest,
+    CreateJobResponse,
+    JobListResponse,
+    JobResponse,
+    JobStageResponse,
+    JobStatus,
+    StatusUpdateRequest,
+    StatusUpdateResponse,
+)
+
+if TYPE_CHECKING:
+    from pulse_engine.core.job_token import JobTokenIssuer
+    from pulse_engine.deployment.repository import RegistrationRepository
+    from pulse_engine.extractor.orchestrator.base import BaseOrchestratorAdapter
+    from pulse_engine.extractor.repository import JobRepository
+    from pulse_engine.extractor.stage_repository import StageRepository
+
+logger = structlog.get_logger(__name__)
+
+TERMINAL_STATES = {JobStatus.completed, JobStatus.failed, JobStatus.cancelled}
+
+_STAGE_MAP = {
+    "extractor": "extraction",
+    "processor": "processing",
+    "storage": "storage",
+}
+_STAGE_ORDER = ["extraction", "processing", "storage"]
+
+
+class JobService:
+    def __init__(
+        self,
+        repository: JobRepository,
+        orchestrator: BaseOrchestratorAdapter,
+        settings: Settings,
+        stage_repository: StageRepository | None = None,
+        deployment_repository: RegistrationRepository | None = None,
+        token_issuer: JobTokenIssuer | None = None,
+        job_launcher: Any | None = None,  # JobLauncher | None
+    ) -> None:
+        self._repository = repository
+        self._orchestrator = orchestrator
+        self._settings = settings
+        self._stage_repo = stage_repository
+        self._deploy_repo = deployment_repository  # kept for legacy path
+        self._token_issuer = token_issuer  # kept for legacy path
+        self._job_launcher = job_launcher  # preferred path
+
+    async def register_job(
+        self, tenant_id: str, request: CreateJobRequest
+    ) -> CreateJobResponse:
+        active = await self._repository.count_active(tenant_id)
+        if active >= self._settings.pulse_max_concurrent_jobs_per_tenant:
+            raise TooManyRequestsError(
+                "Concurrent job limit reached for this tenant",
+                tenant_id=tenant_id,
+                limit=self._settings.pulse_max_concurrent_jobs_per_tenant,
+            )
+
+        job_type = request.job_type or _STAGE_MAP.get(request.stage, request.stage)
+
+        if request.job_id:
+            record = await self._repository.get(request.job_id, tenant_id)
+            if record is None:
+                raise NotFoundError("Job not found", job_id=request.job_id)
+        else:
+            params = {
+                **request.parameters,
+                **request.config,
+                "chain": request.chain,
+            }
+            record = await self._repository.create(
+                tenant_id=tenant_id,
+                job_type=job_type,
+                product=request.product,
+                priority=request.priority.value,
+                parameters=params,
+                orchestrator_run_id=request.orchestrator_run_id,
+                callback_url=request.callback_url,
+            )
+
+        stage_name = _STAGE_MAP.get(request.stage, request.stage)
+        trigger_condition = request.stage is not None and (
+            self._job_launcher is not None
+            or (self._stage_repo and self._deploy_repo and self._token_issuer)
+        )
+        if trigger_condition:
+            await self._trigger_stage(
+                job_id=record.job_id,
+                tenant_id=tenant_id,
+                product=request.product,
+                stage=request.stage,
+                stage_name=stage_name,
+                chain=request.chain,
+                config=request.config,
+                orchestrator=request.orchestrator,
+                compute=request.compute,
+            )
+
+        if record.status == "pending":
+            await self._repository.update_status(
+                record.job_id,
+                tenant_id,
+                "running",
+                started_at=datetime.now(UTC),
+            )
+
+        return CreateJobResponse(
+            job_id=record.job_id,
+            job_type=record.job_type,
+            product=record.product,
+            stage=stage_name,
+            status="running",
+            created_at=record.created_at,
+        )
+
+    async def _trigger_stage(
+        self,
+        job_id: str,
+        tenant_id: str,
+        product: str,
+        stage: str,
+        stage_name: str,
+        chain: bool,
+        config: dict[str, Any],
+        orchestrator: str = "prefect",
+        compute: str = "ecs",
+    ) -> None:
+        if self._job_launcher is not None:
+            from pulse_engine.deployment.job_launcher import LaunchResult
+
+            result: LaunchResult = await self._job_launcher.launch(
+                job_id=job_id,
+                tenant_id=tenant_id,
+                product=product,
+                stage=stage,
+                orchestrator=orchestrator,
+                compute=compute,
+                chain=chain,
+                config=config,
+            )
+            if self._stage_repo:
+                await self._stage_repo.create(
+                    job_id=job_id,
+                    stage=stage_name,
+                    prefect_flow_run_id=result.flow_run_id,
+                )
+            logger.info(
+                "stage_triggered",
+                job_id=job_id,
+                stage=stage_name,
+                flow_run_id=result.flow_run_id,
+            )
+            return
+
+        # Legacy path: direct orchestrator call (backwards compat when launcher absent)
+        if self._deploy_repo is None or self._token_issuer is None:
+            return
+        deployment = await self._deploy_repo.get(
+            product=product,
+            stage=stage,
+        )
+        if deployment is None:
+            raise NotFoundError(
+                f"No deployment registered for {product}/{stage}",
+                product=product,
+                stage=stage,
+            )
+
+        scope = ["jobs:status"]
+        if stage == "storage":
+            scope.append("kb:write")
+        if chain:
+            scope.append("jobs:trigger_next")
+
+        token = self._token_issuer.issue(
+            job_id=job_id,
+            tenant_id=tenant_id,
+            product=product,
+            stage=stage_name,
+            scope=scope,
+        )
+
+        flow_run_id = await self._orchestrator.create_flow_run(
+            deployment_id=deployment.prefect_deployment_id,  # type: ignore[attr-defined]
+            parameters={
+                "job_id": job_id,
+                "chain": chain,
+                "config": config,
+                "pulse_api_token": token,
+                "pulse_engine_url": self._settings.pulse_engine_url,
+            },
+        )
+
+        if self._stage_repo:
+            await self._stage_repo.create(
+                job_id=job_id,
+                stage=stage_name,
+                prefect_flow_run_id=flow_run_id,
+            )
+
+        logger.info(
+            "stage_triggered",
+            job_id=job_id,
+            stage=stage_name,
+            flow_run_id=flow_run_id,
+        )
+
+    async def get_job(self, tenant_id: str, job_id: str) -> JobResponse:
+        record = await self._repository.get(job_id, tenant_id)
+        if record is None:
+            raise NotFoundError("Job not found", job_id=job_id)
+
+        if record.orchestrator_run_id and record.status in (
+            "pending",
+            "running",
+        ):
+            try:
+                run_status = await self._orchestrator.get_run_status(
+                    record.orchestrator_run_id
+                )
+                if (
+                    run_status.status != "unknown"
+                    and run_status.status != record.status
+                ):
+                    record = await self._repository.update_status(
+                        job_id, tenant_id, run_status.status
+                    )
+            except Exception:
+                logger.warning(
+                    "orchestrator_sync_failed",
+                    job_id=job_id,
+                    exc_info=True,
+                )
+
+        response = self._to_response(record)
+
+        if self._stage_repo:
+            stages = await self._stage_repo.get_stages_for_job(job_id)
+            response.stages = [
+                JobStageResponse(
+                    stage=s.stage,
+                    status=s.status,
+                    prefect_flow_run_id=s.prefect_flow_run_id,
+                    started_at=s.started_at,
+                    completed_at=s.completed_at,
+                    error=s.error,
+                )
+                for s in stages
+            ]
+
+        return response
+
+    async def list_jobs(
+        self,
+        tenant_id: str,
+        status: str | None = None,
+        product: str | None = None,
+        job_type: str | None = None,
+        limit: int = 20,
+        offset: int = 0,
+        sort_by: str = "created_at",
+        order: str = "desc",
+    ) -> JobListResponse:
+        jobs, total = await self._repository.list_jobs(
+            tenant_id=tenant_id,
+            status=status,
+            product=product,
+            job_type=job_type,
+            limit=limit,
+            offset=offset,
+            sort_by=sort_by,
+            order=order,
+        )
+        return JobListResponse(
+            total=total,
+            jobs=[self._to_response(j) for j in jobs],
+            limit=limit,
+            offset=offset,
+        )
+
+    async def push_status(
+        self, tenant_id: str, job_id: str, request: StatusUpdateRequest
+    ) -> StatusUpdateResponse:
+        now = datetime.now(UTC)
+        extra_fields: dict[str, object] = {}
+
+        if request.status == JobStatus.running:
+            extra_fields["started_at"] = now
+        if request.status in TERMINAL_STATES:
+            extra_fields["completed_at"] = now
+        if request.result_summary is not None:
+            extra_fields["result_summary"] = request.result_summary
+        if request.error is not None:
+            extra_fields["error"] = request.error
+
+        if request.stage and self._stage_repo:
+            stages = await self._stage_repo.get_stages_for_job(job_id)
+            for s in stages:
+                if s.stage == request.stage:
+                    await self._stage_repo.update_status(
+                        s.id,
+                        request.status.value,
+                        error=request.error,
+                    )
+                    break
+
+            if request.status in TERMINAL_STATES:
+                if request.status == JobStatus.completed and request.stage == "storage":
+                    extra_fields["status"] = "completed"
+                elif request.status == JobStatus.failed:
+                    extra_fields["status"] = "failed"
+
+        record = await self._repository.update_status(
+            job_id, tenant_id, request.status.value, **extra_fields
+        )
+        if record is None:
+            raise NotFoundError("Job not found", job_id=job_id)
+
+        if request.status in TERMINAL_STATES and record.callback_url:
+            await self._fire_callback(
+                record.callback_url,
+                {
+                    "job_id": record.job_id,
+                    "status": record.status,
+                    "result_summary": record.result_summary,
+                    "error": record.error,
+                },
+            )
+
+        return StatusUpdateResponse(
+            job_id=record.job_id,
+            status=record.status,
+            updated_at=now,
+        )
+
+    async def cancel_job(self, tenant_id: str, job_id: str) -> JobResponse:
+        record = await self._repository.get(job_id, tenant_id)
+        if record is None:
+            raise NotFoundError("Job not found", job_id=job_id)
+
+        if record.orchestrator_run_id:
+            try:
+                await self._orchestrator.cancel_run(
+                    record.orchestrator_run_id,
+                )
+            except Exception:
+                logger.warning(
+                    "orchestrator_cancel_failed",
+                    job_id=job_id,
+                    exc_info=True,
+                )
+
+        now = datetime.now(UTC)
+        record = await self._repository.update_status(
+            job_id, tenant_id, "cancelled", completed_at=now
+        )
+
+        if record and record.callback_url:
+            await self._fire_callback(
+                record.callback_url,
+                {
+                    "job_id": record.job_id,
+                    "status": record.status,
+                    "result_summary": record.result_summary,
+                    "error": record.error,
+                },
+            )
+
+        return self._to_response(record)
+
+    async def delete_job(self, tenant_id: str, job_id: str) -> bool:
+        deleted = await self._repository.delete(job_id, tenant_id)
+        if not deleted:
+            raise NotFoundError("Job not found", job_id=job_id)
+        return True
+
+    async def _fire_callback(self, callback_url: str, payload: dict[str, Any]) -> None:
+        try:
+            async with httpx.AsyncClient(
+                timeout=self._settings.pulse_job_callback_timeout
+            ) as client:
+                resp = await client.post(callback_url, json=payload)
+                logger.info(
+                    "callback_fired",
+                    url=callback_url,
+                    status_code=resp.status_code,
+                    job_id=payload.get("job_id"),
+                )
+        except Exception:
+            logger.warning(
+                "callback_failed",
+                url=callback_url,
+                job_id=payload.get("job_id"),
+                exc_info=True,
+            )
+
+    @staticmethod
+    def _to_response(record: object) -> JobResponse:
+        from pulse_engine.extractor.models import JobRecordModel
+
+        r: JobRecordModel = record  # type: ignore[assignment]
+        return JobResponse(
+            job_id=r.job_id,
+            job_type=r.job_type,
+            product=r.product,
+            tenant_id=r.tenant_id,
+            status=r.status,
+            priority=r.priority,
+            parameters=r.parameters,
+            orchestrator_run_id=r.orchestrator_run_id,
+            callback_url=r.callback_url,
+            created_at=r.created_at,
+            started_at=r.started_at,
+            completed_at=r.completed_at,
+            result_summary=r.result_summary,
+            error=r.error,
+        )

pulse_engine/extractor/stage_models.py

@@ -0,0 +1,36 @@
+"""Per-stage tracking model for pipeline jobs."""
+
+import uuid
+from datetime import datetime
+
+import sqlalchemy as sa
+from sqlalchemy.orm import Mapped, mapped_column
+
+from pulse_engine.database import Base
+
+
+class JobStageModel(Base):
+    __tablename__ = "job_stages"
+
+    id: Mapped[str] = mapped_column(
+        sa.String, primary_key=True, default=lambda: str(uuid.uuid4())
+    )
+    job_id: Mapped[str] = mapped_column(
+        sa.String, sa.ForeignKey("job_records.job_id"), nullable=False, index=True
+    )
+    stage: Mapped[str] = mapped_column(sa.String(50), nullable=False)
+    status: Mapped[str] = mapped_column(
+        sa.String(50), nullable=False, default="pending"
+    )
+    prefect_flow_run_id: Mapped[str | None] = mapped_column(
+        sa.String(255), nullable=True, default=None
+    )
+    started_at: Mapped[datetime | None] = mapped_column(
+        sa.DateTime, nullable=True, default=None
+    )
+    completed_at: Mapped[datetime | None] = mapped_column(
+        sa.DateTime, nullable=True, default=None
+    )
+    error: Mapped[str | None] = mapped_column(sa.Text, nullable=True, default=None)
+
+    __table_args__ = (sa.Index("ix_job_stages_job_id_stage", "job_id", "stage"),)

pulse_engine/extractor/stage_repository.py

@@ -0,0 +1,109 @@
+"""Data access layer for job stage records."""
+
+from __future__ import annotations
+
+import uuid
+from datetime import UTC, datetime
+
+import sqlalchemy as sa
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from pulse_engine.extractor.stage_models import JobStageModel
+
+
+class StageRepository:
+    def __init__(self, session: AsyncSession) -> None:
+        self._session = session
+
+    async def create(
+        self,
+        job_id: str,
+        stage: str,
+        prefect_flow_run_id: str | None = None,
+    ) -> JobStageModel:
+        record = JobStageModel(
+            id=str(uuid.uuid4()),
+            job_id=job_id,
+            stage=stage,
+            status="pending",
+            prefect_flow_run_id=prefect_flow_run_id,
+        )
+        self._session.add(record)
+        await self._session.commit()
+        await self._session.refresh(record)
+        return record
+
+    async def update_status(
+        self,
+        stage_id: str,
+        status: str,
+        error: str | None = None,
+    ) -> JobStageModel:
+        stmt = sa.select(JobStageModel).where(JobStageModel.id == stage_id)
+        result = await self._session.execute(stmt)
+        record = result.scalar_one()
+
+        now = datetime.now(UTC)
+        record.status = status
+        if status == "running" and record.started_at is None:
+            record.started_at = now
+        if status in ("completed", "failed", "cancelled"):
+            record.completed_at = now
+        if error is not None:
+            record.error = error
+
+        await self._session.commit()
+        await self._session.refresh(record)
+        return record
+
+    async def get_stages_for_job(self, job_id: str) -> list[JobStageModel]:
+        stmt = (
+            sa.select(JobStageModel)
+            .where(JobStageModel.job_id == job_id)
+            .order_by(JobStageModel.started_at.asc().nulls_last())
+        )
+        result = await self._session.execute(stmt)
+        return list(result.scalars().all())
+
+    async def get_by_flow_run_id(self, flow_run_id: str) -> JobStageModel | None:
+        stmt = sa.select(JobStageModel).where(
+            JobStageModel.prefect_flow_run_id == flow_run_id
+        )
+        result = await self._session.execute(stmt)
+        return result.scalar_one_or_none()
+
+    async def get_running_stages(self) -> list[JobStageModel]:
+        stmt = sa.select(JobStageModel).where(
+            JobStageModel.status == "running",
+            JobStageModel.prefect_flow_run_id.isnot(None),
+        )
+        result = await self._session.execute(stmt)
+        return list(result.scalars().all())
+
+    async def get_completed_unchained_stages(self) -> list[JobStageModel]:
+        """Get completed extraction/processing stages where the next stage
+        has NOT been triggered yet (no row exists for the downstream stage)."""
+        from sqlalchemy.orm import aliased
+
+        NextStage = aliased(JobStageModel)
+        _NEXT = {"extraction": "processing", "processing": "storage"}
+
+        results: list[JobStageModel] = []
+        for current, downstream in _NEXT.items():
+            next_exists = (
+                sa.select(sa.literal(1))
+                .where(
+                    NextStage.job_id == JobStageModel.job_id,
+                    NextStage.stage == downstream,
+                )
+                .correlate(JobStageModel)
+                .exists()
+            )
+            stmt = sa.select(JobStageModel).where(
+                JobStageModel.status == "completed",
+                JobStageModel.stage == current,
+                ~next_exists,
+            )
+            result = await self._session.execute(stmt)
+            results.extend(result.scalars().all())
+        return results