From baac38695d92fd0603f6e0cc8afc09b88d86a4f1 Mon Sep 17 00:00:00 2001 From: Nikhil Jha Date: Sat, 27 Apr 2024 16:58:50 -0700 Subject: [PATCH] feat: improved resources API --- transpire/resources/base.py | 3 + transpire/resources/deployment.py | 101 +----------- transpire/resources/podspec.py | 245 +++++++++++++++++++++++++++++ transpire/resources/pvc.py | 28 ++++ transpire/resources/statefulset.py | 78 +++++++++ 5 files changed, 362 insertions(+), 93 deletions(-) create mode 100644 transpire/resources/podspec.py create mode 100644 transpire/resources/pvc.py create mode 100644 transpire/resources/statefulset.py diff --git a/transpire/resources/base.py b/transpire/resources/base.py index d3d17c6..de3cc9f 100644 --- a/transpire/resources/base.py +++ b/transpire/resources/base.py @@ -12,6 +12,9 @@ class Resource(Generic[T]): def __init__(self): self.patches = [] + def name(self) -> str: + return self.obj.metadata.name + def patch(self, *fns: Callable[[dict], dict]) -> Self: self.patches.extend(fns) return self diff --git a/transpire/resources/deployment.py b/transpire/resources/deployment.py index b773ad5..5ded555 100644 --- a/transpire/resources/deployment.py +++ b/transpire/resources/deployment.py @@ -1,8 +1,9 @@ -from typing import List, Self, Union +from typing import Any, List, Union from kubernetes import client from transpire.resources.base import Resource +from transpire.resources.podspec import PodSpec class Deployment(Resource[client.V1Deployment]): @@ -45,98 +46,12 @@ def __init__( ), ) super().__init__() - - def _init_env(self, container_id: int) -> client.V1Container: - container = self.obj.spec.template.spec.containers[container_id] - - if container.env_from is None: - container.env_from = [] - if container.env is None: - container.env = [] - - return container - - def with_configmap_env( - self, name: str, *, mapping: dict[str, str] | None = None, container_id: int = 0 - ) -> Self: - container = self._init_env(container_id) - if mapping is None: - container.env_from.append( - client.V1EnvFromSource( - config_map_ref=client.V1ConfigMapEnvSource( - name=name, - ) - ) - ) - else: - container.env.extend( - client.V1EnvVar( - name=envvar_name, - value_from=client.V1EnvVarSource( - config_map_key_ref=client.V1ConfigMapKeySelector( - name=name, - key=cm_key, - ) - ), - ) - for envvar_name, cm_key in mapping.items() - ) - return self - - def with_secrets_env( - self, name: str, *, mapping: dict[str, str] | None = None, container_id: int = 0 - ) -> Self: - container = self._init_env(container_id) - if mapping is None: - container.env_from.append( - client.V1EnvFromSource( - secret_ref=client.V1SecretEnvSource( - name=name, - ) - ) - ) - else: - container.env.extend( - client.V1EnvVar( - name=envvar_name, - value_from=client.V1EnvVarSource( - secret_key_ref=client.V1SecretKeySelector( - name=name, - key=secret_key, - ) - ), - ) - for envvar_name, secret_key in mapping.items() - ) - return self - - def get_container( - self, name: str | None = None, *, remove: bool = False - ) -> client.V1Container: - container_list = self.obj.spec.template.spec.containers - - if name is None: - if len(container_list) != 1: - raise ValueError("If multiple containers, must pass name.") - return container_list[0] - - for i, container in enumerate(container_list): - if container.name == name: - if remove: - return container_list.pop(i) - return container - - raise ValueError(f"No such container: {name}") - - def add_container(self, container: client.V1Container) -> int: - container_list = self.obj.spec.template.spec.containers - - names = set(c.name for c in container_list) - if container.name in names: - raise ValueError(f"Can't use name {container.name}, already in use.") - - container_list.append(container) - return len(container_list) - 1 + + def pod_spec(self) -> PodSpec: + return PodSpec(self.obj.spec.template.spec) def get_selector(self) -> dict[str, str]: return {self.SELECTOR_LABEL: self.obj.metadata.name} + + def __getattr__(self, name: str) -> Any: + return getattr(self.pod_spec(), name) diff --git a/transpire/resources/podspec.py b/transpire/resources/podspec.py new file mode 100644 index 0000000..6e045a0 --- /dev/null +++ b/transpire/resources/podspec.py @@ -0,0 +1,245 @@ +from typing import Self +from kubernetes import client + + +# This isn't a resource that can be instantiated directly, it's just a thin wrapper. +class PodSpec: + obj: client.V1PodSpec + + def __init__(self, obj: client.V1PodSpec) -> None: + self.obj = obj + + def _init_env(self, container_name: str | None = None) -> client.V1Container: + container = self.get_container(container_name) + + if container.env_from is None: + container.env_from = [] + if container.env is None: + container.env = [] + + return container + + def with_embedded_env(self, env: dict[str, str], *, container_name: str | None = None) -> Self: + container = self._init_env(container_name=container_name) + container.env.extend( + client.V1EnvVar( + name=envvar_name, + value=envvar_value, + ) + for envvar_name, envvar_value in env.items() + ) + return self + + def with_configmap_env( + self, + name: str, + *, + mapping: dict[str, str] | None = None, + container_name: str | None = None, + ) -> Self: + container = self._init_env(container_name=container_name) + if mapping is None: + container.env_from.append( + client.V1EnvFromSource( + config_map_ref=client.V1ConfigMapEnvSource( + name=name, + ) + ) + ) + else: + container.env.extend( + client.V1EnvVar( + name=envvar_name, + value_from=client.V1EnvVarSource( + config_map_key_ref=client.V1ConfigMapKeySelector( + name=name, + key=cm_key, + ) + ), + ) + for envvar_name, cm_key in mapping.items() + ) + return self + + def with_secret_env( + self, + name: str, + *, + mapping: dict[str, str] | None = None, + container_name: str | None = None, + ) -> Self: + container = self._init_env(container_name=container_name) + if mapping is None: + container.env_from.append( + client.V1EnvFromSource( + secret_ref=client.V1SecretEnvSource( + name=name, + ) + ) + ) + else: + container.env.extend( + client.V1EnvVar( + name=envvar_name, + value_from=client.V1EnvVarSource( + secret_key_ref=client.V1SecretKeySelector( + name=name, + key=secret_key, + ) + ), + ) + for envvar_name, secret_key in mapping.items() + ) + return self + + def with_configmap_volume( + self, name: str, mount_path: str, *, container_name: str | None = None + ) -> Self: + container = self.get_container(container_name) + if container.volume_mounts is None: + container.volume_mounts = [] + + container.volume_mounts.append( + client.V1VolumeMount( + name=name, + mount_path=mount_path, + ) + ) + + container.volumes.append( + client.V1Volume( + name=name, + config_map=client.V1ConfigMapVolumeSource( + name=name, + ), + ) + ) + + return self + + def with_secret_volume( + self, name: str, mount_path: str, *, container_name: str | None = None + ) -> Self: + container = self.get_container(container_name) + if container.volume_mounts is None: + container.volume_mounts = [] + + container.volume_mounts.append( + client.V1VolumeMount( + name=name, + mount_path=mount_path, + ) + ) + + container.volumes.append( + client.V1Volume( + name=name, + secret=client.V1SecretVolumeSource( + secret_name=name, + ), + ) + ) + + return self + + def with_pvc_volume( + self, name: str, mount_path: str, *, container_name: str | None = None + ) -> Self: + container = self.get_container(container_name) + if container.volume_mounts is None: + container.volume_mounts = [] + + container.volume_mounts.append( + client.V1VolumeMount( + name=name, + mount_path=mount_path, + ) + ) + + container.volumes.append( + client.V1Volume( + name=name, + persistent_volume_claim=client.V1PersistentVolumeClaimVolumeSource( + claim_name=name, + ), + ) + ) + + return self + + def with_arbitrary_volume( + self, volume: client.V1Volume, mount_path: str, *, container_name: str | None = None + ) -> Self: + container = self.get_container(container_name) + if container.volume_mounts is None: + container.volume_mounts = [] + + container.volume_mounts.append( + client.V1VolumeMount( + name=volume.name, + mount_path=mount_path, + ) + ) + + container.volumes.append(volume) + return self + + def get_container( + self, name: str | None = None, *, remove: bool = False + ) -> client.V1Container: + container_list = self.obj.spec.template.spec.containers + + if name is None: + if len(container_list) != 1: + raise ValueError("If multiple containers, must pass name.") + return container_list[0] + + for i, container in enumerate(container_list): + if container.name == name: + if remove: + return container_list.pop(i) + return container + + raise ValueError(f"No such container: {name}") + + def add_container(self, name: str, image: str) -> Self: + container_list = self.obj.spec.template.spec.containers + + names = set(c.name for c in container_list) + if name in names: + raise ValueError(f"Can't use name {name}, already in use.") + + container_list.append( + client.V1Container( + name=name, + image=image, + ) + ) + return self + + def add_arbitrary_container(self, container: client.V1Container) -> int: + container_list = self.obj.spec.template.spec.containers + + names = set(c.name for c in container_list) + if container.name in names: + raise ValueError(f"Can't use name {container.name}, already in use.") + + container_list.append(container) + return len(container_list) - 1 + + def with_probes( + self, + *, + liveness: client.V1Probe | None = None, + readiness: client.V1Probe | None = None, + startup: client.V1Probe | None = None, + container_name: str | None = None, + ) -> Self: + container = self.get_container(container_name) + if liveness is not None: + container.liveness_probe = liveness + if readiness is not None: + container.readiness_probe = readiness + if startup is not None: + container.startup_probe = startup + return self diff --git a/transpire/resources/pvc.py b/transpire/resources/pvc.py new file mode 100644 index 0000000..12d1a60 --- /dev/null +++ b/transpire/resources/pvc.py @@ -0,0 +1,28 @@ +from typing import Union + +from kubernetes import client + +from transpire.resources.base import Resource + + +class Service(Resource[client.V1PersistentVolumeClaim]): + def __init__( + self, + name: str, + size: str, + access_modes: list[str], + storage_class_name: str | None = None, + ): + self.obj = client.V1PersistentVolumeClaim( + api_version="v1", + kind="PersistentVolumeClaim", + metadata=client.V1ObjectMeta(name=name), + spec=client.V1PersistentVolumeClaimSpec( + resources=client.V1ResourceRequirements( + requests={"storage": storage} + ), + access_modes=access_modes, + storage_class_name=storage_class_name, + ), + ) + super().__init__() diff --git a/transpire/resources/statefulset.py b/transpire/resources/statefulset.py new file mode 100644 index 0000000..fc07a66 --- /dev/null +++ b/transpire/resources/statefulset.py @@ -0,0 +1,78 @@ +from typing import Any, List, Self, Union + +from kubernetes import client + +from transpire.resources.base import Resource +from transpire.resources.podspec import PodSpec + + +class Deployment(Resource[client.V1StatefulSet]): + SELECTOR_LABEL = "transpire.ocf.io/deployment" + + def __init__( + self, + name: str, + image: str, + ports: List[Union[str, int]], + service_name: str, + *, + args: List[str] | None = None, + ): + self.obj = client.V1StatefulSet( + api_version="apps/v1", + kind="StatefulSet", + metadata=client.V1ObjectMeta(name=name), + spec=client.V1StatefulSetSpec( + replicas=1, + service_name=service_name, + selector=client.V1LabelSelector( + match_labels={self.SELECTOR_LABEL: name} + ), + template=client.V1PodTemplateSpec( + metadata=client.V1ObjectMeta(labels={self.SELECTOR_LABEL: name}), + spec=client.V1PodSpec( + containers=[ + client.V1Container( + name="main", + image=image, + image_pull_policy="IfNotPresent", + args=args, + ports=[ + client.V1ContainerPort(container_port=x) + for x in ports + ], + ) + ] + ), + ), + ), + ) + super().__init__() + + def pod_spec(self) -> PodSpec: + return PodSpec(self.obj.spec.template.spec) + + def get_selector(self) -> dict[str, str]: + return {self.SELECTOR_LABEL: self.obj.metadata.name} + + def with_volume_template(self, name: str, size: str, access_modes: list[str], storage_class_name: str | None = None) -> Self: + self.obj.spec.volume_claim_templates.append( + client.V1PersistentVolumeClaim( + metadata=client.V1ObjectMeta(name=name), + spec=client.V1PersistentVolumeClaimSpec( + access_modes=access_modes, + resources=client.V1ResourceRequirements( + requests={"storage": size} + ), + storage_class_name=storage_class_name, + ), + ) + ) + return self + + def with_arbitrary_volume_template(self, template: client.V1PersistentVolumeClaim) -> Self: + self.obj.spec.volume_claim_templates.append(template) + return self + + def __getattr__(self, name: str) -> Any: + return getattr(self.pod_spec(), name)