DataOps-platform/app/core/orchestration/spec.py

from __future__ import annotations

import copy
import hashlib
import json
import re
from collections import deque
from typing import Any

from app.core.common.identifiers import ensure_governance_uid


SCHEMA_VERSION = "1.0"
NODE_TYPES = {
    "sql.query",
    "sql.execute",
    "rule.apply",
    "quality.check",
    "python",
    "http",
    "condition",
    "parallel",
    "subflow",
    "notify",
}
SECRET_KEY_NAMES = {
    "apikey",
    "authorization",
    "connectionstring",
    "credential",
    "credentials",
    "dsn",
    "password",
    "secret",
    "token",
}
WORKFLOW_ROOT_KEYS = {
    "schema_version",
    "dataflow_uid",
    "name",
    "description",
    "nodes",
    "edges",
    "parameters",
    "labels",
}
NODE_KEYS = {
    "id",
    "type",
    "data_source_uid",
    "purpose",
    "config",
    "idempotency",
}
EDGE_KEYS = {"from", "to", "condition"}
SCHEDULE_ROOT_KEYS = {
    "schema_version",
    "timezone",
    "triggers",
    "max_concurrency",
    "conflict_policy",
    "timeout_seconds",
    "retry",
    "backfill",
}
CONFLICT_POLICIES = {"skip", "queue", "cancel_previous"}
RULE_NODE_CONFIG_KEYS = {
    "component_binding_id",
    "rule_version_id",
    "execution_plan_hash",
    "provenance",
}
RULE_NODE_PROVENANCE_KEYS = {"standard_version_id", "clause_id"}


WORKFLOW_SPEC_SCHEMA = {
    "$schema": "https://json-schema.org/draft/2020-12/schema",
    "$id": "https://dataops.local/schemas/workflow-spec-1.0.json",
    "title": "DataOps WorkflowSpec",
    "type": "object",
    "additionalProperties": False,
    "required": [
        "schema_version",
        "dataflow_uid",
        "name",
        "nodes",
        "edges",
        "parameters",
    ],
    "properties": {
        "schema_version": {"const": SCHEMA_VERSION},
        "dataflow_uid": {"type": "string", "format": "uuid"},
        "name": {"type": "string", "minLength": 1, "maxLength": 200},
        "description": {"type": "string", "maxLength": 2000},
        "nodes": {
            "type": "array",
            "minItems": 1,
            "items": {"$ref": "#/$defs/node"},
        },
        "edges": {"type": "array", "items": {"$ref": "#/$defs/edge"}},
        "parameters": {"type": "object"},
        "labels": {"type": "object"},
    },
    "$defs": {
        "node": {
            "type": "object",
            "additionalProperties": False,
            "required": ["id", "type", "config"],
            "properties": {
                "id": {"type": "string", "pattern": "^[A-Za-z][A-Za-z0-9_-]{0,99}$"},
                "type": {"type": "string", "enum": sorted(NODE_TYPES)},
                "data_source_uid": {"type": "string", "format": "uuid"},
                "purpose": {"type": "string", "enum": ["read", "write"]},
                "config": {"type": "object"},
                "idempotency": {"type": "object"},
            },
        },
        "edge": {
            "type": "object",
            "additionalProperties": False,
            "required": ["from", "to"],
            "properties": {
                "from": {"type": "string"},
                "to": {"type": "string"},
                "condition": {"type": "string"},
            },
        },
    },
}


SCHEDULE_PLAN_SCHEMA = {
    "$schema": "https://json-schema.org/draft/2020-12/schema",
    "$id": "https://dataops.local/schemas/schedule-plan-1.0.json",
    "title": "DataOps SchedulePlan",
    "type": "object",
    "additionalProperties": False,
    "required": [
        "schema_version",
        "timezone",
        "triggers",
        "max_concurrency",
        "conflict_policy",
        "timeout_seconds",
        "retry",
        "backfill",
    ],
    "properties": {
        "schema_version": {"const": SCHEMA_VERSION},
        "timezone": {"type": "string"},
        "triggers": {"type": "array", "minItems": 1},
        "max_concurrency": {"type": "integer", "minimum": 1, "maximum": 100},
        "conflict_policy": {"enum": sorted(CONFLICT_POLICIES)},
        "timeout_seconds": {"type": "integer", "minimum": 1, "maximum": 86400},
        "retry": {"type": "object"},
        "backfill": {"type": "object"},
    },
}


def _closed_object(value: Any, allowed: set[str], label: str) -> dict[str, Any]:
    if not isinstance(value, dict):
        raise ValueError(f"{label} must be an object")
    unknown = sorted(set(value) - allowed)
    if unknown:
        raise ValueError(f"{label} contains unsupported fields: {', '.join(unknown)}")
    return value


def _required_string(value: Any, label: str, *, maximum: int = 200) -> str:
    if not isinstance(value, str) or not value.strip():
        raise ValueError(f"{label} is required")
    normalized = value.strip()
    if len(normalized) > maximum:
        raise ValueError(f"{label} exceeds {maximum} characters")
    return normalized


def _validate_uid(value: Any, label: str) -> str:
    try:
        return ensure_governance_uid({"uid": str(value)})
    except ValueError as exc:
        raise ValueError(f"{label} must be a valid UUIDv7") from exc


def _normalized_secret_key(value: Any) -> str:
    return re.sub(r"[^a-z0-9]", "", str(value).lower())


def _reject_secret_material(value: Any, path: str = "$") -> None:
    if isinstance(value, dict):
        for key, item in value.items():
            if _normalized_secret_key(key) in SECRET_KEY_NAMES:
                raise ValueError(f"secret material is not allowed at {path}.{key}")
            _reject_secret_material(item, f"{path}.{key}")
    elif isinstance(value, list):
        for index, item in enumerate(value):
            _reject_secret_material(item, f"{path}[{index}]")


def _validate_acyclic(node_ids: set[str], edges: list[dict[str, Any]]) -> None:
    indegree = {node_id: 0 for node_id in node_ids}
    adjacency = {node_id: [] for node_id in node_ids}
    seen_edges = set()
    for edge in edges:
        source = edge["from"]
        target = edge["to"]
        if source not in node_ids or target not in node_ids:
            raise ValueError("workflow edge references an unknown node")
        pair = (source, target)
        if pair in seen_edges:
            raise ValueError("workflow graph contains a duplicate edge")
        seen_edges.add(pair)
        adjacency[source].append(target)
        indegree[target] += 1
    queue = deque(node_id for node_id, degree in indegree.items() if degree == 0)
    visited = 0
    while queue:
        current = queue.popleft()
        visited += 1
        for target in adjacency[current]:
            indegree[target] -= 1
            if indegree[target] == 0:
                queue.append(target)
    if visited != len(node_ids):
        raise ValueError("workflow graph must be acyclic")


def _validate_rule_node_config(config: Any) -> dict[str, Any]:
    normalized = copy.deepcopy(
        _closed_object(config, RULE_NODE_CONFIG_KEYS, "rule node config")
    )
    normalized["component_binding_id"] = _validate_uid(
        normalized.get("component_binding_id"), "component_binding_id"
    )
    normalized["rule_version_id"] = _validate_uid(
        normalized.get("rule_version_id"), "rule_version_id"
    )
    plan_hash = _required_string(
        normalized.get("execution_plan_hash"),
        "execution_plan_hash",
        maximum=64,
    )
    if not re.fullmatch(r"[0-9a-f]{64}", plan_hash):
        raise ValueError("execution_plan_hash must be a sha256 hex digest")
    normalized["execution_plan_hash"] = plan_hash
    if "provenance" in normalized:
        provenance = copy.deepcopy(
            _closed_object(
                normalized["provenance"],
                RULE_NODE_PROVENANCE_KEYS,
                "rule node provenance",
            )
        )
        provenance["standard_version_id"] = _validate_uid(
            provenance.get("standard_version_id"), "standard_version_id"
        )
        provenance["clause_id"] = _required_string(
            provenance.get("clause_id"), "standard clause_id", maximum=100
        )
        normalized["provenance"] = provenance
    return normalized


def validate_workflow_spec(spec: Any) -> dict[str, Any]:
    payload = copy.deepcopy(
        _closed_object(spec, WORKFLOW_ROOT_KEYS, "workflow spec")
    )
    if payload.get("schema_version") != SCHEMA_VERSION:
        raise ValueError(f"workflow spec schema_version must be {SCHEMA_VERSION}")
    payload["dataflow_uid"] = _validate_uid(
        payload.get("dataflow_uid"), "dataflow_uid"
    )
    payload["name"] = _required_string(payload.get("name"), "workflow name")
    if "description" in payload:
        payload["description"] = _required_string(
            payload["description"], "workflow description", maximum=2000
        )
    nodes = payload.get("nodes")
    if not isinstance(nodes, list) or not nodes:
        raise ValueError("workflow nodes must be a non-empty array")
    normalized_nodes = []
    node_ids = set()
    for raw_node in nodes:
        node = copy.deepcopy(_closed_object(raw_node, NODE_KEYS, "workflow node"))
        node_id = _required_string(node.get("id"), "node id", maximum=100)
        if not re.fullmatch(r"[A-Za-z][A-Za-z0-9_-]{0,99}", node_id):
            raise ValueError("node id contains unsupported characters")
        if node_id in node_ids:
            raise ValueError("workflow node ids must be unique")
        node_ids.add(node_id)
        node_type = _required_string(node.get("type"), "node type", maximum=100)
        if node_type not in NODE_TYPES:
            raise ValueError(f"unsupported node type: {node_type}")
        config = node.get("config")
        if not isinstance(config, dict):
            raise ValueError("node config must be an object")
        node["id"] = node_id
        node["type"] = node_type
        if node_type in {"sql.query", "sql.execute"}:
            node["data_source_uid"] = _validate_uid(
                node.get("data_source_uid"), "node data_source_uid"
            )
        if node_type == "sql.query" and node.get("purpose") != "read":
            raise ValueError("sql.query node purpose must be read")
        if node_type == "sql.execute":
            if node.get("purpose") != "write":
                raise ValueError("sql.execute node purpose must be write")
            idempotency = node.get("idempotency")
            if not isinstance(idempotency, dict):
                raise ValueError("write node requires idempotency")
            strategy = idempotency.get("strategy")
            key = idempotency.get("key")
            if strategy not in {"partition_replace", "upsert", "deduplication_key"}:
                raise ValueError("write node requires idempotency strategy")
            _required_string(key, "write node idempotency key", maximum=500)
        if node_type in {"rule.apply", "quality.check"}:
            node["config"] = _validate_rule_node_config(config)
            if node_type == "quality.check":
                if node.get("purpose") != "read":
                    raise ValueError("quality.check node purpose must be read")
                if "idempotency" in node:
                    raise ValueError("quality.check node cannot define idempotency")
            else:
                if node.get("purpose") != "write":
                    raise ValueError("rule.apply node purpose must be write")
                idempotency = node.get("idempotency")
                if not isinstance(idempotency, dict):
                    raise ValueError("rule.apply node requires idempotency")
                if idempotency.get("strategy") not in {
                    "partition_replace",
                    "upsert",
                    "deduplication_key",
                }:
                    raise ValueError(
                        "rule.apply node requires idempotency strategy"
                    )
                _required_string(
                    idempotency.get("key"),
                    "rule.apply node idempotency key",
                    maximum=500,
                )
        normalized_nodes.append(node)
    payload["nodes"] = normalized_nodes

    raw_edges = payload.get("edges")
    if not isinstance(raw_edges, list):
        raise ValueError("workflow edges must be an array")
    edges = []
    for raw_edge in raw_edges:
        edge = copy.deepcopy(_closed_object(raw_edge, EDGE_KEYS, "workflow edge"))
        edge["from"] = _required_string(edge.get("from"), "edge from", maximum=100)
        edge["to"] = _required_string(edge.get("to"), "edge to", maximum=100)
        if "condition" in edge:
            edge["condition"] = _required_string(
                edge["condition"], "edge condition", maximum=500
            )
        edges.append(edge)
    _validate_acyclic(node_ids, edges)
    payload["edges"] = edges

    parameters = payload.get("parameters")
    if not isinstance(parameters, dict):
        raise ValueError("workflow parameters must be an object")
    if "labels" in payload and not isinstance(payload["labels"], dict):
        raise ValueError("workflow labels must be an object")
    _reject_secret_material(payload)
    return payload


def workflow_spec_hash(spec: Any) -> str:
    canonical = json.dumps(
        validate_workflow_spec(spec),
        sort_keys=True,
        separators=(",", ":"),
        ensure_ascii=False,
    )
    return hashlib.sha256(canonical.encode("utf-8")).hexdigest()


def _positive_bounded_integer(
    value: Any, label: str, *, minimum: int, maximum: int
) -> int:
    if isinstance(value, bool) or not isinstance(value, int):
        raise ValueError(f"{label} must be an integer")
    if value < minimum or value > maximum:
        raise ValueError(f"{label} must be between {minimum} and {maximum}")
    return value


def validate_schedule_plan(plan: Any) -> dict[str, Any]:
    payload = copy.deepcopy(
        _closed_object(plan, SCHEDULE_ROOT_KEYS, "schedule plan")
    )
    if payload.get("schema_version") != SCHEMA_VERSION:
        raise ValueError(f"schedule plan schema_version must be {SCHEMA_VERSION}")
    payload["timezone"] = _required_string(
        payload.get("timezone"), "schedule timezone", maximum=100
    )
    triggers = payload.get("triggers")
    if not isinstance(triggers, list) or not triggers:
        raise ValueError("schedule triggers must be a non-empty array")
    normalized_triggers = []
    for raw_trigger in triggers:
        trigger = _closed_object(
            copy.deepcopy(raw_trigger),
            {"type", "expression", "at", "event_type"},
            "schedule trigger",
        )
        trigger_type = trigger.get("type")
        if trigger_type not in {"manual", "cron", "at", "event"}:
            raise ValueError("unsupported schedule trigger type")
        if trigger_type == "cron":
            expression = _required_string(
                trigger.get("expression"), "cron expression", maximum=200
            )
            if len(expression.split()) not in {5, 6, 7}:
                raise ValueError("cron expression must have 5, 6, or 7 fields")
        if trigger_type == "at":
            _required_string(trigger.get("at"), "at trigger timestamp", maximum=100)
        if trigger_type == "event":
            _required_string(
                trigger.get("event_type"), "event trigger type", maximum=200
            )
        normalized_triggers.append(trigger)
    payload["triggers"] = normalized_triggers
    payload["max_concurrency"] = _positive_bounded_integer(
        payload.get("max_concurrency"),
        "max_concurrency",
        minimum=1,
        maximum=100,
    )
    if payload.get("conflict_policy") not in CONFLICT_POLICIES:
        raise ValueError("unsupported conflict_policy")
    payload["timeout_seconds"] = _positive_bounded_integer(
        payload.get("timeout_seconds"),
        "timeout_seconds",
        minimum=1,
        maximum=86400,
    )
    retry = _closed_object(
        payload.get("retry"), {"max_attempts", "delay_seconds"}, "retry"
    )
    retry["max_attempts"] = _positive_bounded_integer(
        retry.get("max_attempts"), "retry max_attempts", minimum=1, maximum=10
    )
    retry["delay_seconds"] = _positive_bounded_integer(
        retry.get("delay_seconds"),
        "retry delay_seconds",
        minimum=0,
        maximum=86400,
    )
    backfill = _closed_object(
        payload.get("backfill"), {"max_days", "max_runs"}, "backfill"
    )
    backfill["max_days"] = _positive_bounded_integer(
        backfill.get("max_days"), "backfill max_days", minimum=1, maximum=90
    )
    backfill["max_runs"] = _positive_bounded_integer(
        backfill.get("max_runs"), "backfill max_runs", minimum=1, maximum=1000
    )
    _reject_secret_material(payload)
    return payload