Skip to main content

armada_client/
builder.rs

1use std::collections::HashMap;
2use std::marker::PhantomData;
3
4use crate::api::{IngressConfig, JobSubmitRequestItem, ServiceConfig};
5use crate::k8s::io::api::core::v1::PodSpec;
6
7/// Typestate marker: builder has no pod spec yet — `.build()` is unavailable.
8#[doc(hidden)]
9pub struct NoPodSpec;
10/// Typestate marker: builder has pod spec set — `.build()` is available.
11#[doc(hidden)]
12pub struct HasPodSpec;
13
14/// Typestate builder for [`JobSubmitRequestItem`].
15///
16/// Start with [`JobRequestItemBuilder::new`], set fields, then call
17/// `.pod_spec(spec)` or `.pod_specs(specs)` to transition to the `HasPodSpec`
18/// state, which unlocks `.build()`. Attempting to call `.build()` before
19/// providing a pod spec is a **compile-time error**.
20///
21/// # Example
22///
23/// ```ignore
24/// let item = JobRequestItemBuilder::new()
25///     .namespace("default")
26///     .priority(1.0)
27///     .label("app", "my-app")
28///     .pod_spec(pod_spec)
29///     .build();
30/// ```
31pub struct JobRequestItemBuilder<S> {
32    priority: f64,
33    namespace: String,
34    client_id: String,
35    labels: HashMap<String, String>,
36    annotations: HashMap<String, String>,
37    scheduler: String,
38    ingress: Vec<IngressConfig>,
39    services: Vec<ServiceConfig>,
40    pod_specs: Vec<PodSpec>,
41    _state: PhantomData<S>,
42}
43
44impl JobRequestItemBuilder<NoPodSpec> {
45    /// Create a new builder with all fields set to their zero/empty defaults.
46    ///
47    /// Call setters to configure the job, then supply a pod spec with
48    /// [`.pod_spec()`](JobRequestItemBuilder::pod_spec) or
49    /// [`.pod_specs()`](JobRequestItemBuilder::pod_specs) to unlock
50    /// [`.build()`](JobRequestItemBuilder::build).
51    pub fn new() -> Self {
52        Self {
53            priority: 0.0,
54            namespace: String::new(),
55            client_id: String::new(),
56            labels: HashMap::new(),
57            annotations: HashMap::new(),
58            scheduler: String::new(),
59            ingress: Vec::new(),
60            services: Vec::new(),
61            pod_specs: Vec::new(),
62            _state: PhantomData,
63        }
64    }
65}
66
67impl Default for JobRequestItemBuilder<NoPodSpec> {
68    fn default() -> Self {
69        Self::new()
70    }
71}
72
73// Setters available in all states — preserve the state type via generics.
74impl<S> JobRequestItemBuilder<S> {
75    /// Job priority relative to others in the queue. Higher values are
76    /// scheduled first. Must be non-negative; the server will reject negative
77    /// values. Defaults to `0.0`.
78    #[must_use]
79    pub fn priority(mut self, p: f64) -> Self {
80        self.priority = p;
81        self
82    }
83
84    /// Set the Kubernetes namespace in which the job's pod will run.
85    ///
86    /// Typically `"default"` unless your cluster uses a dedicated namespace
87    /// for batch workloads.
88    #[must_use]
89    pub fn namespace(mut self, ns: impl Into<String>) -> Self {
90        self.namespace = ns.into();
91        self
92    }
93
94    /// Set an opaque client-supplied identifier for idempotency tracking.
95    ///
96    /// If the same `client_id` is submitted twice, Armada will deduplicate
97    /// the request and return the existing job rather than creating a new one.
98    /// Leave empty (the default) to disable deduplication.
99    #[must_use]
100    pub fn client_id(mut self, id: impl Into<String>) -> Self {
101        self.client_id = id.into();
102        self
103    }
104
105    /// Replace the entire labels map.
106    #[must_use]
107    pub fn labels(mut self, l: HashMap<String, String>) -> Self {
108        self.labels = l;
109        self
110    }
111
112    /// Insert a single label. Can be chained multiple times.
113    #[must_use]
114    pub fn label(mut self, k: impl Into<String>, v: impl Into<String>) -> Self {
115        self.labels.insert(k.into(), v.into());
116        self
117    }
118
119    /// Replace the entire annotations map.
120    #[must_use]
121    pub fn annotations(mut self, a: HashMap<String, String>) -> Self {
122        self.annotations = a;
123        self
124    }
125
126    /// Insert a single annotation. Can be chained multiple times.
127    #[must_use]
128    pub fn annotation(mut self, k: impl Into<String>, v: impl Into<String>) -> Self {
129        self.annotations.insert(k.into(), v.into());
130        self
131    }
132
133    /// Override the scheduler for this job. Leave empty to use the cluster default.
134    #[must_use]
135    pub fn scheduler(mut self, s: impl Into<String>) -> Self {
136        self.scheduler = s.into();
137        self
138    }
139
140    /// Replace the entire ingress config list.
141    #[must_use]
142    pub fn ingress(mut self, i: Vec<IngressConfig>) -> Self {
143        self.ingress = i;
144        self
145    }
146
147    /// Append a single ingress config. Can be chained multiple times.
148    #[must_use]
149    pub fn add_ingress(mut self, i: IngressConfig) -> Self {
150        self.ingress.push(i);
151        self
152    }
153
154    /// Replace the entire service config list.
155    #[must_use]
156    pub fn services(mut self, s: Vec<ServiceConfig>) -> Self {
157        self.services = s;
158        self
159    }
160
161    /// Append a single service config. Can be chained multiple times.
162    #[must_use]
163    pub fn add_service(mut self, s: ServiceConfig) -> Self {
164        self.services.push(s);
165        self
166    }
167
168    /// Set a single pod spec and transition the builder to [`HasPodSpec`] state,
169    /// which unlocks `.build()`.
170    ///
171    /// Shorthand for `.pod_specs(vec![spec])`. If called again on a
172    /// `JobRequestItemBuilder<HasPodSpec>`, the previous spec is replaced.
173    #[must_use]
174    pub fn pod_spec(self, spec: PodSpec) -> JobRequestItemBuilder<HasPodSpec> {
175        self.pod_specs(vec![spec])
176    }
177
178    /// Set multiple pod specs and transition the builder to [`HasPodSpec`] state,
179    /// which unlocks `.build()`.
180    ///
181    /// If called again on a `JobRequestItemBuilder<HasPodSpec>`, the previous
182    /// specs are replaced entirely.
183    #[must_use]
184    pub fn pod_specs(self, specs: Vec<PodSpec>) -> JobRequestItemBuilder<HasPodSpec> {
185        JobRequestItemBuilder {
186            priority: self.priority,
187            namespace: self.namespace,
188            client_id: self.client_id,
189            labels: self.labels,
190            annotations: self.annotations,
191            scheduler: self.scheduler,
192            ingress: self.ingress,
193            services: self.services,
194            pod_specs: specs,
195            _state: PhantomData,
196        }
197    }
198}
199
200// `.build()` is only available when the builder is in `HasPodSpec` state.
201impl JobRequestItemBuilder<HasPodSpec> {
202    /// Consume the builder and return a [`JobSubmitRequestItem`] ready to be
203    /// included in a [`crate::JobSubmitRequest`].
204    ///
205    /// This method is only available after calling
206    /// [`.pod_spec()`](JobRequestItemBuilder::pod_spec) or
207    /// [`.pod_specs()`](JobRequestItemBuilder::pod_specs). Calling `.build()`
208    /// on a `JobRequestItemBuilder<NoPodSpec>` is a **compile-time error**.
209    #[must_use]
210    pub fn build(self) -> JobSubmitRequestItem {
211        #[allow(deprecated)]
212        JobSubmitRequestItem {
213            priority: self.priority,
214            namespace: self.namespace,
215            client_id: self.client_id,
216            labels: self.labels,
217            annotations: self.annotations,
218            pod_specs: self.pod_specs,
219            ingress: self.ingress,
220            services: self.services,
221            scheduler: self.scheduler,
222            // Deprecated singular fields — zeroed, not used by this builder
223            pod_spec: None,
224            required_node_labels: HashMap::new(),
225        }
226    }
227}
228
229#[cfg(test)]
230mod tests {
231    use super::*;
232
233    fn minimal_pod_spec() -> PodSpec {
234        PodSpec {
235            containers: vec![],
236            ..Default::default()
237        }
238    }
239
240    #[test]
241    fn builder_with_pod_specs_builds_correctly() {
242        let spec = minimal_pod_spec();
243        let item = JobRequestItemBuilder::new()
244            .namespace("default")
245            .priority(1.0)
246            .pod_specs(vec![spec])
247            .build();
248
249        assert_eq!(item.namespace, "default");
250        assert_eq!(item.priority, 1.0);
251        assert_eq!(item.pod_specs.len(), 1);
252    }
253
254    #[test]
255    fn pod_spec_singular_shorthand() {
256        let item = JobRequestItemBuilder::new()
257            .pod_spec(minimal_pod_spec())
258            .build();
259        assert_eq!(item.pod_specs.len(), 1);
260    }
261
262    #[test]
263    fn pod_spec_called_twice_replaces_previous() {
264        let spec_a = minimal_pod_spec();
265        let mut spec_b = minimal_pod_spec();
266        spec_b.restart_policy = Some("Never".to_string());
267
268        let item = JobRequestItemBuilder::new()
269            .pod_spec(spec_a)
270            .pod_spec(spec_b)
271            .build();
272
273        assert_eq!(item.pod_specs.len(), 1);
274        assert_eq!(item.pod_specs[0].restart_policy.as_deref(), Some("Never"));
275    }
276
277    #[test]
278    fn label_and_annotation_helpers() {
279        let item = JobRequestItemBuilder::new()
280            .label("app", "my-app")
281            .label("env", "prod")
282            .annotation("owner", "team-a")
283            .pod_spec(minimal_pod_spec())
284            .build();
285
286        assert_eq!(item.labels.get("app").map(String::as_str), Some("my-app"));
287        assert_eq!(item.labels.get("env").map(String::as_str), Some("prod"));
288        assert_eq!(
289            item.annotations.get("owner").map(String::as_str),
290            Some("team-a")
291        );
292    }
293
294    #[test]
295    fn optional_fields_default_to_empty() {
296        let item = JobRequestItemBuilder::new()
297            .pod_specs(vec![minimal_pod_spec()])
298            .build();
299
300        assert_eq!(item.priority, 0.0);
301        assert!(item.namespace.is_empty());
302        assert!(item.client_id.is_empty());
303        assert!(item.labels.is_empty());
304        assert!(item.annotations.is_empty());
305        assert!(item.scheduler.is_empty());
306        assert!(item.ingress.is_empty());
307        assert!(item.services.is_empty());
308    }
309}