Scenario 1: Group risk control needs to integrate Subsidiary A's transaction data with Subsidiary B's credit data
Scenario 2: Supply chain optimization needs to federate upstream supplier and downstream distributor inventory data
Scenario 3: Cross-border compliance needs to federate KYC/AML data across jurisdictions, but data cannot leave the country

from dataclasses import dataclass, field
from enum import Enum
from typing import Any

class FederationLevel(Enum):
    INTRA_CLUSTER = "intra_cluster"   # Within cluster: cross-World federation
    INTRA_ORG = "intra_org"           # Within organization: cross-cluster federation
    CROSS_ORG = "cross_org"           # Cross-organization: cross-institution federation

@dataclass
class FederationNode:
    """A node in the federation network."""
    node_id: str
    name: str
    endpoint: str
    level: FederationLevel
    trust_level: str                   # full | limited | minimal
    capabilities: list[str] = field(default_factory=list)
    shared_object_types: list[str] = field(default_factory=list)
    access_policies: dict[str, Any] = field(default_factory=dict)
    metadata: dict[str, Any] = field(default_factory=dict)

@dataclass
class FederationRegistry:
    """Registry of federation nodes."""
    local_node_id: str
    nodes: dict[str, FederationNode] = field(default_factory=dict)

    def register(self, node: FederationNode) -> None:
        self.nodes[node.node_id] = node

    def get_nodes_for_type(self, object_type: str) -> list[FederationNode]:
        return [
            n for n in self.nodes.values()
            if object_type in n.shared_object_types
        ]

class FederationCatalog:
    """Catalog service for federated Ontology schemas."""

    async def publish_schema(
        self, object_type: str, schema: dict, visibility: str = "federation"
    ) -> None:
        """Publish a schema to the federation catalog."""
        await self._catalog_store.save({
            "node_id": self._local_node_id,
            "object_type": object_type,
            "schema": schema,
            "visibility": visibility,
            "published_at": datetime.utcnow().isoformat(),
            "version": schema.get("version", 1),
        })

        for node in self._registry.nodes.values():
            if node.node_id != self._local_node_id:
                await self._notify_schema_update(node, object_type, schema)

    async def discover_schema(self, object_type: str) -> list[dict]:
        """Discover all nodes that provide a specific ObjectType."""
        results = []
        for node in self._registry.nodes.values():
            try:
                schema = await self._fetch_remote_schema(node, object_type)
                if schema:
                    results.append({
                        "node_id": node.node_id,
                        "node_name": node.name,
                        "schema": schema,
                        "trust_level": node.trust_level,
                    })
            except Exception:
                continue
        return results

    async def resolve_schema_conflicts(
        self, object_type: str, schemas: list[dict]
    ) -> dict:
        """Resolve schema conflicts across federation nodes."""
        if not schemas:
            raise ValueError(f"No schemas found for {object_type}")
        base = schemas[0]["schema"]
        for s in schemas[1:]:
            base = self._merge_schemas(base, s["schema"])
        return base

    def _merge_schemas(self, schema_a: dict, schema_b: dict) -> dict:
        merged = dict(schema_a)
        for field_name, field_def in schema_b.get("properties", {}).items():
            if field_name not in merged.get("properties", {}):
                merged.setdefault("properties", {})[field_name] = field_def
        return merged

@dataclass
class FederatedQueryPlan:
    query_id: str
    original_query: str
    sub_queries: list["SubQuery"]
    aggregation: dict[str, Any]
    estimated_cost: float

@dataclass
class SubQuery:
    node_id: str
    query: str
    object_type: str
    filters: list[dict]
    projections: list[str]
    estimated_rows: int

class FederatedQueryPlanner:
    """Plan federated queries across nodes."""

    async def plan(self, query: str) -> FederatedQueryPlan:
        parsed = self._parser.parse(query)
        object_types = parsed.referenced_types

        sub_queries = []
        for obj_type in object_types:
            nodes = self._registry.get_nodes_for_type(obj_type)
            for node in nodes:
                pushable_filters = self._extract_pushable_filters(
                    parsed.filters, obj_type, node
                )
                sub_queries.append(SubQuery(
                    node_id=node.node_id,
                    query=self._build_sub_query(obj_type, pushable_filters, parsed.projections),
                    object_type=obj_type,
                    filters=pushable_filters,
                    projections=self._get_required_projections(parsed, obj_type),
                    estimated_rows=await self._estimate_rows(node, obj_type, pushable_filters),
                ))

        return FederatedQueryPlan(
            query_id=generate_id(),
            original_query=query,
            sub_queries=sub_queries,
            aggregation=self._plan_aggregation(parsed),
            estimated_cost=sum(sq.estimated_rows for sq in sub_queries),
        )

class FederatedQueryExecutor:
    """Execute federated queries across nodes."""

    async def execute(self, plan: FederatedQueryPlan) -> list[dict]:
        import asyncio

        tasks = [self._execute_sub_query(sq) for sq in plan.sub_queries]
        results = await asyncio.gather(*tasks, return_exceptions=True)

        valid_results = []
        for i, result in enumerate(results):
            if isinstance(result, Exception):
                await self._log_sub_query_failure(
                    plan.sub_queries[i].node_id, result
                )
            else:
                valid_results.extend(result)

        return self._aggregate_results(valid_results, plan.aggregation)

    async def _execute_sub_query(self, sub_query: SubQuery) -> list[dict]:
        node = self._registry.nodes[sub_query.node_id]

        if node.node_id == self._registry.local_node_id:
            return await self._local_executor.execute(sub_query.query)

        async with grpc_channel(node.endpoint) as channel:
            stub = FederatedQueryServiceStub(channel)
            response = await stub.ExecuteQuery(
                FederatedQueryRequest(
                    query=sub_query.query,
                    object_type=sub_query.object_type,
                    filters=sub_query.filters,
                    projections=sub_query.projections,
                    caller_node_id=self._registry.local_node_id,
                )
            )
            return [dict(row) for row in response.rows]

    def _aggregate_results(self, results: list[dict], aggregation: dict) -> list[dict]:
        if aggregation.get("type") == "union":
            return results
        elif aggregation.get("type") == "join":
            return self._join_results(results, aggregation["join_key"])
        elif aggregation.get("type") == "group_by":
            return self._group_results(
                results, aggregation["group_key"], aggregation["agg_func"]
            )
        return results

@dataclass
class FederationAccessPolicy:
    policy_id: str
    object_type: str
    allowed_nodes: list[str] | None = None
    denied_nodes: list[str] = field(default_factory=list)
    allowed_fields: list[str] | None = None
    denied_fields: list[str] = field(default_factory=list)
    row_filter: str | None = None
    data_masking: dict[str, str] = field(default_factory=dict)
    max_rows: int | None = None
    audit_required: bool = True

class FederationAccessController:
    """Enforce federation access policies."""

    async def check_access(
        self, caller_node_id: str, object_type: str, requested_fields: list[str]
    ) -> dict:
        policy = await self._get_policy(object_type)

        if not policy:
            return {"allowed": False, "reason": "No federation policy defined"}

        if policy.denied_nodes and caller_node_id in policy.denied_nodes:
            return {"allowed": False, "reason": "Node is denied"}

        if policy.allowed_nodes and caller_node_id not in policy.allowed_nodes:
            return {"allowed": False, "reason": "Node is not in allowed list"}

        denied_fields = set(requested_fields) & set(policy.denied_fields)
        if denied_fields:
            return {"allowed": False, "reason": f"Denied fields: {denied_fields}"}

        masking_rules = {}
        for field_name in requested_fields:
            if field_name in policy.data_masking:
                masking_rules[field_name] = policy.data_masking[field_name]

        return {
            "allowed": True,
            "masking_rules": masking_rules,
            "row_filter": policy.row_filter,
            "max_rows": policy.max_rows,
        }

    async def apply_masking(
        self, data: list[dict], masking_rules: dict[str, str]
    ) -> list[dict]:
        masked = []
        for row in data:
            masked_row = dict(row)
            for field_name, mask_type in masking_rules.items():
                if field_name in masked_row:
                    masked_row[field_name] = self._mask_value(
                        masked_row[field_name], mask_type
                    )
            masked.append(masked_row)
        return masked

    def _mask_value(self, value: Any, mask_type: str) -> Any:
        if mask_type == "hash":
            import hashlib
            return hashlib.sha256(str(value).encode()).hexdigest()[:16]
        elif mask_type == "partial":
            s = str(value)
            return s[:2] + "*" * (len(s) - 4) + s[-2:] if len(s) > 4 else "****"
        elif mask_type == "null":
            return None
        elif mask_type == "category":
            return f"[MASKED:{type(value).__name__}]"
        return value

class CrossBorderFederationPolicy:
    """Ensure compliance with cross-border data regulations."""

    async def enforce(
        self, query_plan: FederatedQueryPlan, data_residency_rules: dict[str, str]
    ) -> FederatedQueryPlan:
        modified_sub_queries = []

        for sq in query_plan.sub_queries:
            node = self._registry.nodes[sq.node_id]
            node_region = node.metadata.get("region", "unknown")
            caller_region = self._local_region

            if node_region != caller_region:
                if not self._is_aggregation_only(sq):
                    sq = self._convert_to_aggregation(sq)
                sq.filters.append({
                    "type": "masking",
                    "policy": "cross_border_default",
                })

            modified_sub_queries.append(sq)

        query_plan.sub_queries = modified_sub_queries
        return query_plan

class FederationSyncService:
    """Event-driven synchronization across federation nodes."""

    async def publish_change(
        self, object_type: str, change_type: str, data: dict
    ) -> None:
        interested_nodes = self._registry.get_nodes_for_type(object_type)

        for node in interested_nodes:
            if node.node_id == self._registry.local_node_id:
                continue

            access = await self._access_controller.check_access(
                node.node_id, object_type, list(data.keys())
            )
            if not access["allowed"]:
                continue

            if access.get("masking_rules"):
                data = (await self._access_controller.apply_masking(
                    [data], access["masking_rules"]
                ))[0]

            await self._event_bus.publish(
                topic=f"federation.{node.node_id}.{object_type}",
                event={
                    "source_node": self._registry.local_node_id,
                    "object_type": object_type,
                    "change_type": change_type,
                    "data": data,
                    "timestamp": datetime.utcnow().isoformat(),
                },
            )

# Cross-subsidiary risk control data query
federated_query = """
SELECT
    t.customer_id,
    t.transaction_amount,
    c.credit_score,
    r.risk_level
FROM
    SubsidiaryA.Transaction t
    JOIN SubsidiaryB.CreditRecord c ON t.customer_id = c.customer_id
    JOIN Headquarters.RiskAssessment r ON t.customer_id = r.customer_id
WHERE
    t.transaction_amount > 100000
    AND c.credit_score < 600
"""
# The federated query engine automatically decomposes into three sub-queries sent to three nodes

# Cross-organization supply chain inventory federated query
supply_chain_query = """
SELECT
    s.supplier_id,
    s.product_id,
    s.available_qty,
    d.demand_forecast,
    d.demand_forecast - s.available_qty AS gap
FROM
    Supplier.Inventory s
    JOIN Internal.DemandForecast d ON s.product_id = d.product_id
WHERE
    d.demand_forecast > s.available_qty
ORDER BY gap DESC
"""