Skip to main content

oximedia_workflow/
builder.rs

1//! Builder pattern for workflow construction.
2
3use crate::error::Result;
4use crate::task::{RetryPolicy, Task, TaskId, TaskPriority, TaskType};
5use crate::workflow::{Workflow, WorkflowConfig};
6use std::collections::HashMap;
7use std::time::Duration;
8
9/// Workflow builder for ergonomic workflow construction.
10pub struct WorkflowBuilder {
11    workflow: Workflow,
12    task_map: HashMap<String, TaskId>,
13}
14
15impl WorkflowBuilder {
16    /// Create a new workflow builder.
17    #[must_use]
18    pub fn new(name: impl Into<String>) -> Self {
19        Self {
20            workflow: Workflow::new(name),
21            task_map: HashMap::new(),
22        }
23    }
24
25    /// Set workflow description.
26    #[must_use]
27    pub fn description(mut self, description: impl Into<String>) -> Self {
28        self.workflow.description = description.into();
29        self
30    }
31
32    /// Set workflow configuration.
33    #[must_use]
34    pub fn config(mut self, config: WorkflowConfig) -> Self {
35        self.workflow.config = config;
36        self
37    }
38
39    /// Set maximum concurrent tasks.
40    #[must_use]
41    pub fn max_concurrent_tasks(mut self, max: usize) -> Self {
42        self.workflow.config.max_concurrent_tasks = max;
43        self
44    }
45
46    /// Set global timeout.
47    #[must_use]
48    pub fn global_timeout(mut self, timeout: Duration) -> Self {
49        self.workflow.config.global_timeout = Some(timeout);
50        self
51    }
52
53    /// Enable fail-fast mode.
54    #[must_use]
55    pub fn fail_fast(mut self, enabled: bool) -> Self {
56        self.workflow.config.fail_fast = enabled;
57        self
58    }
59
60    /// Enable continue-on-error mode.
61    #[must_use]
62    pub fn continue_on_error(mut self, enabled: bool) -> Self {
63        self.workflow.config.continue_on_error = enabled;
64        self
65    }
66
67    /// Set a workflow variable.
68    #[must_use]
69    pub fn variable(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
70        self.workflow.config.variables.insert(key.into(), value);
71        self
72    }
73
74    /// Add a task with a named reference.
75    #[must_use]
76    pub fn task(mut self, name: impl Into<String>, task: Task) -> Self {
77        let name_str = name.into();
78        let task_id = task.id;
79        self.workflow.add_task(task);
80        self.task_map.insert(name_str, task_id);
81        self
82    }
83
84    /// Add a task with `TaskBuilder`.
85    #[must_use]
86    pub fn add_task(mut self, builder: TaskBuilder) -> Self {
87        let (name, task) = builder.build();
88        let task_id = task.id;
89        self.workflow.add_task(task);
90        if let Some(n) = name {
91            self.task_map.insert(n, task_id);
92        }
93        self
94    }
95
96    /// Add a dependency between named tasks.
97    pub fn depends_on(
98        mut self,
99        task: impl AsRef<str>,
100        dependency: impl AsRef<str>,
101    ) -> Result<Self> {
102        let task_id = self
103            .task_map
104            .get(task.as_ref())
105            .ok_or_else(|| crate::error::WorkflowError::TaskNotFound(task.as_ref().to_string()))?;
106
107        let dep_id = self.task_map.get(dependency.as_ref()).ok_or_else(|| {
108            crate::error::WorkflowError::TaskNotFound(dependency.as_ref().to_string())
109        })?;
110
111        self.workflow.add_edge(*dep_id, *task_id)?;
112        Ok(self)
113    }
114
115    /// Add a conditional dependency.
116    pub fn conditional_depends_on(
117        mut self,
118        task: impl AsRef<str>,
119        dependency: impl AsRef<str>,
120        condition: impl Into<String>,
121    ) -> Result<Self> {
122        let task_id = *self
123            .task_map
124            .get(task.as_ref())
125            .ok_or_else(|| crate::error::WorkflowError::TaskNotFound(task.as_ref().to_string()))?;
126
127        let dep_id = *self.task_map.get(dependency.as_ref()).ok_or_else(|| {
128            crate::error::WorkflowError::TaskNotFound(dependency.as_ref().to_string())
129        })?;
130
131        self.workflow
132            .add_conditional_edge(dep_id, task_id, condition.into())?;
133        Ok(self)
134    }
135
136    /// Build the workflow.
137    pub fn build(self) -> Result<Workflow> {
138        self.workflow.validate()?;
139        Ok(self.workflow)
140    }
141}
142
143/// Task builder for ergonomic task construction.
144pub struct TaskBuilder {
145    name: Option<String>,
146    task_name: String,
147    task_type: TaskType,
148    priority: TaskPriority,
149    retry: RetryPolicy,
150    timeout: Duration,
151    metadata: HashMap<String, String>,
152    conditions: Vec<String>,
153}
154
155impl TaskBuilder {
156    /// Create a new task builder.
157    #[must_use]
158    pub fn new(task_name: impl Into<String>, task_type: TaskType) -> Self {
159        Self {
160            name: None,
161            task_name: task_name.into(),
162            task_type,
163            priority: TaskPriority::Normal,
164            retry: RetryPolicy::default(),
165            timeout: Duration::from_secs(3600),
166            metadata: HashMap::new(),
167            conditions: Vec::new(),
168        }
169    }
170
171    /// Set a named reference for this task.
172    #[must_use]
173    pub fn named(mut self, name: impl Into<String>) -> Self {
174        self.name = Some(name.into());
175        self
176    }
177
178    /// Set task priority.
179    #[must_use]
180    pub fn priority(mut self, priority: TaskPriority) -> Self {
181        self.priority = priority;
182        self
183    }
184
185    /// Set retry policy.
186    #[must_use]
187    pub fn retry(mut self, retry: RetryPolicy) -> Self {
188        self.retry = retry;
189        self
190    }
191
192    /// Set retry attempts.
193    #[must_use]
194    pub fn retry_attempts(mut self, attempts: u32) -> Self {
195        self.retry.max_attempts = attempts;
196        self
197    }
198
199    /// Set timeout.
200    #[must_use]
201    pub fn timeout(mut self, timeout: Duration) -> Self {
202        self.timeout = timeout;
203        self
204    }
205
206    /// Add metadata.
207    #[must_use]
208    pub fn metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
209        self.metadata.insert(key.into(), value.into());
210        self
211    }
212
213    /// Add condition.
214    #[must_use]
215    pub fn condition(mut self, condition: impl Into<String>) -> Self {
216        self.conditions.push(condition.into());
217        self
218    }
219
220    /// Build the task.
221    #[must_use]
222    pub fn build(self) -> (Option<String>, Task) {
223        let mut task = Task::new(self.task_name, self.task_type);
224        task.priority = self.priority;
225        task.retry = self.retry;
226        task.timeout = self.timeout;
227        task.metadata = self.metadata;
228        task.conditions = self.conditions;
229
230        (self.name, task)
231    }
232}
233
234/// Transcode task builder.
235pub struct TranscodeTaskBuilder {
236    inner: TaskBuilder,
237}
238
239impl TranscodeTaskBuilder {
240    /// Create a new transcode task builder.
241    #[must_use]
242    pub fn new(
243        name: impl Into<String>,
244        input: impl Into<std::path::PathBuf>,
245        output: impl Into<std::path::PathBuf>,
246        preset: impl Into<String>,
247    ) -> Self {
248        let task_type = TaskType::Transcode {
249            input: input.into(),
250            output: output.into(),
251            preset: preset.into(),
252            params: HashMap::new(),
253        };
254
255        Self {
256            inner: TaskBuilder::new(name, task_type),
257        }
258    }
259
260    /// Set task name reference.
261    #[must_use]
262    pub fn named(mut self, name: impl Into<String>) -> Self {
263        self.inner = self.inner.named(name);
264        self
265    }
266
267    /// Set priority.
268    #[must_use]
269    pub fn priority(mut self, priority: TaskPriority) -> Self {
270        self.inner = self.inner.priority(priority);
271        self
272    }
273
274    /// Set timeout.
275    #[must_use]
276    pub fn timeout(mut self, timeout: Duration) -> Self {
277        self.inner = self.inner.timeout(timeout);
278        self
279    }
280
281    /// Add parameter.
282    #[must_use]
283    pub fn param(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
284        if let TaskType::Transcode { ref mut params, .. } = self.inner.task_type {
285            params.insert(key.into(), value);
286        }
287        self
288    }
289
290    /// Build the task.
291    #[must_use]
292    pub fn build(self) -> (Option<String>, Task) {
293        self.inner.build()
294    }
295}
296
297/// QC task builder.
298pub struct QcTaskBuilder {
299    inner: TaskBuilder,
300}
301
302impl QcTaskBuilder {
303    /// Create a new QC task builder.
304    #[must_use]
305    pub fn new(
306        name: impl Into<String>,
307        input: impl Into<std::path::PathBuf>,
308        profile: impl Into<String>,
309    ) -> Self {
310        let task_type = TaskType::QualityControl {
311            input: input.into(),
312            profile: profile.into(),
313            rules: Vec::new(),
314        };
315
316        Self {
317            inner: TaskBuilder::new(name, task_type),
318        }
319    }
320
321    /// Set task name reference.
322    #[must_use]
323    pub fn named(mut self, name: impl Into<String>) -> Self {
324        self.inner = self.inner.named(name);
325        self
326    }
327
328    /// Add validation rule.
329    #[must_use]
330    pub fn rule(mut self, rule: impl Into<String>) -> Self {
331        if let TaskType::QualityControl { ref mut rules, .. } = self.inner.task_type {
332            rules.push(rule.into());
333        }
334        self
335    }
336
337    /// Build the task.
338    #[must_use]
339    pub fn build(self) -> (Option<String>, Task) {
340        self.inner.build()
341    }
342}
343
344/// Transfer task builder.
345pub struct TransferTaskBuilder {
346    inner: TaskBuilder,
347}
348
349impl TransferTaskBuilder {
350    /// Create a new transfer task builder.
351    #[must_use]
352    pub fn new(
353        name: impl Into<String>,
354        source: impl Into<String>,
355        destination: impl Into<String>,
356        protocol: crate::task::TransferProtocol,
357    ) -> Self {
358        let task_type = TaskType::Transfer {
359            source: source.into(),
360            destination: destination.into(),
361            protocol,
362            options: HashMap::new(),
363        };
364
365        Self {
366            inner: TaskBuilder::new(name, task_type),
367        }
368    }
369
370    /// Set task name reference.
371    #[must_use]
372    pub fn named(mut self, name: impl Into<String>) -> Self {
373        self.inner = self.inner.named(name);
374        self
375    }
376
377    /// Add option.
378    #[must_use]
379    pub fn option(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
380        if let TaskType::Transfer {
381            ref mut options, ..
382        } = self.inner.task_type
383        {
384            options.insert(key.into(), value.into());
385        }
386        self
387    }
388
389    /// Build the task.
390    #[must_use]
391    pub fn build(self) -> (Option<String>, Task) {
392        self.inner.build()
393    }
394}
395
396#[cfg(test)]
397mod tests {
398    use super::*;
399
400    #[test]
401    fn test_workflow_builder() {
402        let workflow = WorkflowBuilder::new("test-workflow")
403            .description("Test workflow")
404            .max_concurrent_tasks(4)
405            .fail_fast(true)
406            .build()
407            .expect("should succeed in test");
408
409        assert_eq!(workflow.name, "test-workflow");
410        assert_eq!(workflow.description, "Test workflow");
411        assert_eq!(workflow.config.max_concurrent_tasks, 4);
412        assert!(workflow.config.fail_fast);
413    }
414
415    #[test]
416    fn test_workflow_builder_with_tasks() {
417        let task1 = Task::new(
418            "task1",
419            TaskType::Wait {
420                duration: Duration::from_secs(1),
421            },
422        );
423        let task2 = Task::new(
424            "task2",
425            TaskType::Wait {
426                duration: Duration::from_secs(1),
427            },
428        );
429
430        let workflow = WorkflowBuilder::new("test")
431            .task("t1", task1)
432            .task("t2", task2)
433            .depends_on("t2", "t1")
434            .expect("should succeed in test")
435            .build()
436            .expect("should succeed in test");
437
438        assert_eq!(workflow.tasks.len(), 2);
439        assert_eq!(workflow.edges.len(), 1);
440    }
441
442    #[test]
443    fn test_task_builder() {
444        let (name, task) = TaskBuilder::new(
445            "test-task",
446            TaskType::Wait {
447                duration: Duration::from_secs(5),
448            },
449        )
450        .named("my-task")
451        .priority(TaskPriority::High)
452        .timeout(Duration::from_secs(10))
453        .metadata("key", "value")
454        .condition("x > 5")
455        .build();
456
457        assert_eq!(name, Some("my-task".to_string()));
458        assert_eq!(task.name, "test-task");
459        assert_eq!(task.priority, TaskPriority::High);
460        assert_eq!(task.timeout, Duration::from_secs(10));
461        assert_eq!(task.metadata.get("key"), Some(&"value".to_string()));
462        assert_eq!(task.conditions.len(), 1);
463    }
464
465    #[test]
466    fn test_transcode_task_builder() {
467        let (name, task) =
468            TranscodeTaskBuilder::new("transcode", "/input.mp4", "/output.mp4", "h264")
469                .named("my-transcode")
470                .priority(TaskPriority::High)
471                .param("bitrate", serde_json::json!(5000000))
472                .build();
473
474        assert_eq!(name, Some("my-transcode".to_string()));
475        assert_eq!(task.name, "transcode");
476
477        if let TaskType::Transcode { params, .. } = &task.task_type {
478            assert_eq!(params.get("bitrate"), Some(&serde_json::json!(5000000)));
479        } else {
480            panic!("Wrong task type");
481        }
482    }
483
484    #[test]
485    fn test_qc_task_builder() {
486        let (_, task) = QcTaskBuilder::new("qc", "/input.mp4", "broadcast")
487            .rule("video_bitrate")
488            .rule("audio_levels")
489            .build();
490
491        if let TaskType::QualityControl { rules, .. } = &task.task_type {
492            assert_eq!(rules.len(), 2);
493            assert!(rules.contains(&"video_bitrate".to_string()));
494        } else {
495            panic!("Wrong task type");
496        }
497    }
498
499    #[test]
500    fn test_transfer_task_builder() {
501        let (_, task) = TransferTaskBuilder::new(
502            "transfer",
503            "/local/file.mp4",
504            "s3://bucket/file.mp4",
505            crate::task::TransferProtocol::S3,
506        )
507        .option("storage_class", "STANDARD")
508        .build();
509
510        if let TaskType::Transfer { options, .. } = &task.task_type {
511            assert_eq!(options.get("storage_class"), Some(&"STANDARD".to_string()));
512        } else {
513            panic!("Wrong task type");
514        }
515    }
516
517    #[test]
518    fn test_workflow_builder_variables() {
519        let workflow = WorkflowBuilder::new("test")
520            .variable("input_dir", serde_json::json!("/inputs"))
521            .variable("output_dir", serde_json::json!("/outputs"))
522            .build()
523            .expect("should succeed in test");
524
525        assert_eq!(workflow.config.variables.len(), 2);
526        assert_eq!(
527            workflow.config.variables.get("input_dir"),
528            Some(&serde_json::json!("/inputs"))
529        );
530    }
531}