pulse-engine 0.2.0__py3-none-any.whl

pulse_engine/runners/prefect_runner.py

@@ -0,0 +1,33 @@
+"""Prefect flow wrapper for product containers.
+
+This module is the ONLY place in the engine that imports Prefect.
+Product containers must NOT import prefect directly — they expose a plain
+`entrypoint:run` function and this runner wraps it with @flow at runtime.
+
+All Prefect-based backends (ECS, Kubernetes) use this as the flow entrypoint:
+    pulse_engine.runners.prefect_runner:main
+"""
+
+from typing import Any
+
+from prefect import flow
+
+
+@flow(name="pulse-stage-runner")  # type: ignore[untyped-decorator]
+def main(
+    job_id: str,
+    chain: bool,
+    config: dict[str, Any],
+    pulse_api_token: str,
+    pulse_engine_url: str,
+) -> None:
+    """Engine-owned Prefect flow. Delegates to the product's plain entrypoint."""
+    from entrypoint import run  # imported at call time from the product container
+
+    run(
+        job_id=job_id,
+        chain=chain,
+        config=config,
+        pulse_api_token=pulse_api_token,
+        pulse_engine_url=pulse_engine_url,
+    )

pulse_engine/s3.py

@@ -0,0 +1,72 @@
+"""S3-backed inter-stage data exchange with NDJSON format and _SUCCESS sentinel."""
+
+from __future__ import annotations
+
+import json
+from typing import Any
+
+from botocore.exceptions import ClientError
+
+
+class StageDataNotReady(Exception):
+    """Raised when the upstream _SUCCESS sentinel is missing."""
+
+
+class S3Stage:
+    """Read/write NDJSON data between pipeline stages via S3."""
+
+    def __init__(self, bucket: str, job_id: str, s3_client: Any) -> None:
+        self._bucket = bucket
+        self._job_id = job_id
+        self._s3 = s3_client
+
+    @property
+    def raw_prefix(self) -> str:
+        return f"jobs/{self._job_id}/raw/"
+
+    @property
+    def processed_prefix(self) -> str:
+        return f"jobs/{self._job_id}/processed/"
+
+    def write_raw(self, documents: list[dict[str, Any]]) -> None:
+        self._write(self.raw_prefix, documents)
+
+    def write_processed(self, documents: list[dict[str, Any]]) -> None:
+        self._write(self.processed_prefix, documents)
+
+    def read_raw(self) -> list[dict[str, Any]]:
+        return self._read(self.raw_prefix)
+
+    def read_processed(self) -> list[dict[str, Any]]:
+        return self._read(self.processed_prefix)
+
+    def _write(self, prefix: str, documents: list[dict[str, Any]]) -> None:
+        ndjson = "\n".join(json.dumps(doc) for doc in documents)
+        self._s3.put_object(
+            Bucket=self._bucket,
+            Key=f"{prefix}data.ndjson",
+            Body=ndjson.encode(),
+            ContentType="application/x-ndjson",
+        )
+        self._s3.put_object(
+            Bucket=self._bucket,
+            Key=f"{prefix}_SUCCESS",
+            Body=b"",
+        )
+
+    def _read(self, prefix: str) -> list[dict[str, Any]]:
+        # Check sentinel — only catch 404; let other S3 errors propagate
+        try:
+            self._s3.head_object(Bucket=self._bucket, Key=f"{prefix}_SUCCESS")
+        except ClientError as e:
+            if e.response["Error"]["Code"] == "404":
+                msg = (
+                    f"Sentinel {prefix}_SUCCESS not found"
+                    " — upstream stage may not be complete"
+                )
+                raise StageDataNotReady(msg) from e
+            raise  # re-raise permissions errors, network errors, etc.
+
+        resp = self._s3.get_object(Bucket=self._bucket, Key=f"{prefix}data.ndjson")
+        body = resp["Body"].read().decode()
+        return [json.loads(line) for line in body.strip().split("\n") if line.strip()]

pulse_engine/secrets.py

@@ -0,0 +1,46 @@
+"""AWS Secrets Manager utilities."""
+
+from __future__ import annotations
+
+import asyncio
+
+import structlog
+
+logger = structlog.get_logger(__name__)
+
+
+async def fetch_secret(secret_id: str, region: str = "ap-south-1") -> str:
+    """Fetch a plaintext secret string from AWS Secrets Manager.
+
+    Runs the boto3 call in a thread so it is safe to await from async code.
+
+    Args:
+        secret_id: The secret name or ARN.
+        region: AWS region for the Secrets Manager endpoint.
+
+    Returns:
+        The secret value as a plain string.
+
+    Raises:
+        RuntimeError: If the secret cannot be retrieved or is empty/binary.
+    """
+    import boto3
+    from botocore.exceptions import ClientError
+
+    def _get() -> str:
+        client = boto3.client("secretsmanager", region_name=region)
+        try:
+            resp = client.get_secret_value(SecretId=secret_id)
+        except ClientError as exc:
+            raise RuntimeError(
+                f"Failed to fetch secret '{secret_id}' from Secrets Manager: {exc}"
+            ) from exc
+        value: str = resp.get("SecretString") or ""
+        if not value:
+            raise RuntimeError(
+                f"Secret '{secret_id}' is empty or binary — expected a plaintext string"
+            )
+        return value
+
+    logger.debug("secrets_manager_fetch", secret_id=secret_id, region=region)
+    return await asyncio.to_thread(_get)

pulse_engine/services/bootstrap.py

@@ -0,0 +1,211 @@
+"""Shared service initialization for FastAPI and MCP entry points."""
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from dataclasses import dataclass, field
+from types import EllipsisType
+from typing import TYPE_CHECKING, Any
+
+import boto3
+
+from pulse_engine.config import Settings
+from pulse_engine.core.logging import setup_logging
+from pulse_engine.database import build_async_engine, build_session_factory
+from pulse_engine.extractor.orchestrator import get_orchestrator_adapter
+from pulse_engine.processor.base import (
+    BaseCoreProcessor,
+    BasePostprocessor,
+    BasePreprocessor,
+)
+from pulse_engine.processor.pipeline import _SENTINEL, ProcessingPipeline
+from pulse_engine.services.opensearch import OpenSearchService
+from pulse_engine.storage.connectors.athena import AthenaConnector
+from pulse_engine.storage.connectors.opensearch import OpenSearchConnector
+from pulse_engine.storage.knowledge_base import KnowledgeBaseService
+
+if TYPE_CHECKING:
+    from celery import Celery
+    from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker
+
+    from pulse_engine.extractor.orchestrator.base import BaseOrchestratorAdapter
+    from pulse_engine.registry import ProductManifest
+
+
+@dataclass
+class ServiceContainer:
+    settings: Settings
+    opensearch: OpenSearchService
+    kb_service: KnowledgeBaseService
+    db_session_factory: async_sessionmaker[AsyncSession] | None
+    orchestrator_adapter: BaseOrchestratorAdapter
+    pipeline: ProcessingPipeline | None = None
+    celery_app: Celery | None = None
+    s3_client: Any = None
+    token_issuer: Any = None
+    pipeline_translators: dict[str, Any] = field(default_factory=dict)
+    pipeline_status_providers: dict[str, Any] = field(default_factory=dict)
+    _engine: AsyncEngine | None = field(default=None, repr=False)
+    _opensearch_connector: OpenSearchConnector | None = field(default=None, repr=False)
+    _athena_connector: AthenaConnector | None = field(default=None, repr=False)
+
+
+def _resolve_stage(
+    value: Any,
+) -> Any:
+    """Convert manifest Ellipsis convention into pipeline sentinel."""
+    if isinstance(value, EllipsisType):
+        return _SENTINEL
+    return value
+
+
+def _build_celery(
+    settings: Settings, manifest: ProductManifest | None = None
+) -> Celery:
+    from pulse_engine.worker import create_celery_app
+
+    return create_celery_app(settings, manifest)
+
+
+@asynccontextmanager
+async def bootstrap_services(
+    settings: Settings,
+    manifest: ProductManifest | None = None,
+) -> AsyncIterator[ServiceContainer]:
+    setup_logging(log_level=settings.log_level, env=settings.app_env)
+
+    # OpenSearch — connect with auth if credentials provided
+    http_auth: tuple[str, str] | None = None
+    if settings.opensearch_username and settings.opensearch_password:
+        http_auth = (settings.opensearch_username, settings.opensearch_password)
+
+    opensearch = OpenSearchService(
+        url=settings.opensearch_url,
+        http_auth=http_auth,
+        use_ssl=settings.opensearch_use_ssl,
+        verify_certs=settings.opensearch_verify_certs,
+    )
+
+    opensearch_connector = OpenSearchConnector(
+        opensearch,
+        settings.embedding_dimension,
+        index_prefix=settings.opensearch_index_prefix,
+    )
+    await opensearch_connector.initialize()
+
+    # Athena — uses its own AWS credentials; product specifies the database
+    athena_connector: AthenaConnector | None = None
+    if settings.athena_output_location:
+        athena_kwargs: dict[str, str] = {"region_name": settings.aws_region}
+        if settings.athena_aws_access_key_id:
+            athena_kwargs["aws_access_key_id"] = settings.athena_aws_access_key_id
+            athena_kwargs["aws_secret_access_key"] = (
+                settings.athena_aws_secret_access_key
+            )
+        athena_client = boto3.client("athena", **athena_kwargs)  # type: ignore[call-overload]
+        athena_connector = AthenaConnector(
+            athena_client=athena_client,
+            database=manifest.athena_database if manifest else "",
+            output_location=settings.athena_output_location,
+            workgroup=settings.athena_workgroup,
+            query_timeout_seconds=settings.athena_query_timeout_seconds,
+        )
+        await athena_connector.initialize()
+
+    kb_service = KnowledgeBaseService(opensearch_connector, athena_connector)
+
+    engine = None
+    db_session_factory = None
+    if settings.database_url:
+        engine = build_async_engine(settings.database_url)
+        db_session_factory = build_session_factory(engine)
+
+    orchestrator_adapter = get_orchestrator_adapter(settings)
+
+    # Celery / Redis
+    celery_app = _build_celery(settings, manifest) if settings.redis_url else None
+
+    # S3 client for inter-stage data
+    s3_client = None
+    if settings.pulse_s3_bucket:
+        s3_client = boto3.client(
+            "s3",
+            region_name=settings.aws_region,
+            aws_access_key_id=settings.aws_access_key_id or None,
+            aws_secret_access_key=settings.aws_secret_access_key or None,
+        )
+
+    # Job token issuer
+    token_issuer = None
+    if settings.pulse_job_token_secret:
+        from pulse_engine.core.job_token import JobTokenIssuer
+
+        token_issuer = JobTokenIssuer(
+            secret=settings.pulse_job_token_secret,
+        )
+
+    # Pipeline orchestration translators and status providers
+    pipeline_translators: dict[str, Any] = {}
+    pipeline_status_providers: dict[str, Any] = {}
+    if settings.prefect_api_url:
+        from pulse_engine.pipeline.translators.prefect_status import (
+            PrefectStatusProvider,
+        )
+        from pulse_engine.pipeline.translators.prefect_translator import (
+            PrefectTranslator,
+        )
+
+        pipeline_translators["prefect"] = PrefectTranslator(
+            prefect_api_url=settings.prefect_api_url,
+            prefect_api_key=settings.prefect_api_key,
+            work_pool_name=settings.prefect_ecs_work_pool_name,
+            engine_image=settings.prefect_engine_image,
+        )
+        pipeline_status_providers["prefect"] = PrefectStatusProvider(
+            prefect_api_url=settings.prefect_api_url,
+            prefect_api_key=settings.prefect_api_key,
+        )
+
+    # Build pluggable pipeline
+    pre: BasePreprocessor | None = _SENTINEL
+    core: BaseCoreProcessor | None = _SENTINEL
+    post: BasePostprocessor | None = _SENTINEL
+    if manifest:
+        pre = _resolve_stage(manifest.preprocessor)
+        core = _resolve_stage(manifest.core_processor)
+        post = _resolve_stage(manifest.postprocessor)
+
+    pipeline = ProcessingPipeline(
+        kb_service=kb_service,
+        preprocessor=pre,
+        core_processor=core,
+        postprocessor=post,
+    )
+
+    container = ServiceContainer(
+        settings=settings,
+        opensearch=opensearch,
+        kb_service=kb_service,
+        db_session_factory=db_session_factory,
+        orchestrator_adapter=orchestrator_adapter,
+        pipeline=pipeline,
+        celery_app=celery_app,
+        s3_client=s3_client,
+        token_issuer=token_issuer,
+        pipeline_translators=pipeline_translators,
+        pipeline_status_providers=pipeline_status_providers,
+        _engine=engine,
+        _opensearch_connector=opensearch_connector,
+        _athena_connector=athena_connector,
+    )
+
+    try:
+        yield container
+    finally:
+        if engine:
+            await engine.dispose()
+        await opensearch_connector.teardown()
+        if athena_connector:
+            await athena_connector.teardown()
+        await opensearch.close()

pulse_engine/services/opensearch.py

@@ -0,0 +1,84 @@
+from typing import Any
+
+import structlog
+from opensearchpy import AsyncOpenSearch
+
+logger = structlog.get_logger()
+
+
+class OpenSearchService:
+    def __init__(
+        self,
+        url: str,
+        http_auth: tuple[str, str] | None = None,
+        use_ssl: bool = False,
+        verify_certs: bool = True,
+        ssl_show_warn: bool = True,
+    ) -> None:
+        hosts = [url]
+        kwargs: dict[str, Any] = {
+            "hosts": hosts,
+            "use_ssl": use_ssl,
+            "verify_certs": verify_certs,
+            "ssl_show_warn": ssl_show_warn,
+        }
+        if http_auth:
+            kwargs["http_auth"] = http_auth
+
+        self._client = AsyncOpenSearch(**kwargs)
+
+    async def ping(self) -> bool:
+        try:
+            return bool(await self._client.ping())
+        except Exception:
+            logger.warning("opensearch_ping_failed")
+            return False
+
+    async def index_document(
+        self, index: str, doc_id: str | None, body: dict[str, Any]
+    ) -> dict[str, Any]:
+        result: dict[str, Any] = await self._client.index(
+            index=index, id=doc_id, body=body
+        )
+        return result
+
+    async def get_document(self, index: str, doc_id: str) -> dict[str, Any]:
+        result: dict[str, Any] = await self._client.get(index=index, id=doc_id)
+        return result
+
+    async def search(self, index: str, query: dict[str, Any]) -> dict[str, Any]:
+        result: dict[str, Any] = await self._client.search(index=index, body=query)
+        return result
+
+    async def delete_document(self, index: str, doc_id: str) -> dict[str, Any]:
+        result: dict[str, Any] = await self._client.delete(index=index, id=doc_id)
+        return result
+
+    async def create_index(
+        self, index: str, body: dict[str, Any]
+    ) -> dict[str, Any] | None:
+        if await self._client.indices.exists(index=index):
+            logger.info("index_already_exists", index=index)
+            return None
+        result: dict[str, Any] = await self._client.indices.create(
+            index=index, body=body
+        )
+        return result
+
+    async def delete_index(self, index: str) -> dict[str, Any]:
+        result: dict[str, Any] = await self._client.indices.delete(index=index)
+        return result
+
+    async def bulk(self, body: list[dict[str, Any]]) -> dict[str, Any]:
+        result: dict[str, Any] = await self._client.bulk(body=body)
+        return result
+
+    async def index_stats(self, index: str) -> dict[str, Any]:
+        result: dict[str, Any] = await self._client.indices.stats(index=index)
+        return result
+
+    async def index_exists(self, index: str) -> bool:
+        return bool(await self._client.indices.exists(index=index))
+
+    async def close(self) -> None:
+        await self._client.close()

pulse_engine/storage/connectors/athena.py

@@ -0,0 +1,226 @@
+import asyncio
+import re
+import time
+from typing import Any
+
+import structlog
+
+from pulse_engine.core.exceptions import (
+    AppError,
+    BadRequestError,
+    ServiceUnavailableError,
+)
+from pulse_engine.storage.connectors.base import BaseStorageConnector
+from pulse_engine.storage.schemas import (
+    ConnectorHealth,
+    Document,
+    QueryResult,
+    SearchQuery,
+    SearchResult,
+    StoreResult,
+)
+
+logger = structlog.get_logger()
+
+_TENANT_ID_RE = re.compile(r"^[a-zA-Z0-9_-]{1,128}$")
+
+# SQL keywords that are forbidden in user-supplied queries
+_FORBIDDEN_SQL_RE = re.compile(
+    r"""
+    \b(
+        INSERT\s+INTO | UPDATE\s+\S+\s+SET | DELETE\s+FROM |
+        DROP\s+(TABLE|DATABASE|INDEX|VIEW|SCHEMA) |
+        ALTER\s+(TABLE|DATABASE|INDEX|VIEW|SCHEMA) |
+        CREATE\s+(TABLE|DATABASE|INDEX|VIEW|SCHEMA) |
+        TRUNCATE | GRANT | REVOKE | EXEC | EXECUTE |
+        MERGE\s+INTO | CALL
+    )\b
+    """,
+    re.IGNORECASE | re.VERBOSE,
+)
+
+# Detect multiple SQL statements (semicolons not inside string literals)
+_MULTI_STATEMENT_RE = re.compile(r";\s*\S")
+
+# Detect UNION-based injection attempts
+_UNION_RE = re.compile(r"\bUNION\b\s+(ALL\s+)?\bSELECT\b", re.IGNORECASE)
+
+
+class AthenaConnector(BaseStorageConnector):
+    def __init__(
+        self,
+        athena_client: Any,
+        database: str,
+        output_location: str,
+        workgroup: str = "primary",
+        query_timeout_seconds: int = 60,
+    ) -> None:
+        self._client = athena_client
+        self._database = database
+        self._output_location = output_location
+        self._workgroup = workgroup
+        self._query_timeout = query_timeout_seconds
+
+    async def initialize(self) -> None:
+        logger.info("athena_connector_initialized")
+
+    async def store(self, tenant_id: str, documents: list[Document]) -> StoreResult:
+        raise NotImplementedError("Athena connector does not support store operations")
+
+    async def retrieve(self, tenant_id: str, doc_id: str) -> Document | None:
+        raise NotImplementedError(
+            "Athena connector does not support retrieve operations"
+        )
+
+    async def search(self, tenant_id: str, query: SearchQuery) -> SearchResult:
+        raise NotImplementedError("Athena connector does not support search operations")
+
+    @staticmethod
+    def _validate_sql(sql: str) -> None:
+        """Reject dangerous SQL patterns before execution.
+
+        Only SELECT statements are permitted. DDL, DML mutations,
+        multi-statement payloads, and UNION injections are blocked.
+        """
+        stripped = sql.strip().rstrip(";").strip()
+
+        if not stripped.upper().startswith("SELECT"):
+            raise BadRequestError("Only SELECT queries are allowed against Athena")
+
+        if _FORBIDDEN_SQL_RE.search(stripped):
+            raise BadRequestError("Query contains forbidden SQL keywords")
+
+        if _MULTI_STATEMENT_RE.search(stripped):
+            raise BadRequestError("Multiple SQL statements are not allowed")
+
+        if _UNION_RE.search(stripped):
+            raise BadRequestError("UNION SELECT queries are not allowed")
+
+    @staticmethod
+    def _inject_tenant_filter(sql: str, tenant_id: str) -> str:
+        if not _TENANT_ID_RE.match(tenant_id):
+            raise BadRequestError("Invalid tenant ID format")
+        where_pattern = re.compile(r"\bWHERE\b", re.IGNORECASE)
+        match = where_pattern.search(sql)
+        tenant_clause = f"tenant_id = '{tenant_id}'"
+
+        if match:
+            insert_pos = match.end()
+            return f"{sql[:insert_pos]} {tenant_clause} AND{sql[insert_pos:]}"
+        else:
+            order_pattern = re.compile(
+                r"\b(ORDER BY|GROUP BY|HAVING|LIMIT)\b", re.IGNORECASE
+            )
+            order_match = order_pattern.search(sql)
+            if order_match:
+                insert_pos = order_match.start()
+                return f"{sql[:insert_pos]}WHERE {tenant_clause} {sql[insert_pos:]}"
+            return f"{sql} WHERE {tenant_clause}"
+
+    async def execute_query(
+        self,
+        tenant_id: str,
+        sql: str,
+        parameters: dict[str, Any] | None = None,
+        database: str | None = None,
+    ) -> QueryResult:
+        """Execute an Athena SQL query.
+
+        Args:
+            database: The Athena database to query. Products specify this at
+                query time; falls back to the connector-level default.
+        """
+        self._validate_sql(sql)
+        filtered_sql = self._inject_tenant_filter(sql, tenant_id)
+        target_database = database or self._database
+
+        try:
+            start = time.monotonic()
+            query_kwargs: dict[str, Any] = {
+                "QueryString": filtered_sql,
+                "ResultConfiguration": {"OutputLocation": self._output_location},
+                "WorkGroup": self._workgroup,
+            }
+            if target_database:
+                query_kwargs["QueryExecutionContext"] = {"Database": target_database}
+
+            response = await asyncio.to_thread(
+                self._client.start_query_execution,
+                **query_kwargs,
+            )
+            query_id = response["QueryExecutionId"]
+
+            elapsed = 0.0
+            wait = 0.5
+            while elapsed < self._query_timeout:
+                status_resp = await asyncio.to_thread(
+                    self._client.get_query_execution,
+                    QueryExecutionId=query_id,
+                )
+                state = status_resp["QueryExecution"]["Status"]["State"]
+
+                if state == "SUCCEEDED":
+                    break
+                elif state == "FAILED":
+                    reason = status_resp["QueryExecution"]["Status"].get(
+                        "StateChangeReason", "Query failed"
+                    )
+                    logger.error(
+                        "athena_query_failed", reason=reason, query_id=query_id
+                    )
+                    raise AppError("Query execution failed")
+                elif state == "CANCELLED":
+                    raise AppError("Athena query was cancelled")
+
+                await asyncio.sleep(wait)
+                elapsed += wait
+                wait = min(wait * 2, 5.0)
+            else:
+                raise ServiceUnavailableError("Athena query timed out")
+
+            results_resp = await asyncio.to_thread(
+                self._client.get_query_results,
+                QueryExecutionId=query_id,
+            )
+
+            result_set = results_resp["ResultSet"]
+            columns = [
+                col["Label"] for col in result_set["ResultSetMetadata"]["ColumnInfo"]
+            ]
+
+            rows: list[list[Any]] = []
+            data_rows = result_set.get("Rows", [])
+            for row in data_rows[1:]:
+                rows.append([d.get("VarCharValue", "") for d in row["Data"]])
+
+            execution_time_ms = int((time.monotonic() - start) * 1000)
+
+            return QueryResult(
+                columns=columns,
+                rows=rows,
+                row_count=len(rows),
+                execution_time_ms=execution_time_ms,
+            )
+        except (AppError, ServiceUnavailableError):
+            raise
+        except Exception as e:
+            logger.error("athena_query_error", exc_info=True)
+            raise ServiceUnavailableError("Query execution unavailable") from e
+
+    async def delete(self, tenant_id: str, doc_id: str) -> bool:
+        raise NotImplementedError("Athena connector does not support delete operations")
+
+    async def health_check(self) -> ConnectorHealth:
+        start = time.monotonic()
+        try:
+            await asyncio.to_thread(self._client.list_work_groups)
+            latency = (time.monotonic() - start) * 1000
+            return ConnectorHealth(connector="athena", status="up", latency_ms=latency)
+        except Exception:
+            latency = (time.monotonic() - start) * 1000
+            return ConnectorHealth(
+                connector="athena", status="down", latency_ms=latency
+            )
+
+    async def teardown(self) -> None:
+        logger.info("athena_connector_teardown")