DataOps-platform/app/core/connectors/runtime.py

"""Connector execution controls: idempotency, bounded retry, cancellation and evidence."""

from __future__ import annotations

import hashlib
import json
import re
import threading
import time
import uuid
from collections import deque
from dataclasses import replace

from app.core.connectors.errors import (
    ConnectorCancelledError,
    ConnectorConfigurationError,
    ConnectorConflictError,
    ConnectorRateLimitError,
    classify_error,
)
from app.core.connectors.sdk import OperationResult
from app.core.connectors.store import UpdateOutcome

MAX_EVIDENCE_BYTES = 32768
MAX_CURSOR_BYTES = 32768
MAX_CHECKPOINT_BYTES = 262144
MAX_RECORDS_BYTES = 1048576
MAX_RECORDS = 10000
MAX_TOTAL_ATTEMPTS = 5
SENSITIVE_KEYS = {
    "password",
    "passwd",
    "secret",
    "token",
    "api_key",
    "authorization",
    "credential",
}


def deterministic_idempotency_key(connector_id, version, request):
    payload = {
        "connector_id": connector_id,
        "version": version,
        "source_uid": request.source_uid,
        "principal_uid": request.principal_uid,
        "business_domain_uid": request.business_domain_uid,
        "environment": request.environment,
        "process_key": request.process_key,
        "source_binding_uid": request.source_binding_uid,
        "source_binding_version": request.source_binding_version,
        "operation": request.operation,
        "config": request.config,
        "scope": request.scope,
        "cursor": request.cursor,
        "checkpoint": request.checkpoint,
        "dry_run": request.dry_run,
        "client_idempotency_hint": request.idempotency_key,
    }
    encoded = json.dumps(
        payload, sort_keys=True, separators=(",", ":"), ensure_ascii=True
    ).encode()
    return hashlib.sha256(encoded).hexdigest()


def redact_evidence(
    value, *, max_bytes=MAX_EVIDENCE_BYTES, summarize=True, truncate_lists=True
):
    def visit(item):
        if isinstance(item, dict):
            return {
                str(key): "[REDACTED]"
                if any(marker in str(key).lower() for marker in SENSITIVE_KEYS)
                else visit(child)
                for key, child in item.items()
            }
        if isinstance(item, (list, tuple)):
            values = item[:1000] if truncate_lists else item
            return [visit(child) for child in values]
        if isinstance(item, str):
            text_value = item[:4096]
            text_value = re.sub(
                r"(?i)\bBearer\s+[A-Za-z0-9._~+/-]+=*", "Bearer [REDACTED]", text_value
            )
            text_value = re.sub(r"\bdopc_[A-Za-z0-9_-]+", "[REDACTED]", text_value)
            text_value = re.sub(
                r"(?i)(password|token|secret|api[_-]?key)=([^&\s]+)",
                r"\1=[REDACTED]",
                text_value,
            )
            text_value = re.sub(
                r"(?i)(https?://)[^/@\s]+@", r"\1[REDACTED]@", text_value
            )
            return text_value
        return item

    safe = visit(value)
    encoded = json.dumps(safe, sort_keys=True, default=str).encode()
    if summarize and len(encoded) > max_bytes:
        return {
            "truncated": True,
            "sha256": hashlib.sha256(encoded).hexdigest(),
            "original_bytes": len(encoded),
        }
    return safe


def _bounded_channel(value, *, max_bytes, channel):
    raw = json.dumps(
        value, sort_keys=True, separators=(",", ":"), default=str
    ).encode()
    if len(raw) > max_bytes:
        raise ConnectorConfigurationError(f"connector {channel} exceeds safe limit")
    safe = redact_evidence(
        value, max_bytes=max_bytes, summarize=False, truncate_lists=False
    )
    encoded = json.dumps(
        safe, sort_keys=True, separators=(",", ":"), default=str
    ).encode()
    if len(encoded) > max_bytes:
        raise ConnectorConfigurationError(f"connector {channel} exceeds safe limit")
    return safe


def sanitize_operation_result(result):
    if not isinstance(result.records, (list, tuple)) or len(result.records) > MAX_RECORDS:
        raise ConnectorConfigurationError("connector records exceed safe limit")
    records = _bounded_channel(
        list(result.records), max_bytes=MAX_RECORDS_BYTES, channel="records"
    )
    if any(not isinstance(item, dict) for item in records):
        raise ConnectorConfigurationError("connector records must be objects")
    cursor = _bounded_channel(
        dict(result.cursor), max_bytes=MAX_CURSOR_BYTES, channel="cursor"
    )
    checkpoint = _bounded_channel(
        dict(result.checkpoint), max_bytes=MAX_CHECKPOINT_BYTES, channel="checkpoint"
    )
    evidence = _bounded_channel(
        dict(result.evidence), max_bytes=MAX_EVIDENCE_BYTES, channel="evidence"
    )
    return replace(
        result,
        records=tuple(records),
        cursor=cursor,
        checkpoint=checkpoint,
        evidence=evidence,
    )


def snapshot_diff(previous, current):
    def identity(item):
        if isinstance(item, dict):
            return (
                str(item.get("asset_key") or item.get("key") or "")
                + ":"
                + str(item.get("field") or "")
            )
        return json.dumps(item, sort_keys=True, default=str)

    before = {identity(item): item for item in previous}
    after = {identity(item): item for item in current}
    shared = before.keys() & after.keys()
    return {
        "added": tuple(after[key] for key in sorted(after.keys() - before.keys())),
        "removed": tuple(before[key] for key in sorted(before.keys() - after.keys())),
        "changed": tuple(
            {"before": before[key], "after": after[key]}
            for key in sorted(shared)
            if before[key] != after[key]
        ),
    }


class InMemoryRunStore:
    """Test/reference store. Production API injects the PostgreSQL repository."""

    def __init__(self):
        self._lock = threading.RLock()
        self._runs = {}
        self._hints = {}

    def claim(self, key, record):
        with self._lock:
            hint = record.get("client_hint_hash")
            if hint and hint in self._hints and self._hints[hint] != key:
                raise ConnectorConflictError("idempotency hint conflicts with another request")
            if key in self._runs:
                if self._runs[key].get("request_hash") != record.get("request_hash"):
                    raise ConnectorConflictError("idempotency request binding conflicts")
                return self._runs[key], False
            self._runs[key] = dict(record)
            if hint:
                self._hints[hint] = key
            return self._runs[key], True

    def update(self, key, **values):
        with self._lock:
            record = self._runs[key]
            current = record.get("status")
            requested = values.get("status")
            expected_attempt = values.pop("expected_attempt", None)
            lease_token = values.pop("lease_token", None)
            acquired = True
            if requested == "resumable":
                acquired = current in {"failed", "cancelled"}
            elif requested == "cancelled":
                acquired = current == "running"
            elif requested == "running":
                attempt = int(values.get("attempt_count", -1))
                acquired = (
                    current in {"running", "resumable", "failed"}
                    and int(record.get("attempt_count", 0)) < attempt
                )
            elif requested in {"succeeded", "dry_run", "failed"}:
                acquired = (
                    current == "running"
                    and not record.get("cancel_requested", False)
                    and int(record.get("attempt_count", -1)) == int(expected_attempt or -1)
                    and record.get("attempt_lease_token") == lease_token
                )
            if not acquired:
                return UpdateOutcome(dict(record), False)
            self._runs[key].update(values)
            if requested == "running":
                self._runs[key]["attempt_lease_token"] = lease_token
            return UpdateOutcome(dict(self._runs[key]), True)

    def cancel(self, key):
        outcome = self.update(key, status="cancelled", cancel_requested=True)
        if not outcome.acquired:
            raise ConnectorConfigurationError("connector run cannot be cancelled")
        return outcome.record

    def is_cancel_requested(self, key):
        with self._lock:
            return bool(self._runs.get(key, {}).get("cancel_requested"))

    def get(self, key):
        with self._lock:
            value = self._runs.get(key)
            if value is None:
                canonical = self._hints.get(hashlib.sha256(str(key).encode()).hexdigest())
                value = self._runs.get(canonical) if canonical else None
            return dict(value) if value else None

    def list(self):
        with self._lock:
            return [dict(item) for item in self._runs.values()]


class SlidingWindowLimiter:
    def __init__(self, limit=30, window_seconds=60, clock=None):
        self.limit = int(limit)
        self.window_seconds = float(window_seconds)
        self.clock = clock or time.monotonic
        self._lock = threading.Lock()
        self._events = {}

    def acquire(self, key):
        now = self.clock()
        with self._lock:
            events = self._events.setdefault(key, deque())
            while events and events[0] <= now - self.window_seconds:
                events.popleft()
            if len(events) >= self.limit:
                raise ConnectorRateLimitError("connector runtime rate limit exceeded")
            events.append(now)


class ConnectorRuntime:
    def __init__(
        self, registry, store=None, limiter=None, max_attempts=3, sleeper=None
    ):
        if max_attempts < 1 or max_attempts > 5:
            raise ConnectorConfigurationError("max_attempts must be between 1 and 5")
        self.registry = registry
        self.store = store or InMemoryRunStore()
        self.limiter = limiter or SlidingWindowLimiter()
        self.max_attempts = max_attempts
        self.sleeper = sleeper or time.sleep

    def _acquire_rate_limit(self, connector_id, source_uid):
        key = f"{connector_id}:{source_uid}"
        if hasattr(self.store, "acquire_rate_limit"):
            self.store.acquire_rate_limit(key)
        else:
            self.limiter.acquire(key)

    def _run_operation(self, key, record, request, operation):
        first_attempt = int(record.get("attempt_count", 0)) + 1
        remaining = MAX_TOTAL_ATTEMPTS - first_attempt + 1
        invocation_attempts = min(self.max_attempts, remaining)
        if invocation_attempts <= 0:
            raise ConnectorConfigurationError(
                "connector run attempt budget is exhausted"
            )
        for attempt in range(first_attempt, first_attempt + invocation_attempts):
            lease_token = str(uuid.uuid4())
            started = self.store.update(
                key,
                attempt_count=attempt,
                status="running",
                error_category=None,
                error_code=None,
                lease_token=lease_token,
            )
            if not started.acquired:
                if started.record.get("status") == "cancelled":
                    raise ConnectorCancelledError()
                raise ConnectorConfigurationError(
                    "connector run attempt is already owned"
                )
            started_record = started.record
            if started_record.get("status") == "cancelled":
                raise ConnectorCancelledError()
            if (
                started_record.get("status") != "running"
                or int(started_record.get("attempt_count", -1)) != attempt
            ):
                raise ConnectorConfigurationError(
                    "connector run attempt could not be claimed"
                )
            try:
                def cancel_probe():
                    return bool(
                        hasattr(self.store, "is_cancel_requested")
                        and self.store.is_cancel_requested(key)
                    )
                attempt_request = replace(
                    request,
                    run_key=key,
                    lease_token=lease_token,
                    cancel_probe=cancel_probe,
                )
                if cancel_probe():
                    raise ConnectorCancelledError()
                result = operation(attempt_request)
                if not isinstance(result, OperationResult):
                    raise ConnectorConfigurationError(
                        "connector returned an invalid result"
                    )
                if result.status != "succeeded":
                    raise ConnectorConfigurationError(
                        "connector operation returned an invalid status"
                    )
                if cancel_probe():
                    raise ConnectorCancelledError()
                safe_result = sanitize_operation_result(result)
                status = "dry_run" if request.dry_run else "succeeded"
                safe_result = replace(safe_result, status=status)
                updated = self.store.update(
                    key,
                    status=status,
                    result=safe_result,
                    checkpoint=dict(safe_result.checkpoint),
                    cursor=dict(safe_result.cursor),
                    error_category=None,
                    error_code=None,
                    expected_attempt=attempt,
                    lease_token=lease_token,
                )
                if not updated.acquired and updated.record.get("status") == "cancelled":
                    raise ConnectorCancelledError()
                if not updated.acquired:
                    raise ConnectorConfigurationError(
                        "connector run completion ownership was lost"
                    )
                return safe_result
            except Exception as error:
                classified = classify_error(error)
                updated = self.store.update(
                    key,
                    status="failed",
                    error_category=classified.category,
                    error_code=type(classified).__name__,
                    expected_attempt=attempt,
                    lease_token=lease_token,
                )
                if not updated.acquired and updated.record.get("status") == "cancelled":
                    raise ConnectorCancelledError() from error
                if not updated.acquired:
                    raise ConnectorConfigurationError(
                        "connector run failure ownership was lost"
                    ) from error
                final_attempt = attempt == first_attempt + invocation_attempts - 1
                if not classified.retryable or final_attempt:
                    raise classified from error
                self.sleeper(min(0.1 * (2 ** (attempt - 1)), 1.0))
        raise ConnectorConfigurationError("connector run did not complete")

    def execute(self, connector_id, version, request):
        connector = self.registry.resolve(connector_id, version, request.operation)
        config = self.registry.validate(connector_id, version, request.config)
        request = replace(request, config=config)
        key = deterministic_idempotency_key(connector_id, version, request)
        client_hint_hash = (
            hashlib.sha256(str(request.idempotency_key).encode()).hexdigest()
            if request.idempotency_key
            else None
        )
        self._acquire_rate_limit(connector_id, request.source_uid)
        record, created = self.store.claim(
            key,
            {
                "idempotency_key": key,
                "request_hash": key,
                "client_hint_hash": client_hint_hash,
                "connector_id": connector_id,
                "connector_version": version,
                "source_uid": request.source_uid,
                "operation": request.operation,
                "config": dict(request.config),
                "scope": dict(request.scope),
                "status": "running",
                "attempt_count": 0,
                "checkpoint": dict(request.checkpoint),
                "cursor": dict(request.cursor),
                "dry_run": request.dry_run,
                "principal_uid": request.principal_uid,
                "business_domain_uid": request.business_domain_uid,
                "environment": request.environment,
                "process_key": request.process_key,
                "source_binding_uid": request.source_binding_uid,
                "source_binding_version": request.source_binding_version,
                "cancel_requested": False,
            },
        )
        if not created:
            if record["status"] in {"succeeded", "dry_run"}:
                return record.get("result") or OperationResult(
                    cursor=record.get("cursor") or {},
                    checkpoint=record.get("checkpoint") or {},
                    evidence={"idempotent_replay": True},
                    status=record["status"],
                )
            if record["status"] == "running":
                raise ConnectorConfigurationError(
                    "idempotent operation is already running"
                )
        if record.get("status") == "cancelled":
            raise ConnectorCancelledError()
        operation = getattr(connector, request.operation)
        if request.dry_run:
            health = connector.health(config)

            def validation_only(_request):
                return OperationResult(
                    evidence={
                        "validation_only": True,
                        "health_status": health.status,
                        "network_requested": False,
                    },
                )

            operation = validation_only
        return self._run_operation(key, record, request, operation)

    def cancel(self, key):
        record = self.store.get(key)
        if not record:
            raise ConnectorConfigurationError("connector run was not found")
        if record["status"] in {"succeeded", "failed", "cancelled"}:
            raise ConnectorConfigurationError("connector run cannot be cancelled")
        key = record.get("idempotency_key", key)
        connector = self.registry.resolve(
            record["connector_id"], record["connector_version"], "cancel"
        )
        from app.core.connectors.sdk import OperationRequest

        request = OperationRequest(
            source_uid=record["source_uid"],
            operation="cancel",
            config=record.get("config") or {},
            scope=record.get("scope") or {},
            cursor=record.get("cursor") or {},
            checkpoint=record.get("checkpoint") or {},
            idempotency_key=key,
            dry_run=bool(record.get("dry_run")),
            principal_uid=record.get("principal_uid"),
            business_domain_uid=record.get("business_domain_uid"),
            environment=record.get("environment"),
            process_key=record.get("process_key"),
            source_binding_uid=record.get("source_binding_uid"),
            source_binding_version=record.get("source_binding_version"),
            run_key=key,
        )
        if hasattr(self.store, "cancel"):
            cancelled = self.store.cancel(key)
        else:
            outcome = self.store.update(
                key, status="cancelled", cancel_requested=True
            )
            if not outcome.acquired:
                raise ConnectorConfigurationError("connector run cannot be cancelled")
            cancelled = outcome.record
        def cancel_probe():
            return bool(
                hasattr(self.store, "is_cancel_requested")
                and self.store.is_cancel_requested(key)
            )
        request = replace(request, cancel_probe=cancel_probe)
        if not request.dry_run:
            connector.cancel(request)
        return cancelled

    def resume(self, key):
        record = self.store.get(key)
        if not record or record["status"] not in {"failed", "cancelled"}:
            raise ConnectorConfigurationError("connector run cannot be resumed")
        key = record.get("idempotency_key", key)
        connector = self.registry.resolve(
            record["connector_id"], record["connector_version"], "resume"
        )
        config = self.registry.validate(
            record["connector_id"],
            record["connector_version"],
            record.get("config") or {},
        )
        self._acquire_rate_limit(record["connector_id"], record["source_uid"])
        from app.core.connectors.sdk import OperationRequest

        request = OperationRequest(
            source_uid=record["source_uid"],
            operation="resume",
            config=config,
            scope=record.get("scope") or {},
            cursor=record.get("cursor") or {},
            checkpoint=record.get("checkpoint") or {},
            idempotency_key=key,
            dry_run=bool(record.get("dry_run")),
            principal_uid=record.get("principal_uid"),
            business_domain_uid=record.get("business_domain_uid"),
            environment=record.get("environment"),
            process_key=record.get("process_key"),
            source_binding_uid=record.get("source_binding_uid"),
            source_binding_version=record.get("source_binding_version"),
        )
        resumable = self.store.update(
            key,
            status="resumable",
            resumed_from=record.get("uid"),
            cancel_requested=False,
        )
        if not resumable.acquired:
            raise ConnectorConfigurationError("connector run cannot be resumed")
        operation = connector.resume
        if request.dry_run:
            health = connector.health(config)

            def validation_only(_request):
                return OperationResult(
                    evidence={
                        "validation_only": True,
                        "health_status": health.status,
                        "network_requested": False,
                        "resumed": True,
                    }
                )

            operation = validation_only
        return self._run_operation(key, resumable.record, request, operation)


__all__ = [
    "ConnectorRuntime",
    "InMemoryRunStore",
    "SlidingWindowLimiter",
    "MAX_TOTAL_ATTEMPTS",
    "deterministic_idempotency_key",
    "snapshot_diff",
    "redact_evidence",
]