Source code for openfactory.schemas.infra

"""
Infrastructure Configuration Schemas for OpenFactory

This module defines Pydantic models used to validate and process infrastructure
configuration files in OpenFactory. These configurations typically describe
networking, nodes, volumes, and IP address management (IPAM) settings for
Docker Swarm-based deployments.

Key Components:
---------------
- **Volumes**: Defines named Docker volumes and associated driver options.
- **Nodes**: Represents Docker Swarm managers and workers with IP and labels.
- **Networks**: Includes support for IPAM configuration and multiple network definitions.
- **InfrastructureSchema**: Top-level model combining nodes, volumes, and networks.
- **get_infrastructure_from_config_file**: Utility function to load and validate YAML-based infrastructure config.

Validation Features:
--------------------
- Ensures that manager and worker node IP addresses are unique across the cluster.
- Uses Pydantic aliases for compatibility with external config file formats.
- Provides rich error messages via user notification hooks.

Example YAML Structure:
-----------------------

.. code-block:: yaml

    nodes:
      managers:
        manager-1:
          ip: 192.168.1.10
      workers:
        worker-1:
          ip: 192.168.1.11

    volumes:
      my-volume:
        driver_opts:
          type: nfs
          o: addr=192.168.1.20,nolock,soft,rw
          device: ":/data"

    networks:
      openfactory-network:
        name: ofa-net
        ipam:
          config:
            - subnet: 192.168.0.0/16
              gateway: 192.168.0.1

Usage:
------

Call `get_infrastructure_from_config_file(path)` to load and validate the infrastructure config.

This module is typically used by OpenFactory's bootstrap and deployment tooling.
"""

from typing import Dict, Optional
from pydantic import BaseModel, Field, RootModel, model_validator, ValidationError
from ipaddress import IPv4Address
from openfactory.models.user_notifications import user_notify
from openfactory.config import load_yaml




class Volume(BaseModel):
    """ Volume Schema. """
    driver_opts: Optional[Dict[str, str]]





class Volumes(RootModel[Dict[str, Optional[Volume]]]):
    """ Volumes Schema. """
    pass





class Node(BaseModel):
    """ Docker Swarm Node Schema. """
    ip: IPv4Address
    labels: Optional[Dict[str, str]] = None





class Managers(RootModel[Dict[str, Node]]):
    """ Docker Swarm Managers Schema. """
    pass





class Workers(RootModel[Dict[str, Node]]):
    """ Docker Swarm Workers Schema. """
    pass





class Nodes(BaseModel):
    """ Docker Swarm Nodes Schema. """
    managers: Optional[Dict[str, Node]] = None
    workers: Optional[Dict[str, Node]] = None



    @model_validator(mode='after')
    def check_unique_ips(values: Dict) -> Dict:
        """
        Validates that IP addresses are unique across managers and workers.

        Args:
            values (Dict): Dictionary of values to validate.

        Returns:
            Dict: Validated values.

        Raises:
            ValueError: If IP addresses are not unique.
        """
        ips = []

        # Collect IP addresses from managers
        if values.managers:  # Access 'managers' directly as an attribute
            ips += [node.ip for node in values.managers.values()]

        # Collect IP addresses from workers
        if values.workers:  # Access 'workers' directly as an attribute
            ips += [node.ip for node in values.workers.values()]

        # Check if all IPs are unique
        if len(ips) != len(set(ips)):
            raise ValueError("IP addresses must be unique across managers and workers.")

        return values






class IPAMConfig(BaseModel):
    """ IPAM Configuration Schema. """
    subnet: str  # Subnet in CIDR format
    gateway: Optional[str] = None  # Optional gateway IP address
    ip_range: Optional[str] = Field(None, alias="ip_range")  # Optional IP range
    aux_addresses: Optional[Dict[str, str]] = Field(None, alias="aux_addresses")  # Optional auxiliary addresses





class IPAM(BaseModel):
    """ IPAM Schema. """
    config: list[IPAMConfig]





class Network(BaseModel):
    """ Network Schema. """
    name: Optional[str] = None
    ipam: IPAM





class Networks(BaseModel):
    """ Networks Schema. """
    openfactory_network: Network = Field(..., alias="openfactory-network")
    docker_ingress_network: Network = Field(..., alias="docker-ingress-network")





class InfrastructureSchema(BaseModel):
    """ Infrastructure Schema. """
    nodes: Nodes
    networks: Optional[Dict[str, Network]] = None
    volumes: Optional[Volumes] = None





def get_infrastructure_from_config_file(infra_yaml_config_file: str) -> Optional[Dict[str, InfrastructureSchema]]:
    """
    Load and validate infrastructure configuration from a YAML file.

    This function reads a YAML file containing infrastructure configuration,
    validates its content using the :class:`InfrastructureSchema` Pydantic model,
    and returns the parsed data as a dictionary.

    Args:
        infra_yaml_config_file (str): Path to the YAML file defining the infrastructure configuration.

    Returns:
        Optional[Dict[str, InfrastructureSchema]]: A dictionary of validated infrastructure configuration data,
                                                   or `None` if validation fails.

    Note:
        In case of validation errors, user notifications will be triggered and `None` will be returned.
    """
    # load yaml description file
    cfg = load_yaml(infra_yaml_config_file)

    # validate and create devices configuration
    try:
        infra = InfrastructureSchema(**cfg)
    except ValidationError as err:
        user_notify.fail(f"Provided YAML configuration file has invalid format\n{err}")
        return None
    except ValueError as err:
        user_notify.fail(f"Provided YAML configuration file has invalid format\n{err}")
        return None
    return infra.model_dump()