dis_spawner/
builder.rs

1use crate::{BackendPodSpec, SessionLivedBackend, SessionLivedBackendSpec, image_pull_policy::ImagePullPolicy};
2use k8s_openapi::api::core::v1::{Container, EnvVar, LocalObjectReference};
3use kube::core::ObjectMeta;
4use std::collections::{BTreeMap, HashMap};
5
6pub const DEFAULT_GRACE_SECONDS: u32 = 300;
7pub const APPLICATION: &str = "application";
8pub const DEFAULT_PREFIX: &str = "spawner-";
9
10pub struct SessionLivedBackendBuilder {
11    image: String,
12    image_pull_secret: Option<String>,
13    env: HashMap<String, String>,
14    image_pull_policy: Option<ImagePullPolicy>,
15    namespace: Option<String>,
16    grace_period_seconds: u32,
17    http_port: Option<u16>,
18    labels: Option<BTreeMap<String, String>>,
19}
20
21impl SessionLivedBackendBuilder {
22    pub fn new(image: &str) -> Self {
23        SessionLivedBackendBuilder {
24            image: image.to_string(),
25            env: HashMap::default(),
26            image_pull_policy: None,
27            image_pull_secret: None,
28            namespace: None,
29            grace_period_seconds: DEFAULT_GRACE_SECONDS,
30            http_port: None,
31            labels: None,
32        }
33    }
34
35    pub fn with_labels(self, labels: Option<BTreeMap<String, String>>) -> Self {
36        SessionLivedBackendBuilder { labels, ..self }
37    }
38
39    pub fn with_env(self, env: HashMap<String, String>) -> Self {
40        SessionLivedBackendBuilder { env, ..self }
41    }
42
43    pub fn with_port(self, http_port: Option<u16>) -> Self {
44        SessionLivedBackendBuilder { http_port, ..self }
45    }
46
47    pub fn with_image_pull_policy(self, image_pull_policy: Option<ImagePullPolicy>) -> Self {
48        SessionLivedBackendBuilder {
49            image_pull_policy,
50            ..self
51        }
52    }
53
54    pub fn with_namespace(self, namespace: Option<String>) -> Self {
55        SessionLivedBackendBuilder { namespace, ..self }
56    }
57
58    pub fn with_grace_period(self, grace_period_seconds: Option<u32>) -> Self {
59        SessionLivedBackendBuilder {
60            grace_period_seconds: grace_period_seconds.unwrap_or(DEFAULT_GRACE_SECONDS),
61            ..self
62        }
63    }
64
65    pub fn with_image_pull_secret(self, image_pull_secret: Option<String>) -> Self {
66        SessionLivedBackendBuilder {
67            image_pull_secret,
68            ..self
69        }
70    }
71
72    pub fn build_spec(&self) -> SessionLivedBackendSpec {
73        let env: Vec<EnvVar> = self
74            .env
75            .iter()
76            .map(|(name, value)| EnvVar {
77                name: name.to_string(),
78                value: Some(value.to_string()),
79                ..EnvVar::default()
80            })
81            .collect();
82
83        SessionLivedBackendSpec {
84            http_port: self.http_port,
85            grace_period_seconds: Some(self.grace_period_seconds),
86            template: BackendPodSpec {
87                containers: vec![Container {
88                    image: Some(self.image.to_string()),
89                    image_pull_policy: self.image_pull_policy.as_ref().map(|d| d.to_string()),
90                    env: Some(env),
91                    name: APPLICATION.to_string(),
92                    ..Default::default()
93                }],
94                image_pull_secrets: self.image_pull_secret.as_ref().map(|d| {
95                    vec![LocalObjectReference {
96                        name: Some(d.to_string()),
97                    }]
98                }),
99                ..Default::default()
100            },
101        }
102    }
103
104    pub fn build_prefixed(&self, prefix: &str) -> SessionLivedBackend {
105        SessionLivedBackend {
106            metadata: ObjectMeta {
107                generate_name: Some(prefix.to_string()),
108                labels: self.labels.clone(),
109                ..Default::default()
110            },
111            spec: self.build_spec(),
112            status: None,
113        }
114    }
115
116    pub fn build_named(&self, name: &str) -> SessionLivedBackend {
117        SessionLivedBackend {
118            metadata: ObjectMeta {
119                name: Some(name.to_string()),
120                namespace: self.namespace.clone(),
121                labels: self.labels.clone(),
122                ..Default::default()
123            },
124            spec: self.build_spec(),
125            status: None,
126        }
127    }
128
129    pub fn build(&self) -> SessionLivedBackend {
130        self.build_prefixed(DEFAULT_PREFIX)
131    }
132}