Coverage for python/pyairflowtester/rules/dag_advanced.py: 95%
181 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-09-19 20:43 +0530
1"""
2Advanced DAG analysis rules (AFW005-AFW015).
4These rules detect more subtle anti-patterns and configuration issues.
5"""
7import re
8from typing import Any, Dict, List
11class BaseRule:
12 """Base rule class."""
14 def __init__(self):
15 self.id = ""
16 self.name = ""
17 self.severity = ""
18 self.category = ""
19 self.execution_mode = ""
21 def evaluate(self, source_code: str, file_name: str = "") -> List[Dict[str, Any]]:
22 """Evaluate rule against source code."""
23 raise NotImplementedError
26class TaskCountRule(BaseRule):
27 """Detect DAGs with excessive task counts."""
29 def __init__(self):
30 super().__init__()
31 self.id = "AFW005"
32 self.name = "Excessive Task Count"
33 self.severity = "medium"
34 self.category = "performance"
35 self.execution_mode = "static"
37 def evaluate(self, source_code: str, file_name: str = "") -> List[Dict[str, Any]]:
38 """Detect DAGs with >500 tasks."""
39 violations = []
41 # Count task definitions
42 task_patterns = [
43 r"@task\(",
44 r"@task_group\(",
45 r"PythonOperator\(",
46 r"BashOperator\(",
47 r"SqlOperator\(",
48 ]
50 task_count = sum(len(re.findall(pattern, source_code)) for pattern in task_patterns)
52 if task_count > 500:
53 violations.append(
54 {
55 "rule_id": self.id,
56 "severity": self.severity,
57 "affected_resource": file_name,
58 "message": f"DAG has {task_count} tasks (exceeds 500 limit)",
59 "remediation": "Split into multiple DAGs or use SubDAGs/task groups",
60 }
61 )
63 return violations
66class CatchupConfigRule(BaseRule):
67 """Detect problematic catchup configurations."""
69 def __init__(self):
70 super().__init__()
71 self.id = "AFW006"
72 self.name = "Risky Catchup Configuration"
73 self.severity = "high"
74 self.category = "reliability"
75 self.execution_mode = "static"
77 def evaluate(self, source_code: str, file_name: str = "") -> List[Dict[str, Any]]:
78 """Detect catchup=True without backfill strategy."""
79 violations = []
81 if re.search(r"catchup\s*=\s*True", source_code):
82 if not re.search(r"backfill|batch_run", source_code):
83 violations.append(
84 {
85 "rule_id": self.id,
86 "severity": self.severity,
87 "affected_resource": file_name,
88 "message": "catchup=True enabled without backfill strategy",
89 "remediation": "Add backfill logic or set catchup=False for production",
90 }
91 )
93 return violations
96class SourceCodePoolConfigurationRule(BaseRule):
97 """Detect default pool usage in DAG source code (AFW007)."""
99 def __init__(self):
100 super().__init__()
101 self.id = "AFW007"
102 self.name = "Default Pool Usage"
103 self.severity = "medium"
104 self.category = "reliability"
105 self.execution_mode = "static"
107 def evaluate(self, source_code: str, file_name: str = "") -> List[Dict[str, Any]]:
108 """Detect tasks using default pool (128 slot limit)."""
109 violations = []
111 # Count tasks with explicit pool vs. without
112 task_with_pool = len(re.findall(r"pool\s*=\s*['\"]", source_code))
113 all_tasks = len(re.findall(r"Operator\(", source_code))
115 if task_with_pool == 0 and all_tasks > 0:
116 violations.append(
117 {
118 "rule_id": self.id,
119 "severity": self.severity,
120 "affected_resource": file_name,
121 "message": "Tasks using default pool (limited to 128 slots)",
122 "remediation": "Create dedicated pool for this DAG or set explicit pool",
123 }
124 )
126 return violations
129class HardcodedConnectionRule(BaseRule):
130 """Detect hardcoded connection IDs."""
132 def __init__(self):
133 super().__init__()
134 self.id = "AFW008"
135 self.name = "Hardcoded Connection ID"
136 self.severity = "high"
137 self.category = "maintainability"
138 self.execution_mode = "static"
140 def evaluate(self, source_code: str, file_name: str = "") -> List[Dict[str, Any]]:
141 """Detect hardcoded connection strings."""
142 violations = []
144 # Look for common hardcoded patterns
145 patterns = [
146 r"conn_id\s*=\s*['\"].*@.*['\"]", # User@host pattern
147 r"connection_string\s*=\s*['\"].*://.*['\"]", # URL pattern
148 r"host\s*=\s*['\"]localhost['\"]",
149 ]
151 for pattern in patterns:
152 if re.search(pattern, source_code):
153 violations.append(
154 {
155 "rule_id": self.id,
156 "severity": self.severity,
157 "affected_resource": file_name,
158 "message": "Hardcoded connection ID or host detected",
159 "remediation": "Use environment variables or Airflow variables",
160 }
161 )
162 break
164 return violations
167class SecretsInCodeRule(BaseRule):
168 """Detect secrets in DAG code."""
170 def __init__(self):
171 super().__init__()
172 self.id = "AFW009"
173 self.name = "Secrets in Code"
174 self.severity = "critical"
175 self.category = "security"
176 self.execution_mode = "static"
178 def evaluate(self, source_code: str, file_name: str = "") -> List[Dict[str, Any]]:
179 """Detect hardcoded secrets."""
180 violations = []
182 secret_patterns = [
183 (r"password\s*=\s*['\"][^'\"]+['\"]", "password"),
184 (r"api_key\s*=\s*['\"][^'\"]+['\"]", "api_key"),
185 (r"token\s*=\s*['\"][^'\"]+['\"]", "token"),
186 (r"secret\s*=\s*['\"][^'\"]+['\"]", "secret"),
187 ]
189 for pattern, secret_type in secret_patterns:
190 if re.search(pattern, source_code):
191 violations.append(
192 {
193 "rule_id": self.id,
194 "severity": self.severity,
195 "affected_resource": file_name,
196 "message": f"Hardcoded {secret_type} detected in source code",
197 "remediation": (
198 "Move to environment variables, AWS Secrets Manager, "
199 "or Airflow Variables"
200 ),
201 }
202 )
203 break
205 return violations
208class RetryConfigurationRule(BaseRule):
209 """Detect problematic retry configurations."""
211 def __init__(self):
212 super().__init__()
213 self.id = "AFW010"
214 self.name = "Excessive Retries"
215 self.severity = "medium"
216 self.category = "reliability"
217 self.execution_mode = "static"
219 def evaluate(self, source_code: str, file_name: str = "") -> List[Dict[str, Any]]:
220 """Detect retries > 5 without backoff."""
221 violations = []
223 # Find retry counts
224 retry_matches = re.findall(r"retries\s*=\s*(\d+)", source_code)
226 for retry_count in retry_matches:
227 if int(retry_count) > 5: 227 ↛ 226line 227 didn't jump to line 226 because the condition on line 227 was always true
228 if "exponential_backoff" not in source_code and "retry_delay" not in source_code:
229 violations.append(
230 {
231 "rule_id": self.id,
232 "severity": self.severity,
233 "affected_resource": file_name,
234 "message": f"High retry count ({retry_count}) without backoff strategy",
235 "remediation": "Add exponential_backoff or configure retry_delay",
236 }
237 )
238 break
240 return violations
243class SensorTimeoutRule(BaseRule):
244 """Detect sensors with problematic timeout settings."""
246 def __init__(self):
247 super().__init__()
248 self.id = "AFW011"
249 self.name = "Sensor Timeout Risk"
250 self.severity = "medium"
251 self.category = "reliability"
252 self.execution_mode = "static"
254 def evaluate(self, source_code: str, file_name: str = "") -> List[Dict[str, Any]]:
255 """Detect sensors with >1h timeout."""
256 violations = []
258 if "Sensor" in source_code or "sensor" in source_code:
259 timeout_matches = re.findall(r"timeout\s*=\s*(\d+)", source_code)
261 for timeout_val in timeout_matches:
262 if int(timeout_val) > 3600: # 1 hour in seconds
263 violations.append(
264 {
265 "rule_id": self.id,
266 "severity": self.severity,
267 "affected_resource": file_name,
268 "message": (
269 f"Sensor timeout > 1 hour ({timeout_val}s) "
270 "without exponential backoff"
271 ),
272 "remediation": "Use exponential_backoff and poke_interval for sensors",
273 }
274 )
275 break
277 return violations
280class BranchComplexityRule(BaseRule):
281 """Detect complex branching logic."""
283 def __init__(self):
284 super().__init__()
285 self.id = "AFW012"
286 self.name = "Complex Branching"
287 self.severity = "low"
288 self.category = "maintainability"
289 self.execution_mode = "static"
291 def evaluate(self, source_code: str, file_name: str = "") -> List[Dict[str, Any]]:
292 """Detect BranchPythonOperator with complex logic."""
293 violations = []
295 if "BranchPythonOperator" in source_code: 295 ↛ 297line 295 didn't jump to line 297 because the condition on line 295 was never true
296 # Count lines in python_callable to estimate complexity
297 branch_blocks = re.findall(
298 r"def\s+\w+\(.*?\):(.+?)(?=\n(?:def|class|\Z))", source_code, re.DOTALL
299 )
301 for block in branch_blocks:
302 line_count = block.count("\n")
303 if line_count > 30:
304 violations.append(
305 {
306 "rule_id": self.id,
307 "severity": self.severity,
308 "affected_resource": file_name,
309 "message": "BranchPythonOperator has complex logic (>30 lines)",
310 "remediation": (
311 "Simplify branching logic or split into separate operations"
312 ),
313 }
314 )
315 break
317 return violations
320class DocumentationRule(BaseRule):
321 """Detect missing DAG documentation."""
323 def __init__(self):
324 super().__init__()
325 self.id = "AFW013"
326 self.name = "Missing Documentation"
327 self.severity = "low"
328 self.category = "maintainability"
329 self.execution_mode = "static"
331 def evaluate(self, source_code: str, file_name: str = "") -> List[Dict[str, Any]]:
332 """Detect DAGs without description."""
333 violations = []
335 if not re.search(r"description\s*=\s*['\"]", source_code):
336 violations.append(
337 {
338 "rule_id": self.id,
339 "severity": self.severity,
340 "affected_resource": file_name,
341 "message": "DAG missing description",
342 "remediation": "Add description parameter to DAG definition",
343 }
344 )
346 return violations
349class AlertingConfigurationRule(BaseRule):
350 """Detect missing alerting."""
352 def __init__(self):
353 super().__init__()
354 self.id = "AFW014"
355 self.name = "No Alerting Configured"
356 self.severity = "medium"
357 self.category = "reliability"
358 self.execution_mode = "static"
360 def evaluate(self, source_code: str, file_name: str = "") -> List[Dict[str, Any]]:
361 """Detect production DAGs without alerts."""
362 violations = []
364 if "production" in file_name.lower() or "prod" in file_name.lower():
365 if not re.search(r"on_failure_callback|on_retry_callback|email", source_code):
366 violations.append(
367 {
368 "rule_id": self.id,
369 "severity": self.severity,
370 "affected_resource": file_name,
371 "message": "Production DAG has no failure alerting configured",
372 "remediation": (
373 "Add on_failure_callback, on_retry_callback, or email notifications"
374 ),
375 }
376 )
378 return violations
381class OperatorDeprecationRule(BaseRule):
382 """Detect deprecated operators."""
384 def __init__(self):
385 super().__init__()
386 self.id = "AFW015"
387 self.name = "Deprecated Operator"
388 self.severity = "medium"
389 self.category = "maintainability"
390 self.execution_mode = "static"
392 def evaluate(self, source_code: str, file_name: str = "") -> List[Dict[str, Any]]:
393 """Detect deprecated Airflow operators."""
394 violations = []
396 deprecated = {
397 "SubDagOperator": "Use task groups instead (@task_group)",
398 "DummyOperator": "Use EmptyOperator instead",
399 "BranchPythonOperator": "Consider using TaskFlow API",
400 }
402 for operator, recommendation in deprecated.items():
403 if operator in source_code:
404 violations.append(
405 {
406 "rule_id": self.id,
407 "severity": self.severity,
408 "affected_resource": operator,
409 "message": f"Deprecated operator used: {operator}",
410 "remediation": recommendation,
411 }
412 )
413 break
415 return violations