Coverage for python/pyairflowtester/rules/config.py: 90%

1""" 

2Airflow configuration audit rules (CFG001-CFG015). 

3 

4These rules check airflow.cfg and runtime configuration for best practices. 

5""" 

6 

7from typing import Any, Dict, List 

8 

9 

10class BaseRule: 

11 """Base configuration rule.""" 

12 

13 def __init__(self): 

14 self.id = "" 

15 self.name = "" 

16 self.severity = "" 

17 self.category = "" 

18 self.execution_mode = "" 

19 

20 def evaluate(self, config: Dict[str, Any]) -> List[Dict[str, Any]]: 

21 """Evaluate rule against configuration.""" 

22 raise NotImplementedError 

23 

24 

25class ExecutorConfigurationRule(BaseRule): 

26 """Check executor type matches workload.""" 

27 

28 def __init__(self): 

29 super().__init__() 

30 self.id = "CFG001" 

31 self.name = "Executor Misconfiguration" 

32 self.severity = "high" 

33 self.category = "performance" 

34 self.execution_mode = "static" 

35 

36 def evaluate(self, config: Dict[str, Any]) -> List[Dict[str, Any]]: 

37 """Check executor configuration.""" 

38 violations = [] 

39 

40 executor = config.get("core", {}).get("executor", "SequentialExecutor") 

41 

42 # SequentialExecutor only for testing 

43 if executor == "SequentialExecutor": 43 ↛ 55line 43 didn't jump to line 55 because the condition on line 43 was always true

44 violations.append( 

45 { 

46 "rule_id": self.id, 

47 "severity": "high", 

48 "affected_resource": "core.executor", 

49 "message": "Using SequentialExecutor in production (no parallelism)", 

50 "remediation": "Use LocalExecutor, CeleryExecutor, or KubernetesExecutor", 

51 } 

52 ) 

53 

54 # Check KubernetesExecutor has resource limits 

55 if executor == "KubernetesExecutor": 55 ↛ 56line 55 didn't jump to line 56 because the condition on line 55 was never true

56 if not config.get("kubernetes", {}).get("enable_security_context"): 

57 violations.append( 

58 { 

59 "rule_id": self.id, 

60 "severity": "medium", 

61 "affected_resource": "core.executor", 

62 "message": "KubernetesExecutor without security context", 

63 "remediation": "Enable security context and resource limits", 

64 } 

65 ) 

66 

67 return violations 

68 

69 

70class AirflowCfgPoolConfigurationRule(BaseRule): 

71 """Check pool configuration in airflow.cfg (CFG002).""" 

72 

73 def __init__(self): 

74 super().__init__() 

75 self.id = "CFG002" 

76 self.name = "Pool Size Mismatch" 

77 self.severity = "medium" 

78 self.category = "reliability" 

79 self.execution_mode = "static" 

80 

81 def evaluate(self, config: Dict[str, Any]) -> List[Dict[str, Any]]: 

82 """Check pool sizes.""" 

83 violations = [] 

84 

85 default_pool_size = config.get("core", {}).get("default_pool_size", 128) 

86 max_active_runs = config.get("core", {}).get("max_active_dag_runs", 16) 

87 

88 if int(default_pool_size) < int(max_active_runs) * 5: 88 ↛ 89line 88 didn't jump to line 89 because the condition on line 88 was never true

89 violations.append( 

90 { 

91 "rule_id": self.id, 

92 "severity": self.severity, 

93 "affected_resource": "core.default_pool_size", 

94 "message": ( 

95 f"Pool size ({default_pool_size}) too small for " 

96 f"max runs ({max_active_runs})" 

97 ), 

98 "remediation": "Increase default_pool_size or reduce max_active_dag_runs", 

99 } 

100 ) 

101 

102 return violations 

103 

104 

105class ConcurrencyConfigurationRule(BaseRule): 

106 """Check DAG/task concurrency settings.""" 

107 

108 def __init__(self): 

109 super().__init__() 

110 self.id = "CFG003" 

111 self.name = "Concurrency Mismatch" 

112 self.severity = "medium" 

113 self.category = "performance" 

114 self.execution_mode = "static" 

115 

116 def evaluate(self, config: Dict[str, Any]) -> List[Dict[str, Any]]: 

117 """Check concurrency configuration.""" 

118 violations = [] 

119 

120 dag_concurrency = int(config.get("core", {}).get("dag_concurrency", 16)) 

121 parallelism = int(config.get("core", {}).get("parallelism", 32)) 

122 

123 if dag_concurrency > parallelism: 123 ↛ 124line 123 didn't jump to line 124 because the condition on line 123 was never true

124 violations.append( 

125 { 

126 "rule_id": self.id, 

127 "severity": self.severity, 

128 "affected_resource": "core.dag_concurrency", 

129 "message": ( 

130 f"dag_concurrency ({dag_concurrency}) exceeds parallelism ({parallelism})" 

131 ), 

132 "remediation": "Reduce dag_concurrency or increase parallelism", 

133 } 

134 ) 

135 

136 return violations 

137 

138 

139class QueueConfigurationRule(BaseRule): 

140 """Check queue configuration.""" 

141 

142 def __init__(self): 

143 super().__init__() 

144 self.id = "CFG004" 

145 self.name = "Queue Bottleneck" 

146 self.severity = "medium" 

147 self.category = "reliability" 

148 self.execution_mode = "static" 

149 

150 def evaluate(self, config: Dict[str, Any]) -> List[Dict[str, Any]]: 

151 """Check queue configuration.""" 

152 violations = [] 

153 

154 # Check if using single queue for all DAGs 

155 if "default_queue" in config.get("celery", {}): 155 ↛ 156line 155 didn't jump to line 156 because the condition on line 155 was never true

156 violations.append( 

157 { 

158 "rule_id": self.id, 

159 "severity": self.severity, 

160 "affected_resource": "celery.default_queue", 

161 "message": "Single default queue (no prioritization)", 

162 "remediation": "Create separate queues for priority levels and DAG types", 

163 } 

164 ) 

165 

166 return violations 

167 

168 

169class MaxActiveRunsRule(BaseRule): 

170 """Check max_active_dag_runs.""" 

171 

172 def __init__(self): 

173 super().__init__() 

174 self.id = "CFG005" 

175 self.name = "Max Active Runs Not Limited" 

176 self.severity = "high" 

177 self.category = "reliability" 

178 self.execution_mode = "static" 

179 

180 def evaluate(self, config: Dict[str, Any]) -> List[Dict[str, Any]]: 

181 """Check max active runs limit.""" 

182 violations = [] 

183 

184 max_active = int(config.get("core", {}).get("max_active_dag_runs", 16)) 

185 

186 if max_active > 32: 186 ↛ 187line 186 didn't jump to line 187 because the condition on line 186 was never true

187 violations.append( 

188 { 

189 "rule_id": self.id, 

190 "severity": self.severity, 

191 "affected_resource": "core.max_active_dag_runs", 

192 "message": f"Max active runs too high ({max_active}), risk of runaway backfill", 

193 "remediation": "Reduce max_active_dag_runs to 8-16", 

194 } 

195 ) 

196 

197 return violations 

198 

199 

200class XComConfigurationRule(BaseRule): 

201 """Check XCom backend.""" 

202 

203 def __init__(self): 

204 super().__init__() 

205 self.id = "CFG006" 

206 self.name = "XCom Backend Issue" 

207 self.severity = "medium" 

208 self.category = "performance" 

209 self.execution_mode = "static" 

210 

211 def evaluate(self, config: Dict[str, Any]) -> List[Dict[str, Any]]: 

212 """Check XCom configuration.""" 

213 violations = [] 

214 

215 xcom_backend = config.get("core", {}).get("xcom_backend", "airflow.models.xcom.BaseXCom") 

216 

217 if "BaseXCom" in xcom_backend or "DbXCom" in xcom_backend: 217 ↛ 228line 217 didn't jump to line 228 because the condition on line 217 was always true

218 violations.append( 

219 { 

220 "rule_id": self.id, 

221 "severity": self.severity, 

222 "affected_resource": "core.xcom_backend", 

223 "message": "XCom using database backend (causes bloat)", 

224 "remediation": "Use S3XCom or custom cloud storage backend", 

225 } 

226 ) 

227 

228 return violations 

229 

230 

231class LogRetentionRule(BaseRule): 

232 """Check log retention.""" 

233 

234 def __init__(self): 

235 super().__init__() 

236 self.id = "CFG007" 

237 self.name = "Insufficient Log Retention" 

238 self.severity = "medium" 

239 self.category = "compliance" 

240 self.execution_mode = "static" 

241 

242 def evaluate(self, config: Dict[str, Any]) -> List[Dict[str, Any]]: 

243 """Check log retention.""" 

244 violations = [] 

245 

246 log_retention = int(config.get("logging", {}).get("log_retention_days", 30)) 

247 

248 if log_retention < 30: 248 ↛ 249line 248 didn't jump to line 249 because the condition on line 248 was never true

249 violations.append( 

250 { 

251 "rule_id": self.id, 

252 "severity": self.severity, 

253 "affected_resource": "logging.log_retention_days", 

254 "message": f"Log retention < 30 days ({log_retention}), compliance risk", 

255 "remediation": "Increase log_retention_days to minimum 30-90 days", 

256 } 

257 ) 

258 

259 return violations 

260 

261 

262class EncryptionConfigurationRule(BaseRule): 

263 """Check encryption settings.""" 

264 

265 def __init__(self): 

266 super().__init__() 

267 self.id = "CFG008" 

268 self.name = "Encryption Not Configured" 

269 self.severity = "high" 

270 self.category = "security" 

271 self.execution_mode = "static" 

272 

273 def evaluate(self, config: Dict[str, Any]) -> List[Dict[str, Any]]: 

274 """Check encryption.""" 

275 violations = [] 

276 

277 fernet_key = config.get("core", {}).get("fernet_key") 

278 

279 if not fernet_key or fernet_key == "[NOT CONFIGURED]": 279 ↛ 293line 279 didn't jump to line 293 because the condition on line 279 was always true

280 violations.append( 

281 { 

282 "rule_id": self.id, 

283 "severity": self.severity, 

284 "affected_resource": "core.fernet_key", 

285 "message": "Fernet encryption not configured", 

286 "remediation": ( 

287 "Generate Fernet key: python -c 'from cryptography.fernet import " 

288 "Fernet; print(Fernet.generate_key())'" 

289 ), 

290 } 

291 ) 

292 

293 return violations 

294 

295 

296class TLSConfigurationRule(BaseRule): 

297 """Check TLS settings.""" 

298 

299 def __init__(self): 

300 super().__init__() 

301 self.id = "CFG009" 

302 self.name = "TLS Not Enabled" 

303 self.severity = "high" 

304 self.category = "security" 

305 self.execution_mode = "static" 

306 

307 def evaluate(self, config: Dict[str, Any]) -> List[Dict[str, Any]]: 

308 """Check TLS configuration.""" 

309 violations = [] 

310 

311 enable_ssl = config.get("webserver", {}).get("enable_ssl", False) 

312 

313 if not enable_ssl: 

314 violations.append( 

315 { 

316 "rule_id": self.id, 

317 "severity": self.severity, 

318 "affected_resource": "webserver.enable_ssl", 

319 "message": "TLS/SSL not enabled for Airflow web UI", 

320 "remediation": "Set enable_ssl=True and provide cert_file/key_file", 

321 } 

322 ) 

323 

324 return violations 

325 

326 

327class RBACConfigurationRule(BaseRule): 

328 """Check RBAC settings.""" 

329 

330 def __init__(self): 

331 super().__init__() 

332 self.id = "CFG010" 

333 self.name = "RBAC Disabled" 

334 self.severity = "high" 

335 self.category = "security" 

336 self.execution_mode = "static" 

337 

338 def evaluate(self, config: Dict[str, Any]) -> List[Dict[str, Any]]: 

339 """Check RBAC configuration.""" 

340 violations = [] 

341 

342 auth_backend = config.get("webserver", {}).get("rbac", True) 

343 

344 if not auth_backend: 

345 violations.append( 

346 { 

347 "rule_id": self.id, 

348 "severity": self.severity, 

349 "affected_resource": "webserver.rbac", 

350 "message": "RBAC disabled, no access control", 

351 "remediation": "Enable RBAC and configure authentication backend", 

352 } 

353 ) 

354 

355 return violations 

356 

357 

358class SchedulerConfigurationRule(BaseRule): 

359 """Check scheduler settings.""" 

360 

361 def __init__(self): 

362 super().__init__() 

363 self.id = "CFG011" 

364 self.name = "Scheduler Misconfiguration" 

365 self.severity = "medium" 

366 self.category = "performance" 

367 self.execution_mode = "static" 

368 

369 def evaluate(self, config: Dict[str, Any]) -> List[Dict[str, Any]]: 

370 """Check scheduler configuration.""" 

371 violations = [] 

372 

373 heartbeat_sec = int(config.get("scheduler", {}).get("scheduler_heartbeat_sec", 5)) 

374 

375 if heartbeat_sec < 2: 375 ↛ 376line 375 didn't jump to line 376 because the condition on line 375 was never true

376 violations.append( 

377 { 

378 "rule_id": self.id, 

379 "severity": self.severity, 

380 "affected_resource": "scheduler.scheduler_heartbeat_sec", 

381 "message": f"Scheduler heartbeat too frequent ({heartbeat_sec}s), high CPU", 

382 "remediation": "Increase heartbeat_sec to 5-10 seconds", 

383 } 

384 ) 

385 

386 return violations 

387 

388 

389class WorkerConfigurationRule(BaseRule): 

390 """Check worker settings.""" 

391 

392 def __init__(self): 

393 super().__init__() 

394 self.id = "CFG012" 

395 self.name = "Worker Misconfiguration" 

396 self.severity = "medium" 

397 self.category = "reliability" 

398 self.execution_mode = "static" 

399 

400 def evaluate(self, config: Dict[str, Any]) -> List[Dict[str, Any]]: 

401 """Check worker configuration.""" 

402 violations = [] 

403 

404 prefetch = int(config.get("celery", {}).get("worker_prefetch_multiplier", 1)) 

405 

406 if prefetch > 2: 406 ↛ 407line 406 didn't jump to line 407 because the condition on line 406 was never true

407 violations.append( 

408 { 

409 "rule_id": self.id, 

410 "severity": self.severity, 

411 "affected_resource": "celery.worker_prefetch_multiplier", 

412 "message": f"Worker prefetch multiplier > 2 ({prefetch}), memory risk", 

413 "remediation": "Reduce to 1-2 and monitor memory usage", 

414 } 

415 ) 

416 

417 return violations 

418 

419 

420class LogStorageRule(BaseRule): 

421 """Check log storage location.""" 

422 

423 def __init__(self): 

424 super().__init__() 

425 self.id = "CFG013" 

426 self.name = "Local Log Storage" 

427 self.severity = "high" 

428 self.category = "reliability" 

429 self.execution_mode = "static" 

430 

431 def evaluate(self, config: Dict[str, Any]) -> List[Dict[str, Any]]: 

432 """Check log storage.""" 

433 violations = [] 

434 

435 log_folder = config.get("logging", {}).get("base_log_folder", "/var/log/airflow") 

436 

437 if "/var/log" in log_folder or "/tmp" in log_folder: 437 ↛ 448line 437 didn't jump to line 448 because the condition on line 437 was always true

438 violations.append( 

439 { 

440 "rule_id": self.id, 

441 "severity": self.severity, 

442 "affected_resource": "logging.base_log_folder", 

443 "message": "Logs stored on local disk (no HA, no persistence)", 

444 "remediation": "Use S3, GCS, or HDFS for log storage", 

445 } 

446 ) 

447 

448 return violations 

449 

450 

451class DatabaseBackupRule(BaseRule): 

452 """Check database backup configuration.""" 

453 

454 def __init__(self): 

455 super().__init__() 

456 self.id = "CFG014" 

457 self.name = "No Database Backup" 

458 self.severity = "high" 

459 self.category = "reliability" 

460 self.execution_mode = "static" 

461 

462 def evaluate(self, config: Dict[str, Any]) -> List[Dict[str, Any]]: 

463 """Check database backup.""" 

464 violations = [] 

465 

466 # Check if backup configuration present 

467 if "backup" not in str(config).lower(): 467 ↛ 478line 467 didn't jump to line 478 because the condition on line 467 was always true

468 violations.append( 

469 { 

470 "rule_id": self.id, 

471 "severity": self.severity, 

472 "affected_resource": "database", 

473 "message": "No database backup configuration detected", 

474 "remediation": "Configure daily backups of Airflow metadata database", 

475 } 

476 ) 

477 

478 return violations 

479 

480 

481class DAGFolderConfigurationRule(BaseRule): 

482 """Check DAG folder configuration.""" 

483 

484 def __init__(self): 

485 super().__init__() 

486 self.id = "CFG015" 

487 self.name = "DAG Folder On NFS" 

488 self.severity = "medium" 

489 self.category = "performance" 

490 self.execution_mode = "static" 

491 

492 def evaluate(self, config: Dict[str, Any]) -> List[Dict[str, Any]]: 

493 """Check DAG folder location.""" 

494 violations = [] 

495 

496 dag_folder = config.get("core", {}).get("dags_folder", "/airflow/dags") 

497 

498 if "/mnt" in dag_folder or "/nfs" in dag_folder: 498 ↛ 499line 498 didn't jump to line 499 because the condition on line 498 was never true

499 violations.append( 

500 { 

501 "rule_id": self.id, 

502 "severity": self.severity, 

503 "affected_resource": "core.dags_folder", 

504 "message": "DAG folder on NFS mount (parse time issues)", 

505 "remediation": "Use local SSD storage or container volume mounts", 

506 } 

507 ) 

508 

509 return violations