Coverage for python/pyairflowtester/dependency_intelligence/parsers.py: 72%

1"""Parsers for different dependency sources (Airflow, dbt, datasets).""" 

2 

3import ast 

4import json 

5import logging 

6from pathlib import Path 

7from typing import Any, Dict, List, Optional, Tuple 

8 

9from .models import ( 

10 DependencyGraph, 

11 Edge, 

12 Node, 

13 NodeType, 

14 RelationshipType, 

15) 

16 

17logger = logging.getLogger(__name__) 

18 

19 

20class AirflowDAGParser: 

21 """Parse Airflow DAG files to extract dependencies.""" 

22 

23 @staticmethod 

24 def parse_dag_file(file_path: str) -> Tuple[Optional[str], List[str], List[Tuple[str, str]]]: 

25 """ 

26 Parse a Python DAG file and extract DAG ID, task IDs, and dependencies. 

27 

28 Returns: 

29 Tuple of (dag_id, task_ids, dependencies) 

30 """ 

31 try: 

32 with open(file_path, "r") as f: 

33 source_code = f.read() 

34 return AirflowDAGParser.parse_dag_code(source_code) 

35 except Exception as e: 

36 logger.error(f"Error parsing DAG file {file_path}: {e}") 

37 return None, [], [] 

38 

39 @staticmethod 

40 def parse_dag_code(source_code: str) -> Tuple[Optional[str], List[str], List[Tuple[str, str]]]: 

41 """ 

42 Parse DAG Python code and extract DAG ID, tasks, and dependencies. 

43 

44 Returns: 

45 Tuple of (dag_id, task_ids, dependencies) 

46 """ 

47 try: 

48 tree = ast.parse(source_code) 

49 except SyntaxError as e: 

50 logger.error(f"Syntax error parsing DAG code: {e}") 

51 return None, [], [] 

52 

53 dag_id = None 

54 task_ids = set() 

55 dependencies = [] 

56 

57 class DAGVisitor(ast.NodeVisitor): 

58 def visit_Call(self, node): 

59 nonlocal dag_id 

60 

61 # Look for DAG() instantiation 

62 if isinstance(node.func, ast.Name) and node.func.id == "DAG": 

63 # dag_id is idiomatically passed positionally: DAG('my_dag', ...) 

64 if node.args and isinstance(node.args[0], ast.Constant): 

65 dag_id = node.args[0].value 

66 elif node.args and isinstance(node.args[0], ast.Str): 66 ↛ 67line 66 didn't jump to line 67 because the condition on line 66 was never true

67 dag_id = node.args[0].s 

68 

69 for keyword in node.keywords: 

70 if keyword.arg == "dag_id" and isinstance(keyword.value, ast.Constant): 

71 dag_id = keyword.value.value 

72 elif keyword.arg == "dag_id" and isinstance(keyword.value, ast.Str): 72 ↛ 73line 72 didn't jump to line 73 because the condition on line 72 was never true

73 dag_id = keyword.value.s 

74 

75 # Look for task assignments 

76 if isinstance(node.func, ast.Name) and "Operator" in node.func.id: 

77 for keyword in node.keywords: 

78 if keyword.arg == "task_id" and isinstance(keyword.value, ast.Constant): 

79 task_ids.add(keyword.value.value) 

80 elif keyword.arg == "task_id" and isinstance(keyword.value, ast.Str): 80 ↛ 81line 80 didn't jump to line 81 because the condition on line 80 was never true

81 task_ids.add(keyword.value.s) 

82 

83 # Look for task dependencies (set_upstream/set_downstream) 

84 if isinstance(node.func, ast.Attribute): 

85 if node.func.attr in ("set_upstream", "set_downstream"): 85 ↛ 95line 85 didn't jump to line 95 because the condition on line 85 was always true

86 if isinstance(node.func.value, ast.Name): 86 ↛ 95line 86 didn't jump to line 95 because the condition on line 86 was always true

87 task1 = node.func.value.id 

88 if len(node.args) > 0 and isinstance(node.args[0], ast.Name): 88 ↛ 95line 88 didn't jump to line 95 because the condition on line 88 was always true

89 task2 = node.args[0].id 

90 if node.func.attr == "set_upstream": 90 ↛ 91line 90 didn't jump to line 91 because the condition on line 90 was never true

91 dependencies.append((task2, task1)) 

92 else: 

93 dependencies.append((task1, task2)) 

94 

95 self.generic_visit(node) 

96 

97 DAGVisitor().visit(tree) 

98 

99 return dag_id, sorted(list(task_ids)), dependencies 

100 

101 @staticmethod 

102 def build_graph(dag_files: List[str]) -> DependencyGraph: 

103 """ 

104 Build a dependency graph from multiple DAG files. 

105 

106 Args: 

107 dag_files: List of paths to DAG Python files 

108 

109 Returns: 

110 DependencyGraph with all DAGs, tasks, and dependencies 

111 """ 

112 graph = DependencyGraph() 

113 all_dependencies = [] 

114 

115 for dag_file in dag_files: 

116 dag_id, task_ids, dependencies = AirflowDAGParser.parse_dag_file(dag_file) 

117 

118 if dag_id: 118 ↛ 115line 118 didn't jump to line 115 because the condition on line 118 was always true

119 # Add DAG node 

120 dag_node = Node( 

121 id=f"dag_{dag_id}", 

122 name=dag_id, 

123 type=NodeType.DAG, 

124 owner="airflow", 

125 ) 

126 graph.add_node(dag_node) 

127 

128 # Add task nodes 

129 for task_id in task_ids: 129 ↛ 130line 129 didn't jump to line 130 because the loop on line 129 never started

130 task_node = Node( 

131 id=f"task_{dag_id}_{task_id}", 

132 name=f"{dag_id}.{task_id}", 

133 type=NodeType.TASK, 

134 owner="airflow", 

135 metadata={"dag_id": dag_id, "task_id": task_id}, 

136 ) 

137 graph.add_node(task_node) 

138 

139 # Add edge from DAG to task 

140 edge = Edge( 

141 source=f"dag_{dag_id}", 

142 target=f"task_{dag_id}_{task_id}", 

143 relationship_type=RelationshipType.CALLS, 

144 ) 

145 graph.add_edge(edge) 

146 

147 # Store dependencies for later processing 

148 for source, target in dependencies: 148 ↛ 149line 148 didn't jump to line 149 because the loop on line 148 never started

149 all_dependencies.append((dag_id, source, target)) 

150 

151 # Add task-to-task dependencies 

152 for dag_id, source_task, target_task in all_dependencies: 152 ↛ 153line 152 didn't jump to line 153 because the loop on line 152 never started

153 source_id = f"task_{dag_id}_{source_task}" 

154 target_id = f"task_{dag_id}_{target_task}" 

155 

156 if source_id in graph.nodes and target_id in graph.nodes: 

157 edge = Edge( 

158 source=source_id, 

159 target=target_id, 

160 relationship_type=RelationshipType.DEPENDS_ON, 

161 ) 

162 graph.add_edge(edge) 

163 

164 return graph 

165 

166 

167class dbtManifestParser: 

168 """Parse dbt manifest.json to extract model lineage.""" 

169 

170 @staticmethod 

171 def parse_manifest(manifest_path: str) -> DependencyGraph: 

172 """ 

173 Parse dbt manifest.json and build dependency graph. 

174 

175 Args: 

176 manifest_path: Path to manifest.json 

177 

178 Returns: 

179 DependencyGraph with dbt models and lineage 

180 """ 

181 graph = DependencyGraph() 

182 

183 try: 

184 with open(manifest_path, "r") as f: 

185 manifest = json.load(f) 

186 except Exception as e: 

187 logger.error(f"Error reading manifest: {e}") 

188 return graph 

189 

190 # Parse nodes (models, tests, sources, etc.) 

191 nodes_data = manifest.get("nodes", {}) 

192 

193 for node_id, node_data in nodes_data.items(): 

194 # Determine node type 

195 if "model." in node_id: 

196 node_type = NodeType.DBT_MODEL 

197 elif "test." in node_id: 

198 node_type = NodeType.DBT_TEST 

199 elif "source." in node_id: 199 ↛ 201line 199 didn't jump to line 201 because the condition on line 199 was always true

200 node_type = NodeType.DBT_SOURCE 

201 elif "snapshot." in node_id: 

202 node_type = NodeType.DBT_SNAPSHOT 

203 elif "exposure." in node_id: 

204 node_type = NodeType.DBT_EXPOSURE 

205 else: 

206 continue 

207 

208 # Create node 

209 node = Node( 

210 id=node_id, 

211 name=node_data.get("name", node_id), 

212 type=node_type, 

213 owner=node_data.get("meta", {}).get("owner", "dbt"), 

214 description=node_data.get("description", ""), 

215 metadata={ 

216 "package": node_data.get("package_name", ""), 

217 "database": node_data.get("database", ""), 

218 "schema": node_data.get("schema", ""), 

219 "tags": node_data.get("tags", []), 

220 }, 

221 ) 

222 graph.add_node(node) 

223 

224 # Parse dependencies (edges) 

225 for node_id, node_data in nodes_data.items(): 

226 depends_on = node_data.get("depends_on", {}).get("nodes", []) 

227 

228 for dependency_id in depends_on: 

229 if dependency_id in graph.nodes: 229 ↛ 228line 229 didn't jump to line 228 because the condition on line 229 was always true

230 edge = Edge( 

231 source=dependency_id, 

232 target=node_id, 

233 relationship_type=RelationshipType.DEPENDS_ON, 

234 ) 

235 graph.add_edge(edge) 

236 

237 # Parse exposures 

238 exposures = manifest.get("exposures", {}) 

239 for exposure_id, exposure_data in exposures.items(): 

240 node = Node( 

241 id=exposure_id, 

242 name=exposure_data.get("name", exposure_id), 

243 type=NodeType.DBT_EXPOSURE, 

244 owner=exposure_data.get("meta", {}).get("owner", "dbt"), 

245 description=exposure_data.get("description", ""), 

246 ) 

247 graph.add_node(node) 

248 

249 for dependency_id in exposure_data.get("depends_on", {}).get("nodes", []): 

250 if dependency_id in graph.nodes: 250 ↛ 249line 250 didn't jump to line 249 because the condition on line 250 was always true

251 edge = Edge( 

252 source=dependency_id, 

253 target=exposure_id, 

254 relationship_type=RelationshipType.EXPOSES, 

255 ) 

256 graph.add_edge(edge) 

257 

258 return graph 

259 

260 @staticmethod 

261 def parse_model_node(node_data: Dict[str, Any]) -> Node: 

262 """Parse a single dbt model node.""" 

263 return Node( 

264 id=node_data.get("unique_id"), 

265 name=node_data.get("name"), 

266 type=NodeType.DBT_MODEL, 

267 owner=node_data.get("meta", {}).get("owner", "dbt"), 

268 description=node_data.get("description", ""), 

269 metadata={ 

270 "materialized": node_data.get("config", {}).get("materialized"), 

271 "tags": node_data.get("tags", []), 

272 }, 

273 ) 

274 

275 

276class AirflowDatasetParser: 

277 """Parse Airflow dataset definitions and dependencies.""" 

278 

279 @staticmethod 

280 def parse_dataset_connections(source_code: str) -> Tuple[List[str], List[Tuple[str, str]]]: 

281 """ 

282 Parse Airflow code for dataset producers and consumers. 

283 

284 Returns: 

285 Tuple of (dataset_ids, dependencies) 

286 """ 

287 try: 

288 tree = ast.parse(source_code) 

289 except SyntaxError: 

290 return [], [] 

291 

292 datasets = set() 

293 dependencies = [] 

294 

295 class DatasetVisitor(ast.NodeVisitor): 

296 def visit_Call(self, node): 

297 # Look for Dataset() instantiation 

298 if isinstance(node.func, ast.Name) and node.func.id == "Dataset": 

299 for keyword in node.keywords: 299 ↛ 300line 299 didn't jump to line 300 because the loop on line 299 never started

300 if keyword.arg == "uri" and isinstance(keyword.value, ast.Constant): 

301 datasets.add(keyword.value.value) 

302 

303 # Look for dataset_triggers 

304 if isinstance(node.func, ast.Name) and node.func.id == "DAG": 

305 for keyword in node.keywords: 

306 if keyword.arg == "start_date": 306 ↛ 307line 306 didn't jump to line 307 because the condition on line 306 was never true

307 pass # Found DAG start_date 

308 

309 self.generic_visit(node) 

310 

311 DatasetVisitor().visit(tree) 

312 

313 return sorted(list(datasets)), dependencies 

314 

315 @staticmethod 

316 def build_dataset_graph(airflow_files: List[str]) -> DependencyGraph: 

317 """ 

318 Build dataset dependency graph from Airflow files. 

319 

320 Args: 

321 airflow_files: List of paths to Airflow Python files 

322 

323 Returns: 

324 DependencyGraph with datasets as nodes 

325 """ 

326 graph = DependencyGraph() 

327 

328 for airflow_file in airflow_files: 

329 try: 

330 with open(airflow_file, "r") as f: 

331 source_code = f.read() 

332 

333 datasets, deps = AirflowDatasetParser.parse_dataset_connections(source_code) 

334 

335 for dataset_uri in datasets: 

336 node = Node( 

337 id=f"dataset_{hash(dataset_uri)}", 

338 name=dataset_uri, 

339 type=NodeType.DATASET, 

340 owner="airflow", 

341 metadata={"uri": dataset_uri}, 

342 ) 

343 graph.add_node(node) 

344 

345 except Exception as e: 

346 logger.error(f"Error parsing dataset file {airflow_file}: {e}") 

347 

348 return graph 

349 

350 

351class UnifiedGraphBuilder: 

352 """Build unified graph from multiple sources (Airflow + dbt + datasets).""" 

353 

354 @staticmethod 

355 def build_unified_graph( 

356 dag_files: List[str] = None, 

357 dbt_manifest: str = None, 

358 dataset_files: List[str] = None, 

359 ) -> DependencyGraph: 

360 """ 

361 Build a unified dependency graph from multiple sources. 

362 

363 Returns: 

364 Unified DependencyGraph 

365 """ 

366 unified_graph = DependencyGraph() 

367 

368 # Parse Airflow DAGs 

369 if dag_files: 

370 dag_graph = AirflowDAGParser.build_graph(dag_files) 

371 unified_graph.nodes.update(dag_graph.nodes) 

372 unified_graph.edges.extend(dag_graph.edges) 

373 

374 # Parse dbt manifest 

375 if dbt_manifest and Path(dbt_manifest).exists(): 

376 dbt_graph = dbtManifestParser.parse_manifest(dbt_manifest) 

377 unified_graph.nodes.update(dbt_graph.nodes) 

378 unified_graph.edges.extend(dbt_graph.edges) 

379 

380 # Parse datasets 

381 if dataset_files: 381 ↛ 382line 381 didn't jump to line 382 because the condition on line 381 was never true

382 dataset_graph = AirflowDatasetParser.build_dataset_graph(dataset_files) 

383 unified_graph.nodes.update(dataset_graph.nodes) 

384 unified_graph.edges.extend(dataset_graph.edges) 

385 

386 return unified_graph