Improve generators

Author: Alessandro De Maria

kubernetes/argocd.py

@@ -10,10 +10,10 @@
 
 class ArgoCDApplicationConfigSpec(KubernetesResourceSpec):
     project: str = "default"
-    destination: Dict = None
+    destination: dict = None
     source: dict
     sync_policy: dict = None
-    ignore_differences: dict = None
+    ignore_differences: list[dict] = None
 
 
 class ArgoCDApplication(KubernetesResource):
@@ -27,7 +27,7 @@ def body(self):
         self.root.spec.project = self.config.project
         destination = self.config.destination
         self.root.spec.destination.name = self.cluster.display_name or destination.name
-        self.root.spec.destination.namespace = destination.namespace
+        self.root.spec.destination.namespace = destination.get("namespace")
         self.root.spec.source = self.config.source
         self.root.spec.syncPolicy = self.config.sync_policy
 
@@ -44,7 +44,6 @@ def body(self):
     path="generators.argocd.applications",
     global_generator=True,
     activation_path="argocd.app_of_apps",
-    apply_patches=["generators.argocd.defaults.application"],
 )
 class GenArgoCDApplication(kgenlib.BaseStore):
     def body(self):
@@ -53,6 +52,11 @@ def body(self):
         name = config.get("name", self.name)
         enabled = config.get("enabled", True)
         clusters = self.inventory.parameters.get("clusters", {})
+        single_cluster = config.get("single_cluster", False)
+        if single_cluster:
+            cluster = self.inventory.parameters.get("cluster", {})
+            cluster_id = cluster.get("user", None)
+            clusters = {cluster_id: cluster}
 
         if enabled:
             for cluster in clusters.keys():
@@ -72,40 +76,6 @@ def body(self):
                 self.add(argo_application)
 
 
-@kgenlib.register_generator(
-    path="isoflow.graphs",
-    global_generator=True,
-    activation_path="argocd.app_of_apps",
-)
-class GenArgoCDApplicationPerResource(kgenlib.BaseStore):
-    def body(self):
-        graph_config = self.config
-        argocd_application_config = graph_config.get("argocd_application_config", {})
-
-        # Applies generator defaults to the config of the argocd application
-        kgenlib.patch_config(
-            argocd_application_config,
-            self.inventory,
-            "parameters.generators.argocd.defaults.application",
-        )
-
-        graph_name = graph_config.get("name", self.name)
-        argocd_application_config.setdefault("source", {}).setdefault("directory", {})[
-            "include"
-        ] = f"{graph_name}.yml"
-        argocd_application_config["source"].pop("plugin", None)
-
-        namespace = argocd_application_config["destination"]["namespace"]
-        cluster = argocd_application_config["destination"]["name"]
-
-        graph_name = f"{graph_name}.{namespace}.{cluster}"
-
-        argo_application = ArgoCDApplication(
-            name=graph_name, config=argocd_application_config
-        )
-        self.add(argo_application)
-
-
 class ArgoCDProjectConfigSpec(KubernetesResourceSpec):
     source_repos: list = []
     destinations: list = []
@@ -130,7 +100,8 @@ def body(self):
 
 @kgenlib.register_generator(
     path="generators.argocd.projects",
-    apply_patches=["generators.argocd.defaults.project"],
+    global_generator=True,
+    activation_path="argocd.app_of_apps",
 )
 class GenArgoCDProject(kgenlib.BaseStore):
     def body(self):

kubernetes/autoscaling.py

@@ -2,7 +2,7 @@
 
 logger = logging.getLogger(__name__)
 
-from .common import KubernetesResource
+from .common import KubernetesResource, kgenlib
 
 
 class KedaScaledObject(KubernetesResource):
@@ -17,6 +17,16 @@ def body(self):
         self.root.spec.scaleTargetRef.kind = workload.root.kind
         self.root.spec.scaleTargetRef.apiVersion = workload.root.apiVersion
         self.root.spec.update(config.keda_scaled_object)
+        if self.root.spec.maxReplicaCount == 0:
+            # keda supports pausing autoscaling
+            # but doesn't support setting maxReplicaCount to 0
+            self.root.metadata.annotations.update(
+                {
+                    "autoscaling.keda.sh/paused-replicas": "0",
+                    "autoscaling.keda.sh/paused": "true",
+                }
+            )
+            self.root.spec.maxReplicaCount = 1
 
         # remove replica from workload because HPA is managing it
         workload.root.spec.pop("replicas")
@@ -83,3 +93,19 @@ def body(self):
 
         # remove replica from workload because HPA is managing it
         workload.root.spec.pop("replicas")
+
+
+@kgenlib.register_generator(path="generators.kubernetes.vpa")
+class VerticalPodAutoscalerGenerator(kgenlib.BaseStore):
+    def body(self):
+        name = self.config.get("name", self.name)
+        self.config.vpa.update_mode = self.config.update_mode
+        self.config.vpa.resource_policy = self.config.resource_policy
+
+        workload = KubernetesResource(
+            name=name, kind=self.config.kind, api_version=self.config.api_version
+        )
+
+        self.add(
+            VerticalPodAutoscaler(name=name, config=self.config, workload=workload)
+        )

kubernetes/common.py

@@ -2,10 +2,11 @@
 
 logger = logging.getLogger(__name__)
 from enum import StrEnum, auto
-from typing import Any, Dict, List, Optional
+from typing import Annotated, Any, Dict, List, Optional, Union
 
 from kapitan.inputs.kadet import BaseObj, load_from_search_paths
-from pydantic import DirectoryPath, Field, FilePath
+from pydantic import ConfigDict, DirectoryPath, Field, FilePath, model_validator
+from pydantic.alias_generators import to_camel
 
 kgenlib = load_from_search_paths("kgenlib")
 
@@ -68,6 +69,8 @@ def body(self):
         self.root.kind = self.kind
 
         self.root.metadata.name = self.rendered_name or self.name
+        self.root.metadata.labels
+        self.root.metadata.annotations
         self.add_label("name", self.name)
         if self.config:
             self.add_labels(self.config.labels)
@@ -141,13 +144,17 @@ def body(self):
     def set_namespace(self, namespace: str):
         self.root.metadata.namespace = namespace
 
+    def remove_namespace(self):
+        self.root.metadata.pop("namespace")
+
 
 class WorkloadTypes(StrEnum):
     DEPLOYMENT = auto()
     STATEFULSET = auto()
     DAEMONSET = auto()
     JOB = auto()
     CRONJOB = auto()
+    CLOUD_RUN_SERVICE = auto()
 
 
 class RestartPolicy(StrEnum):
@@ -156,6 +163,12 @@ class RestartPolicy(StrEnum):
     NEVER = "Never"
 
 
+class ConcurrentPolicy(StrEnum):
+    ALLOW = "Allow"
+    FORBID = "Forbid"
+    REPLACE = "Replace"
+
+
 class ImagePullPolicy(StrEnum):
     ALWAYS = "Always"
     IF_NOT_PRESENT = "IfNotPresent"
@@ -286,7 +299,9 @@ class ServiceAccountConfigSpec(KubernetesResourceSpec):
     namespace: Optional[str] = None
     rendered_name: Optional[str] = None
     name: Optional[str] = None
-    roles: Optional[List[Dict[str, Any]]] = None
+    roles: Optional[
+        list[dict[str, list[str]]] | dict[str, dict[str, list[str]] | list[str]]
+    ] = None
 
 
 class ContainerSpec(kgenlib.BaseModel):
@@ -298,7 +313,7 @@ class ContainerSpec(kgenlib.BaseModel):
     healthcheck: Optional[HealthCheckConfigSpec] = None
     image: str = None
     image_pull_policy: Optional[ImagePullPolicy] = ImagePullPolicy.IF_NOT_PRESENT
-    lifecycle: dict = {}
+    lifecycle: Optional[dict] = None
     pod_annotations: dict = {}
     pod_labels: dict = {}
     ports: Dict[str, ContainerPortSpec] = {}
@@ -308,6 +323,10 @@ class ContainerSpec(kgenlib.BaseModel):
     volume_mounts: dict = {}
 
 
+class InitContainerSpec(ContainerSpec):
+    sidecar: Optional[bool] = None
+
+
 class ServiceTypes(StrEnum):
     EXTERNAL_NAME = "ExternalName"
     CLUSTER_IP = "ClusterIP"
@@ -322,7 +341,7 @@ class SessionAffinity(StrEnum):
 
 class RoleBindingConfigSpec(KubernetesResourceSpec):
     roleRef: Optional[Dict[str, Any]] = None
-    subject: Optional[List[Dict[str, Any]]] = None
+    subjects: Optional[List[Dict[str, Any]]] = None
 
 
 class ServiceConfigSpec(KubernetesResourceSpec):
@@ -355,8 +374,9 @@ class NetworkPolicySpec(KubernetesResourceSpec):
 
 class WorkloadConfigSpec(KubernetesResourceSpec, ContainerSpec):
     type: Optional[WorkloadTypes] = WorkloadTypes.DEPLOYMENT
+    namespace: str | None = None
     schedule: Optional[str] = None
-    additional_containers: Optional[Dict[str, ContainerSpec]] = {}
+    additional_containers: Optional[Dict[str, Union[ContainerSpec, None]]] = {}
     additional_services: Optional[Dict[str, ServiceConfigSpec]] = {}
     annotations: dict = {}
     application: Optional[str] = None
@@ -371,7 +391,7 @@ class WorkloadConfigSpec(KubernetesResourceSpec, ContainerSpec):
     host_pid: Optional[bool] = None
     hpa: dict = {}
     image_pull_secrets: list = []
-    init_containers: Optional[Dict[str, ContainerSpec]] = {}
+    init_containers: Optional[Dict[str, Union[InitContainerSpec, None]]] = {}
     istio_policy: dict = {}
     keda_scaled_object: dict = {}
     labels: Dict[str, str] = {}
@@ -399,16 +419,25 @@ class WorkloadConfigSpec(KubernetesResourceSpec, ContainerSpec):
     workload_security_context: dict = {}
 
 
+class CloudRunServiceConfigSpec(WorkloadConfigSpec):
+    type: Optional[WorkloadTypes] = WorkloadTypes.CLOUD_RUN_SERVICE
+
+
 class DeploymentConfigSpec(WorkloadConfigSpec):
     type: Optional[WorkloadTypes] = WorkloadTypes.DEPLOYMENT
     update_strategy: Optional[dict] = {}
     strategy: Optional[dict] = {}
 
 
+class PodManagementPolicy(StrEnum):
+    ORDERED_READY = "OrderedReady"
+    PARALLEL = "Parallel"
+
+
 class StatefulSetConfigSpec(WorkloadConfigSpec):
     type: WorkloadTypes = WorkloadTypes.STATEFULSET
+    pod_management_policy: str = PodManagementPolicy.ORDERED_READY
     update_strategy: dict = {}
-    strategy: dict = {}
 
 
 class DaemonSetConfigSpec(WorkloadConfigSpec):
@@ -426,3 +455,94 @@ class JobConfigSpec(WorkloadConfigSpec):
 class CronJobConfigSpec(JobConfigSpec):
     type: WorkloadTypes = WorkloadTypes.CRONJOB
     schedule: str
+    concurrency_policy: Optional[ConcurrentPolicy] = ConcurrentPolicy.ALLOW
+
+
+ContainerProbeSpecTypes = Annotated[
+    Union[ContainerEXECProbeSpec, ContainerHTTPProbeSpec, ContainerTCPProbeSpec],
+    Field(discriminator="type"),
+]
+
+
+class ContainerProbes(kgenlib.BaseModel):
+    model_config = ConfigDict(
+        alias_generator=to_camel,
+        populate_by_name=True,
+        from_attributes=True,
+        extra="ignore",
+    )
+    type: ProbeTypes = Field(exclude=True)
+    initial_delay_seconds: int = 0
+    period_seconds: int = 10
+    timeout_seconds: int = 1
+    success_threshold: int = 1
+    failure_threshold: int = 3
+
+    # Define a class method for creating probes from data
+    @classmethod
+    def from_spec(cls, spec: ContainerProbeSpecTypes) -> Union["ContainerProbes", None]:
+        probe = None
+        if not spec or not spec.enabled:
+            return probe
+        probe_type = spec.type
+        if probe_type == ProbeTypes.TCP:
+            probe = TCPProbe.model_validate(spec)
+        elif probe_type == ProbeTypes.HTTP:
+            probe = HTTPProbe.model_validate(spec)
+        elif probe_type == ProbeTypes.EXEC:
+            probe = ExecProbe.model_validate(spec)
+        else:
+            raise ValueError(f"Invalid probe type: {probe_type}")
+        return probe.model_dump(by_alias=True)
+
+
+class HttpGet(kgenlib.BaseModel):
+    model_config = ConfigDict(
+        alias_generator=to_camel,
+        populate_by_name=True,
+        from_attributes=True,
+    )
+    path: str = "/"
+    port: str | int = 80
+    httpHeaders: Optional[List[dict]] = None
+    scheme: Optional[ProbeSchemeSpec] = ProbeSchemeSpec.HTTP
+
+
+class HTTPProbe(ContainerProbes):
+    httpGet: Optional[HttpGet] = None
+    port: str | int = Field(exclude=True)
+    path: str = Field(exclude=True)
+    scheme: Optional[ProbeSchemeSpec] = Field(exclude=True)
+    httpHeaders: Optional[List[dict]] = Field(exclude=True)
+
+    # Use a validator to handle the 'port' mapping
+    @model_validator(mode="after")
+    def map_port_to_http_config(self):
+        self.httpGet = HttpGet.model_validate(self)
+        return self
+
+
+class TCPProbe(ContainerProbes):
+    port: str | int = Field(exclude=True)
+    tcpSocket: Optional[dict] = None
+
+    # Use a validator to handle the 'port' mapping
+    @model_validator(mode="after")
+    def map_port_to_tcp_socket(self):
+        self.tcpSocket = {"port": self.port}
+        return self
+
+
+class ExecProbe(ContainerProbes):
+    command: List = Field(exclude=True)
+    exec: Optional[dict] = None
+
+    @model_validator(mode="after")
+    def map_command_to_exec_socket(self):
+        self.exec = {"command": self.command}
+        return self
+
+
+ContainerProbeTypes = Annotated[
+    Union[HTTPProbe, TCPProbe, ExecProbe], Field(discriminator="type")
+]

kubernetes/files.py

@@ -0,0 +1,17 @@
+import logging
+
+from .common import kgenlib
+
+logger = logging.getLogger(__name__)
+
+
+@kgenlib.register_generator(path="generators.kubernetes.raw")
+class RawManifestFilesGenerator(kgenlib.BaseStore):
+    def body(self):
+        for file in self.config.files:
+            self.add(kgenlib.BaseStore.from_yaml_file(file))
+        if self.config.filename:
+            [
+                setattr(content, "filename", self.config.filename)
+                for content in self.content_list
+            ]