stack-orchestrator/stack_orchestrator/deploy/spec.py

# Copyright © 2022, 2023 Vulcanize

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.

# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http:#www.gnu.org/licenses/>.

import typing
from pathlib import Path

import humanfriendly

from stack_orchestrator import constants
from stack_orchestrator.util import get_yaml


class ResourceLimits:
    cpus: float | None = None
    memory: int | None = None
    storage: int | None = None

    def __init__(self, obj=None):
        if obj is None:
            obj = {}
        if "cpus" in obj:
            self.cpus = float(obj["cpus"])
        if "memory" in obj:
            self.memory = humanfriendly.parse_size(obj["memory"])
        if "storage" in obj:
            self.storage = humanfriendly.parse_size(obj["storage"])

    def __len__(self):
        return len(self.__dict__)

    def __iter__(self):
        for k in self.__dict__:
            yield k, self.__dict__[k]

    def __repr__(self):
        return str(self.__dict__)


class Resources:
    limits: ResourceLimits | None = None
    reservations: ResourceLimits | None = None

    def __init__(self, obj=None):
        if obj is None:
            obj = {}
        if "reservations" in obj:
            self.reservations = ResourceLimits(obj["reservations"])
        if "limits" in obj:
            self.limits = ResourceLimits(obj["limits"])

    def __len__(self):
        return len(self.__dict__)

    def __iter__(self):
        for k in self.__dict__:
            yield k, self.__dict__[k]

    def __repr__(self):
        return str(self.__dict__)


class Spec:
    obj: typing.Any
    file_path: Path | None

    def __init__(self, file_path: Path | None = None, obj=None) -> None:
        if obj is None:
            obj = {}
        self.file_path = file_path
        self.obj = obj

    def __getitem__(self, item):
        return self.obj[item]

    def __contains__(self, item):
        return item in self.obj

    def get(self, item, default=None):
        return self.obj.get(item, default)

    def init_from_file(self, file_path: Path):
        self.obj = get_yaml().load(open(file_path))
        self.file_path = file_path

    def get_image_registry(self):
        return self.obj.get(constants.image_registry_key)

    def get_image_registry_config(self) -> dict | None:
        """Returns registry auth config: {server, username, token-env}.

        Used for private container registries like GHCR. The token-env field
        specifies an environment variable containing the API token/PAT.

        Note: Uses 'registry-credentials' key to avoid collision with
        'image-registry' key which is for pushing images.
        """
        result: dict[str, str] | None = self.obj.get("registry-credentials")
        return result

    def get_volumes(self):
        return self.obj.get(constants.volumes_key, {})

    def get_configmaps(self):
        return self.obj.get(constants.configmaps_key, {})

    def get_container_resources(self):
        return Resources(self.obj.get(constants.resources_key, {}).get("containers", {}))

    def get_container_resources_for(self, container_name: str) -> Resources | None:
        """Look up per-container resource overrides from spec.yml.

        Checks resources.containers.<container_name> in the spec. Returns None
        if no per-container override exists (caller falls back to other sources).
        """
        containers_block = self.obj.get(constants.resources_key, {}).get("containers", {})
        if container_name in containers_block:
            entry = containers_block[container_name]
            # Only treat it as a per-container override if it's a dict with
            # reservations/limits nested inside (not a top-level global key)
            if isinstance(entry, dict) and ("reservations" in entry or "limits" in entry):
                return Resources(entry)
        return None

    def get_volume_resources(self):
        return Resources(self.obj.get(constants.resources_key, {}).get(constants.volumes_key, {}))

    def get_http_proxy(self):
        return self.obj.get(constants.network_key, {}).get(constants.http_proxy_key, [])

    def get_annotations(self):
        return self.obj.get(constants.annotations_key, {})

    def get_replicas(self):
        return self.obj.get(constants.replicas_key, 1)

    def get_node_affinities(self):
        return self.obj.get(constants.node_affinities_key, [])

    def get_node_tolerations(self):
        return self.obj.get(constants.node_tolerations_key, [])

    def get_labels(self):
        return self.obj.get(constants.labels_key, {})

    def get_privileged(self):
        return (
            "true"
            == str(self.obj.get(constants.security_key, {}).get("privileged", "false")).lower()
        )

    def get_capabilities(self):
        return self.obj.get(constants.security_key, {}).get("capabilities", [])

    def get_unlimited_memlock(self):
        return (
            "true"
            == str(
                self.obj.get(constants.security_key, {}).get(
                    constants.unlimited_memlock_key, "false"
                )
            ).lower()
        )

    def get_runtime_class(self):
        """Get runtime class name from spec, or derive from security settings.

        The runtime class determines which containerd runtime handler to use,
        allowing different pods to have different rlimit profiles (e.g., for
        unlimited RLIMIT_MEMLOCK).

        Returns:
            Runtime class name string, or None to use default runtime.
        """
        # Explicit runtime class takes precedence
        explicit = self.obj.get(constants.security_key, {}).get(constants.runtime_class_key, None)
        if explicit:
            return explicit

        # Auto-derive from unlimited-memlock setting
        if self.get_unlimited_memlock():
            return constants.high_memlock_runtime

        return None  # Use default runtime

    def get_deployment_type(self):
        return self.obj.get(constants.deploy_to_key)

    def get_acme_email(self):
        return self.obj.get(constants.network_key, {}).get(constants.acme_email_key, "")

    def is_kubernetes_deployment(self):
        return self.get_deployment_type() in [
            constants.k8s_kind_deploy_type,
            constants.k8s_deploy_type,
        ]

    def is_kind_deployment(self):
        return self.get_deployment_type() in [constants.k8s_kind_deploy_type]

    def is_docker_deployment(self):
        return self.get_deployment_type() in [constants.compose_deploy_type]