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

armada_client/
builder.rs
