1use 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
9pub struct WorkflowBuilder {
11 workflow: Workflow,
12 task_map: HashMap<String, TaskId>,
13}
14
15impl WorkflowBuilder {
16 #[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 #[must_use]
27 pub fn description(mut self, description: impl Into<String>) -> Self {
28 self.workflow.description = description.into();
29 self
30 }
31
32 #[must_use]
34 pub fn config(mut self, config: WorkflowConfig) -> Self {
35 self.workflow.config = config;
36 self
37 }
38
39 #[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 #[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 #[must_use]
55 pub fn fail_fast(mut self, enabled: bool) -> Self {
56 self.workflow.config.fail_fast = enabled;
57 self
58 }
59
60 #[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 #[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 #[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 #[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 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 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 pub fn build(self) -> Result<Workflow> {
138 self.workflow.validate()?;
139 Ok(self.workflow)
140 }
141}
142
143pub 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 #[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 #[must_use]
173 pub fn named(mut self, name: impl Into<String>) -> Self {
174 self.name = Some(name.into());
175 self
176 }
177
178 #[must_use]
180 pub fn priority(mut self, priority: TaskPriority) -> Self {
181 self.priority = priority;
182 self
183 }
184
185 #[must_use]
187 pub fn retry(mut self, retry: RetryPolicy) -> Self {
188 self.retry = retry;
189 self
190 }
191
192 #[must_use]
194 pub fn retry_attempts(mut self, attempts: u32) -> Self {
195 self.retry.max_attempts = attempts;
196 self
197 }
198
199 #[must_use]
201 pub fn timeout(mut self, timeout: Duration) -> Self {
202 self.timeout = timeout;
203 self
204 }
205
206 #[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 #[must_use]
215 pub fn condition(mut self, condition: impl Into<String>) -> Self {
216 self.conditions.push(condition.into());
217 self
218 }
219
220 #[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
234pub struct TranscodeTaskBuilder {
236 inner: TaskBuilder,
237}
238
239impl TranscodeTaskBuilder {
240 #[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 #[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 #[must_use]
269 pub fn priority(mut self, priority: TaskPriority) -> Self {
270 self.inner = self.inner.priority(priority);
271 self
272 }
273
274 #[must_use]
276 pub fn timeout(mut self, timeout: Duration) -> Self {
277 self.inner = self.inner.timeout(timeout);
278 self
279 }
280
281 #[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 #[must_use]
292 pub fn build(self) -> (Option<String>, Task) {
293 self.inner.build()
294 }
295}
296
297pub struct QcTaskBuilder {
299 inner: TaskBuilder,
300}
301
302impl QcTaskBuilder {
303 #[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 #[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 #[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 #[must_use]
339 pub fn build(self) -> (Option<String>, Task) {
340 self.inner.build()
341 }
342}
343
344pub struct TransferTaskBuilder {
346 inner: TaskBuilder,
347}
348
349impl TransferTaskBuilder {
350 #[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 #[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 #[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 #[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}