Coverage for cli / commands / dag_cmd.py: 100.00%

1"""DAG pipeline commands.""" 

2 

3import sys 

4from typing import Any 

5 

6import click 

7 

8from ..config import GCOConfig 

9from ..output import get_output_formatter 

10 

11pass_config = click.make_pass_decorator(GCOConfig, ensure=True) 

12 

13 

14@click.group() 

15@pass_config 

16def dag(config: Any) -> None: 

17 """Run multi-step job pipelines with dependencies.""" 

18 

19 

20@dag.command("run") 

21@click.argument("dag_file", type=click.Path(exists=True)) 

22@click.option("--region", "-r", help="Region to run in (default: from DAG file or first deployed)") 

23@click.option( 

24 "--timeout", "-t", default=3600, type=int, help="Timeout per step in seconds (default: 3600)" 

25) 

26@click.option("--dry-run", is_flag=True, help="Validate the DAG without running it") 

27@pass_config 

28def dag_run(config: Any, dag_file: Any, region: Any, timeout: Any, dry_run: Any) -> None: 

29 """Run a DAG pipeline from a YAML definition. 

30 

31 The DAG file defines steps with dependencies. Steps run in order, 

32 and downstream steps are skipped if a dependency fails. 

33 

34 Examples: 

35 gco dag run pipeline.yaml 

36 gco dag run pipeline.yaml -r us-east-1 

37 gco dag run pipeline.yaml --dry-run 

38 """ 

39 from ..dag import get_dag_runner, load_dag 

40 

41 formatter = get_output_formatter(config) 

42 

43 try: 

44 dag_def = load_dag(dag_file) 

45 errors = dag_def.validate() 

46 

47 if errors: 

48 for err in errors: 

49 formatter.print_error(err) 

50 sys.exit(1) 

51 

52 if dry_run: 

53 formatter.print_success(f"DAG '{dag_def.name}' is valid ({len(dag_def.steps)} steps)") 

54 print("\n Execution order:") 

55 # Show topological order 

56 completed: set[str] = set() 

57 order = 1 

58 remaining = list(dag_def.steps) 

59 while remaining: 

60 batch = [s for s in remaining if all(d in completed for d in s.depends_on)] 

61 for step in batch: 

62 deps = f" (after: {', '.join(step.depends_on)})" if step.depends_on else "" 

63 print(f" {order}. {step.name} → {step.manifest}{deps}") 

64 completed.add(step.name) 

65 remaining.remove(step) 

66 order += 1 

67 print() 

68 return 

69 

70 runner = get_dag_runner(config) 

71 

72 def on_progress(step_name: str, status: str, msg: str) -> None: 

73 if status == "started": 

74 formatter.print_info(msg) 

75 elif status == "running": 

76 formatter.print_info(f" [{step_name}] {msg}") 

77 elif status == "succeeded": 

78 formatter.print_success(f" [{step_name}] ✓ Completed") 

79 elif status == "failed": 

80 formatter.print_error(f" [{step_name}] ✗ {msg}") 

81 elif status == "skipped": 

82 formatter.print_info(f" [{step_name}] ⊘ {msg}") 

83 elif "completed" in status: 

84 print() 

85 formatter.print_info(msg) 

86 

87 result = runner.run( 

88 dag_def, 

89 region=region, 

90 timeout_per_step=timeout, 

91 progress_callback=on_progress, 

92 ) 

93 

94 if result.has_failures(): 

95 sys.exit(1) 

96 

97 except Exception as e: 

98 formatter.print_error(f"DAG execution failed: {e}") 

99 sys.exit(1) 

100 

101 

102@dag.command("validate") 

103@click.argument("dag_file", type=click.Path(exists=True)) 

104@pass_config 

105def dag_validate(config: Any, dag_file: Any) -> None: 

106 """Validate a DAG definition without running it. 

107 

108 Examples: 

109 gco dag validate pipeline.yaml 

110 """ 

111 from ..dag import load_dag 

112 

113 formatter = get_output_formatter(config) 

114 

115 try: 

116 dag_def = load_dag(dag_file) 

117 errors = dag_def.validate() 

118 

119 if errors: 

120 formatter.print_error(f"DAG '{dag_def.name}' has {len(errors)} error(s):") 

121 for err in errors: 

122 formatter.print_error(f" - {err}") 

123 sys.exit(1) 

124 

125 formatter.print_success(f"DAG '{dag_def.name}' is valid") 

126 print(f" Steps: {len(dag_def.steps)}") 

127 print(f" Region: {dag_def.region or '(auto-detect)'}") 

128 print(f" Namespace: {dag_def.namespace}") 

129 for step in dag_def.steps: 

130 deps = f" → depends on: {', '.join(step.depends_on)}" if step.depends_on else "" 

131 print(f" - {step.name}: {step.manifest}{deps}") 

132 print() 

133 

134 except Exception as e: 

135 formatter.print_error(f"Failed to load DAG: {e}") 

136 sys.exit(1)