pulse-engine 0.2.0__py3-none-any.whl

pulse_engine/extractor/base.py

@@ -0,0 +1,48 @@
+"""Base extractor ABC and supporting types for product extractors."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Any
+
+
+@dataclass
+class ExtractorConfig:
+    """Configuration for a product extractor."""
+
+    name: str
+    job_type: str
+    schedule: str | None = None
+    timeout_seconds: int = 3600
+    max_retries: int = 3
+
+
+@dataclass
+class ExtractionResult:
+    """A single extracted item from a source."""
+
+    raw_content: str
+    source_id: str
+    source_type: str
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+
+class BaseExtractor(ABC):
+    """Abstract base class that all product extractors must implement."""
+
+    @abstractmethod
+    def get_config(self) -> ExtractorConfig:
+        """Return the static configuration for this extractor."""
+
+    @abstractmethod
+    async def extract(
+        self, tenant_id: str, parameters: dict[str, Any]
+    ) -> list[ExtractionResult]:
+        """Run extraction and return a list of results."""
+
+    async def on_success(self, tenant_id: str, results: list[ExtractionResult]) -> None:
+        """Hook called after successful extraction. Override for custom behaviour."""
+
+    async def on_failure(self, tenant_id: str, error: Exception) -> None:
+        """Hook called after extraction failure. Override for custom behaviour."""

pulse_engine/extractor/models.py

@@ -0,0 +1,50 @@
+"""Job record ORM model."""
+
+import uuid
+from datetime import datetime
+from typing import Any
+
+import sqlalchemy as sa
+from sqlalchemy.orm import Mapped, mapped_column
+
+from pulse_engine.database import Base
+
+
+class JobRecordModel(Base):
+    __tablename__ = "job_records"
+
+    job_id: Mapped[str] = mapped_column(
+        sa.String, primary_key=True, default=lambda: str(uuid.uuid4())
+    )
+    job_type: Mapped[str] = mapped_column(sa.String, nullable=False)
+    product: Mapped[str] = mapped_column(sa.String, nullable=False)
+    tenant_id: Mapped[str] = mapped_column(sa.String, nullable=False, index=True)
+    status: Mapped[str] = mapped_column(sa.String, nullable=False, default="pending")
+    priority: Mapped[str] = mapped_column(sa.String, nullable=False, default="normal")
+    parameters: Mapped[dict[str, Any]] = mapped_column(
+        sa.JSON, nullable=False, default=dict
+    )
+    orchestrator_run_id: Mapped[str | None] = mapped_column(
+        sa.String, nullable=True, default=None
+    )
+    callback_url: Mapped[str | None] = mapped_column(
+        sa.String, nullable=True, default=None
+    )
+    created_at: Mapped[datetime] = mapped_column(
+        sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False
+    )
+    started_at: Mapped[datetime | None] = mapped_column(
+        sa.DateTime(timezone=True), nullable=True, default=None
+    )
+    completed_at: Mapped[datetime | None] = mapped_column(
+        sa.DateTime(timezone=True), nullable=True, default=None
+    )
+    result_summary: Mapped[dict[str, Any] | None] = mapped_column(
+        sa.JSON, nullable=True, default=None
+    )
+    error: Mapped[str | None] = mapped_column(sa.String, nullable=True, default=None)
+
+    __table_args__ = (
+        sa.Index("ix_job_records_tenant_status", "tenant_id", "status"),
+        sa.Index("ix_job_records_tenant_product", "tenant_id", "product"),
+    )

pulse_engine/extractor/orchestrator/__init__.py

@@ -0,0 +1,15 @@
+from pulse_engine.config import Settings
+from pulse_engine.extractor.orchestrator.base import BaseOrchestratorAdapter
+from pulse_engine.extractor.orchestrator.noop import NoopAdapter
+
+
+def get_orchestrator_adapter(settings: Settings) -> BaseOrchestratorAdapter:
+    backend = settings.pulse_orchestrator_backend.lower()
+    if backend == "prefect":
+        from pulse_engine.extractor.orchestrator.prefect import PrefectAdapter
+
+        return PrefectAdapter(
+            api_url=settings.prefect_api_url,
+            api_key=settings.prefect_api_key or None,
+        )
+    return NoopAdapter()

pulse_engine/extractor/orchestrator/base.py

@@ -0,0 +1,34 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass
+class OrchestratorRunStatus:
+    run_id: str
+    status: str  # pending, running, completed, failed, cancelled, unknown
+    raw_state: str | None = None
+
+
+class BaseOrchestratorAdapter(ABC):
+    """Abstraction over any external job orchestrator (Prefect, Airflow, etc.)."""
+
+    @abstractmethod
+    async def get_run_status(self, run_id: str) -> OrchestratorRunStatus:
+        """Fetch the current status of an orchestrator run by its external ID."""
+
+    @abstractmethod
+    async def cancel_run(self, run_id: str) -> bool:
+        """Request cancellation of an orchestrator run."""
+
+    @abstractmethod
+    async def health_check(self) -> bool:
+        """Return True if the orchestrator is reachable."""
+
+    @abstractmethod
+    async def create_flow_run(
+        self,
+        deployment_id: str,
+        parameters: dict[str, Any] | None = None,
+    ) -> str:
+        """Create a new flow run for the given deployment and return its run ID."""

pulse_engine/extractor/orchestrator/noop.py

@@ -0,0 +1,37 @@
+import uuid
+from typing import Any
+
+from pulse_engine.extractor.orchestrator.base import (
+    BaseOrchestratorAdapter,
+    OrchestratorRunStatus,
+)
+
+
+class NoopAdapter(BaseOrchestratorAdapter):
+    """No-op adapter used when no orchestrator is configured."""
+
+    async def get_run_status(self, run_id: str) -> OrchestratorRunStatus:
+        return OrchestratorRunStatus(run_id=run_id, status="unknown")
+
+    async def cancel_run(self, run_id: str) -> bool:
+        return False
+
+    async def health_check(self) -> bool:
+        return True
+
+    async def create_flow_run(
+        self,
+        deployment_id: str,
+        parameters: dict[str, Any] | None = None,
+    ) -> str:
+        return ""
+
+    async def create_or_update_deployment(
+        self,
+        name: str,
+        flow_entrypoint: str,
+        image: str,
+        work_pool_name: str = "products-worker-pool",
+    ) -> tuple[str, str]:
+        """Return a generated ID and the name — no actual orchestrator registration."""
+        return str(uuid.uuid4()), name

pulse_engine/extractor/orchestrator/prefect.py

@@ -0,0 +1,163 @@
+from __future__ import annotations
+
+import base64
+from typing import Any
+
+import httpx
+import structlog
+
+from pulse_engine.extractor.base import BaseExtractor
+from pulse_engine.extractor.orchestrator.base import (
+    BaseOrchestratorAdapter,
+    OrchestratorRunStatus,
+)
+
+logger = structlog.get_logger(__name__)
+
+# Map Prefect flow-run state types to canonical job statuses
+_PREFECT_STATE_MAP: dict[str, str] = {
+    "COMPLETED": "completed",
+    "FAILED": "failed",
+    "CRASHED": "failed",
+    "CANCELLED": "cancelled",
+    "CANCELLING": "cancelled",
+    "RUNNING": "running",
+    "PENDING": "pending",
+    "SCHEDULED": "pending",
+}
+
+
+class PrefectAdapter(BaseOrchestratorAdapter):
+    """Adapter for querying Prefect server API."""
+
+    def __init__(self, api_url: str, api_key: str | None = None) -> None:
+        self._api_url = api_url.rstrip("/")
+        headers: dict[str, str] = {"Content-Type": "application/json"}
+        if api_key:
+            encoded = base64.b64encode(api_key.encode()).decode()
+            headers["Authorization"] = f"Basic {encoded}"
+        self._client = httpx.AsyncClient(
+            base_url=self._api_url, headers=headers, timeout=10.0
+        )
+
+    async def get_run_status(self, run_id: str) -> OrchestratorRunStatus:
+        try:
+            resp = await self._client.get(f"/flow_runs/{run_id}")
+            resp.raise_for_status()
+            data = resp.json()
+            raw_state = data.get("state", {}).get("type", "UNKNOWN")
+            canonical = _PREFECT_STATE_MAP.get(raw_state.upper(), "unknown")
+            return OrchestratorRunStatus(
+                run_id=run_id, status=canonical, raw_state=raw_state
+            )
+        except Exception:
+            logger.warning("prefect_status_fetch_failed", run_id=run_id, exc_info=True)
+            return OrchestratorRunStatus(run_id=run_id, status="unknown")
+
+    async def cancel_run(self, run_id: str) -> bool:
+        try:
+            resp = await self._client.post(
+                f"/flow_runs/{run_id}/set_state",
+                json={"state": {"type": "CANCELLED"}},
+            )
+            return resp.is_success
+        except Exception:
+            logger.warning("prefect_cancel_failed", run_id=run_id, exc_info=True)
+            return False
+
+    async def health_check(self) -> bool:
+        try:
+            resp = await self._client.get("/health")
+            return resp.is_success
+        except Exception:
+            return False
+
+    async def register_extractors(
+        self,
+        extractors: list[type[BaseExtractor]],
+        pipeline: Any = None,
+    ) -> list[str]:
+        """Create Prefect deployments for each registered extractor.
+
+        Returns a list of deployment IDs created.
+        """
+        deployment_ids: list[str] = []
+        for ext_cls in extractors:
+            ext = ext_cls()
+            config = ext.get_config()
+            payload: dict[str, Any] = {
+                "name": config.name,
+                "flow_name": f"extract_{config.job_type}",
+                "parameters": {
+                    "extractor_cls": f"{ext_cls.__module__}.{ext_cls.__qualname__}",
+                },
+                "tags": ["pulse-engine", config.job_type],
+            }
+            if config.schedule:
+                payload["schedule"] = {"cron": config.schedule}
+
+            try:
+                resp = await self._client.post("/deployments/", json=payload)
+                resp.raise_for_status()
+                dep_id = resp.json().get("id", "")
+                deployment_ids.append(dep_id)
+                logger.info(
+                    "extractor_registered",
+                    extractor=config.name,
+                    job_type=config.job_type,
+                    deployment_id=dep_id,
+                )
+            except Exception:
+                logger.warning(
+                    "extractor_registration_failed",
+                    extractor=config.name,
+                    exc_info=True,
+                )
+
+        return deployment_ids
+
+    async def _get_or_create_flow_id(self, flow_name: str) -> str:
+        """Return the UUID of a Prefect flow, creating it if it doesn't exist."""
+        resp = await self._client.post("/flows/", json={"name": flow_name})
+        resp.raise_for_status()
+        return str(resp.json()["id"])
+
+    async def create_or_update_deployment(
+        self,
+        name: str,
+        flow_entrypoint: str,
+        image: str,
+        work_pool_name: str = "products-worker-pool",
+    ) -> tuple[str, str]:
+        """Create or update a Prefect deployment.
+
+        Returns (deployment_id, deployment_name).
+        """
+        flow_id = await self._get_or_create_flow_id(name)
+        payload: dict[str, Any] = {
+            "name": name,
+            "flow_id": flow_id,
+            "entrypoint": flow_entrypoint,
+            "work_pool_name": work_pool_name,
+            "job_variables": {"image": image},
+        }
+        resp = await self._client.post("/deployments/", json=payload)
+        resp.raise_for_status()
+        data: dict[str, Any] = resp.json()
+        return str(data["id"]), str(data.get("name", name))
+
+    async def create_flow_run(
+        self,
+        deployment_id: str,
+        parameters: dict[str, Any] | None = None,
+    ) -> str:
+        """Trigger a flow run from a deployment. Returns the flow run ID."""
+        payload: dict[str, Any] = {}
+        if parameters:
+            payload["parameters"] = parameters
+        resp = await self._client.post(
+            f"/deployments/{deployment_id}/create_flow_run",
+            json=payload,
+        )
+        resp.raise_for_status()
+        return str(resp.json()["id"])

pulse_engine/extractor/repository.py

@@ -0,0 +1,163 @@
+"""Job repository — data access layer for job records."""
+
+from __future__ import annotations
+
+import uuid
+from datetime import UTC, datetime
+from typing import Any
+
+import sqlalchemy as sa
+from sqlalchemy.engine import CursorResult
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from pulse_engine.core.exceptions import BadRequestError
+from pulse_engine.extractor.models import JobRecordModel
+
+# Explicit allowlist of columns permitted for sort_by queries
+_ALLOWED_SORT_FIELDS: set[str] = {
+    "created_at",
+    "status",
+    "product",
+    "job_type",
+    "priority",
+}
+
+
+class JobRepository:
+    def __init__(self, session: AsyncSession) -> None:
+        self._session = session
+
+    async def create(
+        self,
+        tenant_id: str,
+        job_type: str,
+        product: str,
+        priority: str = "normal",
+        parameters: dict[str, Any] | None = None,
+        orchestrator_run_id: str | None = None,
+        callback_url: str | None = None,
+    ) -> JobRecordModel:
+        record = JobRecordModel(
+            job_id=str(uuid.uuid4()),
+            tenant_id=tenant_id,
+            job_type=job_type,
+            product=product,
+            priority=priority,
+            parameters=parameters or {},
+            orchestrator_run_id=orchestrator_run_id,
+            callback_url=callback_url,
+            status="pending",
+            created_at=datetime.now(UTC),
+        )
+        self._session.add(record)
+        await self._session.commit()
+        await self._session.refresh(record)
+        return record
+
+    async def get(self, job_id: str, tenant_id: str) -> JobRecordModel | None:
+        stmt = sa.select(JobRecordModel).where(
+            JobRecordModel.job_id == job_id,
+            JobRecordModel.tenant_id == tenant_id,
+        )
+        result = await self._session.execute(stmt)
+        return result.scalar_one_or_none()
+
+    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",
+    ) -> tuple[list[JobRecordModel], int]:
+        base = sa.select(JobRecordModel).where(
+            JobRecordModel.tenant_id == tenant_id,
+        )
+        count_base = (
+            sa.select(sa.func.count())
+            .select_from(JobRecordModel)
+            .where(
+                JobRecordModel.tenant_id == tenant_id,
+            )
+        )
+
+        if status:
+            base = base.where(JobRecordModel.status == status)
+            count_base = count_base.where(JobRecordModel.status == status)
+        if product:
+            base = base.where(JobRecordModel.product == product)
+            count_base = count_base.where(JobRecordModel.product == product)
+        if job_type:
+            base = base.where(JobRecordModel.job_type == job_type)
+            count_base = count_base.where(JobRecordModel.job_type == job_type)
+
+        # Sorting — validate against allowlist to prevent column-name injection
+        if sort_by not in _ALLOWED_SORT_FIELDS:
+            raise BadRequestError(
+                f"Invalid sort field: {sort_by!r}. "
+                f"Allowed: {', '.join(sorted(_ALLOWED_SORT_FIELDS))}"
+            )
+        sort_col = getattr(JobRecordModel, sort_by)
+        if order == "asc":
+            base = base.order_by(sa.asc(sort_col))
+        else:
+            base = base.order_by(sa.desc(sort_col))
+
+        base = base.offset(offset).limit(limit)
+
+        result = await self._session.execute(base)
+        jobs = list(result.scalars().all())
+
+        count_result = await self._session.execute(count_base)
+        total = count_result.scalar_one()
+
+        return jobs, total
+
+    async def update_status(
+        self,
+        job_id: str,
+        tenant_id: str,
+        status: str,
+        **fields: object,
+    ) -> JobRecordModel | None:
+        record = await self.get(job_id, tenant_id)
+        if record is None:
+            return None
+
+        record.status = status
+        for key, value in fields.items():
+            if hasattr(record, key):
+                setattr(record, key, value)
+
+        await self._session.commit()
+        await self._session.refresh(record)
+        return record
+
+    async def delete(self, job_id: str, tenant_id: str) -> bool:
+        stmt = sa.delete(JobRecordModel).where(
+            JobRecordModel.job_id == job_id,
+            JobRecordModel.tenant_id == tenant_id,
+        )
+        result: CursorResult[Any] = await self._session.execute(stmt)  # type: ignore[assignment]
+        await self._session.commit()
+        return bool(result.rowcount > 0)
+
+    async def count_active(self, tenant_id: str) -> int:
+        stmt = (
+            sa.select(sa.func.count())
+            .select_from(JobRecordModel)
+            .where(
+                JobRecordModel.tenant_id == tenant_id,
+                JobRecordModel.status.in_(["pending", "running"]),
+            )
+        )
+        result = await self._session.execute(stmt)
+        return result.scalar_one()
+
+    async def get_by_id(self, job_id: str) -> JobRecordModel | None:
+        stmt = sa.select(JobRecordModel).where(JobRecordModel.job_id == job_id)
+        result = await self._session.execute(stmt)
+        return result.scalar_one_or_none()

pulse_engine/extractor/router.py

@@ -0,0 +1,102 @@
+"""Jobs API router."""
+
+from typing import Literal
+
+from fastapi import APIRouter, Depends, Query
+from starlette.responses import JSONResponse
+
+from pulse_engine.core.scope import require_scope
+from pulse_engine.dependencies import get_job_service
+from pulse_engine.extractor.schemas import (
+    CreateJobRequest,
+    CreateJobResponse,
+    JobListResponse,
+    JobResponse,
+    StatusUpdateRequest,
+    StatusUpdateResponse,
+)
+from pulse_engine.extractor.service import JobService
+from pulse_engine.middleware.tenant import get_tenant_id
+
+router = APIRouter(prefix="/jobs", tags=["Jobs"])
+
+
+@router.post(
+    "/",
+    response_model=CreateJobResponse,
+    status_code=202,
+    dependencies=[require_scope("jobs:trigger_next")],
+)
+async def register_job(
+    body: CreateJobRequest,
+    tenant_id: str = Depends(get_tenant_id),
+    service: JobService = Depends(get_job_service),
+) -> CreateJobResponse:
+    return await service.register_job(tenant_id, body)
+
+
+@router.get("/{job_id}", response_model=JobResponse, status_code=200)
+async def get_job(
+    job_id: str,
+    tenant_id: str = Depends(get_tenant_id),
+    service: JobService = Depends(get_job_service),
+) -> JobResponse:
+    return await service.get_job(tenant_id, job_id)
+
+
+@router.get("/", response_model=JobListResponse, status_code=200)
+async def list_jobs(
+    tenant_id: str = Depends(get_tenant_id),
+    service: JobService = Depends(get_job_service),
+    status: str | None = Query(None),
+    product: str | None = Query(None),
+    job_type: str | None = Query(None),
+    limit: int = Query(20, ge=1, le=100),
+    offset: int = Query(0, ge=0),
+    sort_by: str = Query("created_at"),
+    order: Literal["asc", "desc"] = Query("desc"),
+) -> JobListResponse:
+    return await service.list_jobs(
+        tenant_id,
+        status=status,
+        product=product,
+        job_type=job_type,
+        limit=limit,
+        offset=offset,
+        sort_by=sort_by,
+        order=order,
+    )
+
+
+@router.post(
+    "/{job_id}/status",
+    response_model=StatusUpdateResponse,
+    status_code=200,
+    dependencies=[require_scope("jobs:status")],
+)
+async def push_status(
+    job_id: str,
+    body: StatusUpdateRequest,
+    tenant_id: str = Depends(get_tenant_id),
+    service: JobService = Depends(get_job_service),
+) -> StatusUpdateResponse:
+    return await service.push_status(tenant_id, job_id, body)
+
+
+@router.post("/{job_id}/cancel", response_model=JobResponse, status_code=200)
+async def cancel_job(
+    job_id: str,
+    tenant_id: str = Depends(get_tenant_id),
+    service: JobService = Depends(get_job_service),
+) -> JobResponse:
+    return await service.cancel_job(tenant_id, job_id)
+
+
+@router.delete("/{job_id}", status_code=204)
+async def delete_job(
+    job_id: str,
+    tenant_id: str = Depends(get_tenant_id),
+    service: JobService = Depends(get_job_service),
+) -> JSONResponse:
+    await service.delete_job(tenant_id, job_id)
+    return JSONResponse(status_code=204, content=None)