pulse-engine 0.2.0__py3-none-any.whl

pulse_engine/core/job_token.py

@@ -0,0 +1,109 @@
+# src/pulse_engine/core/job_token.py
+"""Job-scoped JWT issuance and verification for container callbacks."""
+
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass, field
+from typing import Any
+
+from jose import JWTError, jwt
+
+from pulse_engine.core.exceptions import UnauthorizedError
+
+_ALGORITHM = "HS256"
+_DEFAULT_TTL = 86400  # 24 hours
+
+
+@dataclass
+class JobClaims:
+    """Claims extracted from a job-scoped JWT."""
+
+    sub: str  # "job:{job_id}"
+    tenant_id: str
+    product: str
+    stage: str  # "extraction" | "processing" | "storage"
+    scope: list[str]
+    orchestrator: str = "prefect"
+    compute: str = "ecs"
+    raw: dict[str, Any] = field(default_factory=dict)
+
+
+class JobTokenIssuer:
+    """Issues HMAC-signed JWTs scoped to a single job stage."""
+
+    def __init__(self, secret: str) -> None:
+        self._secret = secret
+
+    def issue(
+        self,
+        job_id: str,
+        tenant_id: str,
+        product: str,
+        stage: str,
+        scope: list[str],
+        orchestrator: str = "prefect",
+        compute: str = "ecs",
+        ttl_seconds: int = _DEFAULT_TTL,
+    ) -> str:
+        now = int(time.time())
+        payload: dict[str, Any] = {
+            "sub": f"job:{job_id}",
+            "custom:tenant_id": tenant_id,
+            "product": product,
+            "stage": stage,
+            "scope": scope,
+            "orchestrator": orchestrator,
+            "compute": compute,
+            "iat": now,
+            "exp": now + ttl_seconds,
+        }
+        return jwt.encode(payload, self._secret, algorithm=_ALGORITHM)
+
+    def issue_token(
+        self,
+        pipeline_run_id: str,
+        tenant_id: str,
+        ttl_seconds: int = _DEFAULT_TTL,
+    ) -> str:
+        """Issue a run-scoped token for pipeline orchestration."""
+        return self.issue(
+            job_id=pipeline_run_id,
+            tenant_id=tenant_id,
+            product="pipeline",
+            stage="pipeline",
+            scope=["pipeline:run", "kb:write"],
+            ttl_seconds=ttl_seconds,
+        )
+
+
+class JobTokenVerifier:
+    """Verifies HMAC-signed job-scoped JWTs."""
+
+    def __init__(self, secret: str) -> None:
+        self._secret = secret
+
+    async def verify(self, token: str) -> JobClaims:
+        try:
+            payload = jwt.decode(token, self._secret, algorithms=[_ALGORITHM])
+        except JWTError as e:
+            raise UnauthorizedError(f"Invalid token: {e}") from e
+
+        sub = payload.get("sub", "")
+        if not sub.startswith("job:"):
+            raise UnauthorizedError("Not a job-scoped token")
+
+        tenant_id = payload.get("custom:tenant_id")
+        if not tenant_id:
+            raise UnauthorizedError("Token missing tenant_id claim")
+
+        return JobClaims(
+            sub=sub,
+            tenant_id=tenant_id,
+            product=payload.get("product", ""),
+            stage=payload.get("stage", ""),
+            scope=payload.get("scope", []),
+            orchestrator=payload.get("orchestrator", "prefect"),
+            compute=payload.get("compute", "ecs"),
+            raw=payload,
+        )

pulse_engine/core/logging.py

@@ -0,0 +1,45 @@
+import logging
+import sys
+
+import structlog
+
+
+def setup_logging(log_level: str = "INFO", env: str = "development") -> None:
+    shared_processors: list[structlog.types.Processor] = [
+        structlog.contextvars.merge_contextvars,
+        structlog.stdlib.add_log_level,
+        structlog.stdlib.add_logger_name,
+        structlog.processors.TimeStamper(fmt="iso"),
+        structlog.processors.StackInfoRenderer(),
+        structlog.processors.UnicodeDecoder(),
+    ]
+
+    if env == "development":
+        renderer: structlog.types.Processor = structlog.dev.ConsoleRenderer()
+    else:
+        renderer = structlog.processors.JSONRenderer()
+
+    structlog.configure(
+        processors=[
+            *shared_processors,
+            structlog.stdlib.ProcessorFormatter.wrap_for_formatter,
+        ],
+        logger_factory=structlog.stdlib.LoggerFactory(),
+        wrapper_class=structlog.stdlib.BoundLogger,
+        cache_logger_on_first_use=True,
+    )
+
+    formatter = structlog.stdlib.ProcessorFormatter(
+        processors=[
+            structlog.stdlib.ProcessorFormatter.remove_processors_meta,
+            renderer,
+        ],
+    )
+
+    handler = logging.StreamHandler(sys.stdout)
+    handler.setFormatter(formatter)
+
+    root_logger = logging.getLogger()
+    root_logger.handlers.clear()
+    root_logger.addHandler(handler)
+    root_logger.setLevel(getattr(logging, log_level.upper(), logging.INFO))

pulse_engine/core/scope.py

@@ -0,0 +1,23 @@
+"""Scope enforcement for job-scoped tokens."""
+
+from typing import Any
+
+from fastapi import Depends, Request
+
+from pulse_engine.core.exceptions import ForbiddenError
+from pulse_engine.core.job_token import JobClaims
+
+
+def require_scope(scope: str) -> Any:
+    """FastAPI dependency that enforces scope on job-scoped tokens.
+
+    Cognito-authenticated requests bypass the check (no scope field).
+    """
+
+    def _check(request: Request) -> None:
+        claims = getattr(request.state, "user_claims", None)
+        if isinstance(claims, JobClaims):
+            if scope not in claims.scope:
+                raise ForbiddenError(f"Token missing required scope: {scope}")
+
+    return Depends(_check)

pulse_engine/core/security.py

@@ -0,0 +1,130 @@
+import time as _time
+from dataclasses import dataclass
+from typing import Any, Protocol
+
+import httpx
+from jose import JWTError, jwt
+
+from pulse_engine.core.exceptions import UnauthorizedError
+from pulse_engine.core.job_token import JobClaims
+
+_JWKS_TTL_SECONDS = 3600  # Re-fetch JWKS every hour
+
+
+@dataclass
+class CognitoClaims:
+    sub: str
+    email: str
+    tenant_id: str
+    raw: dict[str, Any]
+
+
+class TokenVerifier(Protocol):
+    async def verify(self, token: str) -> CognitoClaims | JobClaims: ...
+
+
+class CognitoTokenVerifier:
+    def __init__(self, jwks_url: str, issuer: str, audience: str) -> None:
+        self._jwks_url = jwks_url
+        self._issuer = issuer
+        self._audience = audience
+        self._jwks: dict[str, Any] | None = None
+        self._jwks_fetched_at: float = 0.0
+
+    async def _get_jwks(self, *, force_refresh: bool = False) -> dict[str, Any]:
+        """Return cached JWKS, refreshing when the TTL has elapsed or forced."""
+        now = _time.monotonic()
+        if (
+            self._jwks is None
+            or force_refresh
+            or (now - self._jwks_fetched_at) > _JWKS_TTL_SECONDS
+        ):
+            async with httpx.AsyncClient() as client:
+                resp = await client.get(self._jwks_url)
+                resp.raise_for_status()
+                self._jwks = resp.json()
+                self._jwks_fetched_at = now
+        assert self._jwks is not None  # guaranteed by the branch above
+        return self._jwks
+
+    async def verify(self, token: str) -> CognitoClaims:
+        try:
+            jwks = await self._get_jwks()
+            unverified_header = jwt.get_unverified_header(token)
+            kid = unverified_header.get("kid")
+            key = self._find_key(jwks, kid)
+
+            # On key miss, attempt a one-time forced refresh (handles key rotation)
+            if key is None:
+                jwks = await self._get_jwks(force_refresh=True)
+                key = self._find_key(jwks, kid)
+
+            if key is None:
+                raise UnauthorizedError("Invalid token: signing key not found")
+
+            payload = jwt.decode(
+                token,
+                key,
+                algorithms=["RS256"],
+                audience=self._audience,
+                issuer=self._issuer,
+            )
+        except JWTError as e:
+            raise UnauthorizedError("Invalid token") from e
+
+        tenant_id = payload.get("custom:tenant_id")
+        if not tenant_id:
+            raise UnauthorizedError("Token missing tenant_id claim")
+
+        return CognitoClaims(
+            sub=payload.get("sub", ""),
+            email=payload.get("email", ""),
+            tenant_id=tenant_id,
+            raw=payload,
+        )
+
+    @staticmethod
+    def _find_key(jwks: dict[str, Any], kid: str | None) -> dict[str, Any] | None:
+        for k in jwks.get("keys", []):
+            if k.get("kid") == kid:
+                return k  # type: ignore[no-any-return]
+        return None
+
+
+class MockTokenVerifier:
+    def __init__(self, secret: str, audience: str, issuer: str) -> None:
+        self._secret = secret
+        self._audience = audience
+        self._issuer = issuer
+
+    async def verify(self, token: str) -> CognitoClaims:
+        try:
+            payload = jwt.decode(
+                token,
+                self._secret,
+                algorithms=["HS256"],
+                audience=self._audience,
+                issuer=self._issuer,
+            )
+        except JWTError as e:
+            raise UnauthorizedError("Invalid token") from e
+
+        tenant_id = payload.get("custom:tenant_id")
+        if not tenant_id:
+            raise UnauthorizedError("Token missing tenant_id claim")
+
+        return CognitoClaims(
+            sub=payload.get("sub", ""),
+            email=payload.get("email", ""),
+            tenant_id=tenant_id,
+            raw=payload,
+        )
+
+
+def extract_bearer_token(authorization: str | None) -> str:
+    if not authorization:
+        raise UnauthorizedError("Missing authorization header")
+    parts = authorization.split()
+    if len(parts) != 2 or parts[0].lower() != "bearer":
+        raise UnauthorizedError("Invalid authorization header format")
+    return parts[1]

pulse_engine/database.py

@@ -0,0 +1,30 @@
+from sqlalchemy.ext.asyncio import (
+    AsyncEngine,
+    AsyncSession,
+    async_sessionmaker,
+    create_async_engine,
+)
+from sqlalchemy.orm import DeclarativeBase
+
+
+class Base(DeclarativeBase):
+    pass
+
+
+def build_async_engine(database_url: str) -> AsyncEngine:
+    return create_async_engine(
+        database_url,
+        echo=False,
+        future=True,
+        connect_args={
+            "statement_cache_size": 0,
+            "prepared_statement_cache_size": 0,
+            "prepared_statement_name_func": lambda: "",
+        },
+    )
+
+
+def build_session_factory(
+    engine: AsyncEngine,
+) -> async_sessionmaker[AsyncSession]:
+    return async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)

pulse_engine/dependencies.py

@@ -0,0 +1,166 @@
+from collections.abc import AsyncIterator
+from typing import TYPE_CHECKING
+
+from fastapi import Depends, Request
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from pulse_engine.config import Settings, get_settings
+from pulse_engine.core.job_token import JobClaims
+from pulse_engine.core.security import (
+    CognitoClaims,
+    CognitoTokenVerifier,
+    TokenVerifier,
+    extract_bearer_token,
+)
+from pulse_engine.services.opensearch import OpenSearchService
+from pulse_engine.storage.knowledge_base import KnowledgeBaseService
+
+if TYPE_CHECKING:
+    from pulse_engine.deployment.backends.registry import BackendRegistry
+    from pulse_engine.deployment.job_launcher import JobLauncher
+    from pulse_engine.deployment.service import DeploymentService
+    from pulse_engine.extractor.orchestrator.base import BaseOrchestratorAdapter
+    from pulse_engine.extractor.repository import JobRepository
+    from pulse_engine.extractor.service import JobService
+    from pulse_engine.processor.pipeline import ProcessingPipeline
+
+
+def get_token_verifier(
+    settings: Settings = Depends(get_settings),
+) -> TokenVerifier:
+    return CognitoTokenVerifier(
+        jwks_url=settings.cognito_jwks_url,
+        issuer=settings.cognito_issuer,
+        audience=settings.cognito_app_client_id,
+    )
+
+
+async def get_current_user(
+    request: Request,
+    verifier: TokenVerifier = Depends(get_token_verifier),
+) -> CognitoClaims | JobClaims:
+    token = extract_bearer_token(request.headers.get("authorization"))
+    return await verifier.verify(token)
+
+
+def get_opensearch(request: Request) -> OpenSearchService:
+    return request.app.state.opensearch  # type: ignore[no-any-return]
+
+
+def get_knowledge_base(request: Request) -> KnowledgeBaseService:
+    return request.app.state.knowledge_base  # type: ignore[no-any-return]
+
+
+async def get_db_session(request: Request) -> AsyncIterator[AsyncSession]:
+    factory = request.app.state.db_session_factory
+    if factory is None:
+        from pulse_engine.core.exceptions import ServiceUnavailableError
+
+        raise ServiceUnavailableError("Database not configured")
+    async with factory() as session:
+        try:
+            yield session
+            await session.commit()
+        except Exception:
+            await session.rollback()
+            raise
+
+
+def get_orchestrator(request: Request) -> "BaseOrchestratorAdapter":
+    return request.app.state.orchestrator_adapter  # type: ignore[no-any-return]
+
+
+def get_job_repository(
+    session: AsyncSession = Depends(get_db_session),
+) -> "JobRepository":
+    from pulse_engine.extractor.repository import JobRepository
+
+    return JobRepository(session)
+
+
+def get_backend_registry(
+    request: Request,
+    settings: Settings = Depends(get_settings),
+) -> "BackendRegistry":
+    # Cache on app state to avoid recreating on each request
+    if not hasattr(request.app.state, "backend_registry"):
+        from pulse_engine.deployment.backends.registry import BackendRegistry
+
+        request.app.state.backend_registry = BackendRegistry(settings)
+    return request.app.state.backend_registry  # type: ignore[no-any-return]
+
+
+def get_job_launcher(
+    request: Request,
+    session: AsyncSession = Depends(get_db_session),
+    registry: "BackendRegistry" = Depends(get_backend_registry),
+    settings: Settings = Depends(get_settings),
+) -> "JobLauncher":
+    from pulse_engine.core.exceptions import ServiceUnavailableError
+    from pulse_engine.core.job_token import JobTokenIssuer
+    from pulse_engine.deployment.backend_deployment_repository import (
+        BackendDeploymentRepository,
+    )
+    from pulse_engine.deployment.job_launcher import JobLauncher
+    from pulse_engine.deployment.repository import RegistrationRepository
+
+    reg_repo = RegistrationRepository(session)
+    backend_repo = BackendDeploymentRepository(session)
+    if not settings.pulse_job_token_secret:
+        raise ServiceUnavailableError("PULSE_JOB_TOKEN_SECRET is not configured")
+    token_issuer = JobTokenIssuer(secret=settings.pulse_job_token_secret)
+    return JobLauncher(
+        registration_repo=reg_repo,
+        backend_deployment_repo=backend_repo,
+        registry=registry,
+        token_issuer=token_issuer,
+        settings=settings,
+    )
+
+
+def get_job_service(
+    request: Request,
+    repo: "JobRepository" = Depends(get_job_repository),
+    orchestrator: "BaseOrchestratorAdapter" = Depends(get_orchestrator),
+    settings: Settings = Depends(get_settings),
+    session: AsyncSession = Depends(get_db_session),
+    job_launcher: "JobLauncher" = Depends(get_job_launcher),
+) -> "JobService":
+    from pulse_engine.extractor.service import JobService
+    from pulse_engine.extractor.stage_repository import StageRepository
+
+    stage_repo = StageRepository(session)
+    # NOTE: deployment_repository and token_issuer are intentionally NOT passed here.
+    # JobLauncher handles all triggering. The legacy path in _trigger_stage is
+    # disabled in production by omitting these args.
+    return JobService(
+        repository=repo,
+        orchestrator=orchestrator,
+        settings=settings,
+        stage_repository=stage_repo,
+        job_launcher=job_launcher,
+    )
+
+
+def get_processing_pipeline(request: Request) -> "ProcessingPipeline":
+    from pulse_engine.processor.pipeline import ProcessingPipeline
+
+    kb = request.app.state.knowledge_base
+    return ProcessingPipeline(kb_service=kb)
+
+
+def get_deployment_service(
+    session: AsyncSession = Depends(get_db_session),
+) -> "DeploymentService":
+    from pulse_engine.deployment.backend_deployment_repository import (
+        BackendDeploymentRepository,
+    )
+    from pulse_engine.deployment.repository import RegistrationRepository
+    from pulse_engine.deployment.service import DeploymentService
+
+    reg_repo = RegistrationRepository(session)
+    backend_repo = BackendDeploymentRepository(session)
+    return DeploymentService(
+        registration_repo=reg_repo,
+        backend_deployment_repo=backend_repo,
+    )

pulse_engine/deployment/backend_deployment_repository.py

@@ -0,0 +1,83 @@
+"""Data access layer for the product_deployments lazy cache table."""
+
+from __future__ import annotations
+
+from datetime import UTC, datetime
+
+import sqlalchemy as sa
+from sqlalchemy.dialects.postgresql import insert as pg_insert
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from pulse_engine.deployment.models import ProductBackendDeploymentModel
+
+
+class BackendDeploymentRepository:
+    """CRUD for product_deployments — the per-backend Prefect deployment ID cache.
+
+    Owns all reads, writes, and deletes for this table.
+    """
+
+    def __init__(self, session: AsyncSession) -> None:
+        self._session = session
+
+    async def get(
+        self,
+        product: str,
+        stage: str,
+        orchestrator: str,
+        compute: str,
+    ) -> ProductBackendDeploymentModel | None:
+        stmt = sa.select(ProductBackendDeploymentModel).where(
+            ProductBackendDeploymentModel.product == product,
+            ProductBackendDeploymentModel.stage == stage,
+            ProductBackendDeploymentModel.orchestrator == orchestrator,
+            ProductBackendDeploymentModel.compute == compute,
+        )
+        result = await self._session.execute(stmt)
+        return result.scalar_one_or_none()
+
+    async def upsert(
+        self,
+        product: str,
+        stage: str,
+        orchestrator: str,
+        compute: str,
+        deployment_id: str,
+    ) -> None:
+        """Insert or update deployment_id using SQL-level ON CONFLICT DO UPDATE.
+
+        Safe under concurrent triggers — two concurrent callers both complete
+        without IntegrityError; the last write wins.
+        """
+        now = datetime.now(UTC)
+        stmt = (
+            pg_insert(ProductBackendDeploymentModel)
+            .values(
+                product=product,
+                stage=stage,
+                orchestrator=orchestrator,
+                compute=compute,
+                deployment_id=deployment_id,
+                created_at=now,
+                updated_at=now,
+            )
+            .on_conflict_do_update(
+                index_elements=["product", "stage", "orchestrator", "compute"],
+                set_={"deployment_id": deployment_id, "updated_at": now},
+            )
+        )
+        await self._session.execute(stmt)
+        await self._session.commit()
+
+    async def delete_by_product_stage(self, product: str, stage: str) -> None:
+        """Delete all cached deployment IDs for this product+stage.
+
+        Called by DeploymentService when a product re-registers with a new image,
+        so JobLauncher recreates Prefect deployments on the next trigger.
+        """
+        stmt = sa.delete(ProductBackendDeploymentModel).where(
+            ProductBackendDeploymentModel.product == product,
+            ProductBackendDeploymentModel.stage == stage,
+        )
+        await self._session.execute(stmt)
+        await self._session.commit()

pulse_engine/deployment/backends/base.py

@@ -0,0 +1,50 @@
+"""Abstract base class for runner backends."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+from pulse_engine.extractor.orchestrator.base import OrchestratorRunStatus
+
+
+class BaseRunnerBackend(ABC):
+    """One implementation per (orchestrator, compute) pair.
+
+    Responsible for: backend setup, run-unit registration,
+    run triggering, status querying, and cancellation.
+    """
+
+    @abstractmethod
+    async def prepare(self) -> None:
+        """One-time setup before triggering (e.g. work pool creation).
+        No-op for backends that don't require it."""
+
+    @abstractmethod
+    async def register(
+        self,
+        product: str,
+        stage: str,
+        image: str,
+        entrypoint: str | None = None,
+    ) -> str:
+        """Register or update a runnable unit for this product+stage.
+
+        Returns a stable handle the launcher caches.
+        - Prefect backends: Prefect deployment UUID
+        - Native Lambda: "{product}-{stage}" (function name by convention)
+
+        entrypoint is required for Prefect backends, ignored for native backends.
+        """
+
+    @abstractmethod
+    async def trigger(self, handle: str, parameters: dict[str, Any]) -> str:
+        """Trigger a run. Returns a run ID for correlation."""
+
+    @abstractmethod
+    async def get_run_status(self, run_id: str) -> OrchestratorRunStatus:
+        """Fetch current status of a run."""
+
+    @abstractmethod
+    async def cancel_run(self, run_id: str) -> bool:
+        """Request cancellation. Return True if accepted."""

pulse_engine/deployment/backends/exceptions.py

@@ -0,0 +1,20 @@
+"""Exceptions for runner backends."""
+
+
+class BackendNotAvailableError(Exception):
+    """Raised when a requested (orchestrator, compute) backend is not available."""
+
+    def __init__(
+        self,
+        orchestrator: str,
+        compute: str,
+        available: list[str],
+    ) -> None:
+        self.orchestrator = orchestrator
+        self.compute = compute
+        self.available = available
+        msg = (
+            f"Backend not available: orchestrator={orchestrator!r}, "
+            f"compute={compute!r}. Available: {', '.join(available)}"
+        )
+        super().__init__(msg)