Coverage for python/pyairflowtester/dependency_intelligence/runtime_import.py: 86%

1"""Sandboxed runtime-import fallback for dynamically-generated DAGs. 

2 

3`AirflowDAGParser` (parsers.py) is pure AST pattern-matching: it looks for 

4literal `DAG(...)`/`*Operator(...)` calls and `set_upstream`/`set_downstream` 

5attribute calls in the source text. DAGs built via factory functions, 

6dynamic loops, or `exec`/`eval` don't have any of those literal shapes in 

7their source -- the task graph only exists after the module actually runs. 

8 

9This module is the fallback for that case: import the DAG file for real, in 

10an isolated subprocess (so a buggy or malicious DAG file can't corrupt or 

11hang this process), and read back whatever real `airflow.models.DAG`/task 

12objects resulted -- which, having actually executed, have their `>>`/ 

13`set_upstream` wiring fully resolved regardless of how dynamically the 

14tasks were constructed. 

15 

16Per this project's own design (see README/CLAUDE.md), Airflow is not a hard 

17runtime dependency -- static analysis works without it installed. This 

18fallback is opt-in and requires it: if `airflow` isn't importable in the 

19sandboxed subprocess, `parse_dag_via_runtime_import` reports that plainly 

20rather than guessing at a dependency graph with no framework to resolve it 

21against. 

22""" 

23 

24from __future__ import annotations 

25 

26import json 

27import subprocess 

28import sys 

29from pathlib import Path 

30from typing import List, Optional, Tuple 

31 

32# Memory ceiling for the sandboxed subprocess. Real Airflow DAG modules doing 

33# import-time work (reading connections, building large task fan-outs) can 

34# legitimately need more than a few dozen MB, but this still bounds a runaway 

35# or hostile DAG file from consuming unbounded memory on the host. 

36_DEFAULT_MEMORY_LIMIT_BYTES = 512 * 1024 * 1024 

37_DEFAULT_TIMEOUT_SECONDS = 30.0 

38 

39_SANDBOX_SCRIPT = r""" 

40import importlib.util 

41import json 

42import sys 

43 

44try: 

45 import resource 

46 _soft, _hard = resource.getrlimit(resource.RLIMIT_AS) 

47 _limit = {memory_limit_bytes} 

48 _hard_limit = _hard if _hard != resource.RLIM_INFINITY else _limit 

49 resource.setrlimit(resource.RLIMIT_AS, (_limit, _hard_limit)) 

50except Exception: 

51 pass # RLIMIT_AS isn't available everywhere; degrade to timeout-only sandboxing. 

52 

53result = {{"dag_id": None, "task_ids": [], "dependencies": [], "error": None}} 

54 

55try: 

56 import airflow # noqa: F401 

57 from airflow.models import DAG 

58except ImportError: 

59 result["error"] = "airflow_not_installed" 

60 print(json.dumps(result)) 

61 sys.exit(0) 

62 

63file_path = {file_path!r} 

64 

65try: 

66 spec = importlib.util.spec_from_file_location("_pyairflowtester_sandboxed_dag", file_path) 

67 module = importlib.util.module_from_spec(spec) 

68 spec.loader.exec_module(module) 

69except Exception as exc: 

70 result["error"] = f"import_failed: {{type(exc).__name__}}: {{exc}}" 

71 print(json.dumps(result)) 

72 sys.exit(0) 

73 

74dags = [obj for obj in vars(module).values() if isinstance(obj, DAG)] 

75if not dags: 

76 result["error"] = "no_dag_found_after_import" 

77 print(json.dumps(result)) 

78 sys.exit(0) 

79 

80dag = dags[0] 

81result["dag_id"] = dag.dag_id 

82result["task_ids"] = sorted(dag.task_dict.keys()) 

83 

84dependencies = [] 

85for task_id, task in dag.task_dict.items(): 

86 for upstream_id in sorted(getattr(task, "upstream_task_ids", []) or []): 

87 dependencies.append([upstream_id, task_id]) 

88result["dependencies"] = dependencies 

89 

90print(json.dumps(result)) 

91""" 

92 

93 

94def parse_dag_via_runtime_import( 

95 file_path: str, 

96 timeout_seconds: float = _DEFAULT_TIMEOUT_SECONDS, 

97 memory_limit_bytes: int = _DEFAULT_MEMORY_LIMIT_BYTES, 

98) -> Tuple[Optional[str], List[str], List[Tuple[str, str]]]: 

99 """Import `file_path` in a sandboxed subprocess and extract the real, 

100 fully-resolved DAG task graph -- the fallback for DAGs 

101 `AirflowDAGParser`'s static AST parsing can't see into. 

102 

103 Returns the same `(dag_id, task_ids, dependencies)` shape as 

104 `AirflowDAGParser.parse_dag_code`, so this is a drop-in fallback: call 

105 the static parser first, and only fall back to this (slower, requires 

106 `airflow` installed, actually executes the file) when it comes back 

107 empty or incomplete. 

108 

109 Never raises for expected failure modes (airflow not installed, the 

110 file failing to import, no DAG object found, timeout) -- returns 

111 `(None, [], [])` and logs instead, since a fallback that can itself 

112 crash the caller isn't a usable fallback. 

113 """ 

114 import logging 

115 

116 logger = logging.getLogger(__name__) 

117 

118 resolved_path = str(Path(file_path).resolve()) 

119 script = _SANDBOX_SCRIPT.format(file_path=resolved_path, memory_limit_bytes=memory_limit_bytes) 

120 

121 try: 

122 proc = subprocess.run( 

123 [sys.executable, "-c", script], 

124 capture_output=True, 

125 text=True, 

126 timeout=timeout_seconds, 

127 ) 

128 except subprocess.TimeoutExpired: 

129 logger.warning( 

130 "Runtime-import fallback timed out after %ss parsing %s", timeout_seconds, file_path 

131 ) 

132 return None, [], [] 

133 

134 if proc.returncode != 0 or not proc.stdout.strip(): 134 ↛ 135line 134 didn't jump to line 135 because the condition on line 134 was never true

135 logger.warning( 

136 "Runtime-import fallback failed for %s (exit %s): %s", 

137 file_path, 

138 proc.returncode, 

139 proc.stderr.strip()[-2000:], 

140 ) 

141 return None, [], [] 

142 

143 try: 

144 payload = json.loads(proc.stdout.strip().splitlines()[-1]) 

145 except (json.JSONDecodeError, IndexError): 

146 logger.warning("Runtime-import fallback produced unparseable output for %s", file_path) 

147 return None, [], [] 

148 

149 if payload.get("error"): 

150 logger.info("Runtime-import fallback for %s: %s", file_path, payload["error"]) 

151 return None, [], [] 

152 

153 dependencies = [(pair[0], pair[1]) for pair in payload.get("dependencies", [])] 

154 return payload.get("dag_id"), list(payload.get("task_ids", [])), dependencies 

155 

156 

157def parse_dag_file_with_fallback( 

158 file_path: str, 

159 timeout_seconds: float = _DEFAULT_TIMEOUT_SECONDS, 

160) -> Tuple[Optional[str], List[str], List[Tuple[str, str]]]: 

161 """Parse a DAG file with `AirflowDAGParser` first; if that finds no 

162 tasks (the signature of a dynamically-generated DAG the static parser 

163 can't see into), fall back to `parse_dag_via_runtime_import`. 

164 

165 This is the function most callers want -- it gets the AST parser's 

166 speed and zero-dependency behavior for the common case, and only pays 

167 for a sandboxed subprocess import when static analysis genuinely came 

168 up empty. 

169 """ 

170 from .parsers import AirflowDAGParser 

171 

172 dag_id, task_ids, dependencies = AirflowDAGParser.parse_dag_file(file_path) 

173 if task_ids: 

174 return dag_id, task_ids, dependencies 

175 

176 return parse_dag_via_runtime_import(file_path, timeout_seconds=timeout_seconds)