Coverage for python/pyairflowtester/models.py: 97%

1""" 

2Data models for PyAirflowTester. 

3""" 

4 

5from dataclasses import dataclass, field 

6from datetime import datetime 

7from typing import Any, Dict, List, Optional 

8 

9# Severity weights for scoring 

10SEVERITY_WEIGHTS = { 

11 "critical": 1.0, 

12 "high": 0.75, 

13 "medium": 0.5, 

14 "low": 0.25, 

15 "info": 0.1, 

16} 

17 

18 

19@dataclass 

20class Rule: 

21 """Rule definition.""" 

22 

23 id: str 

24 name: str 

25 severity: str 

26 category: str 

27 execution_mode: str 

28 description: str 

29 remediation: str 

30 tags: List[str] = field(default_factory=list) 

31 

32 

33@dataclass 

34class RuleViolation: 

35 """Rule violation result.""" 

36 

37 rule_id: str 

38 severity: str 

39 affected_resource: str 

40 message: str 

41 remediation: str 

42 context: Dict[str, Any] = field(default_factory=dict) 

43 timestamp: datetime = field(default_factory=datetime.utcnow) 

44 

45 def to_dict(self) -> Dict[str, Any]: 

46 """Convert to dictionary.""" 

47 return { 

48 "rule_id": self.rule_id, 

49 "severity": self.severity, 

50 "affected_resource": self.affected_resource, 

51 "message": self.message, 

52 "remediation": self.remediation, 

53 "context": self.context, 

54 "timestamp": self.timestamp.isoformat(), 

55 } 

56 

57 

58@dataclass 

59class DagDefinition: 

60 """Airflow DAG definition.""" 

61 

62 dag_id: str 

63 file_path: str 

64 source_code: str 

65 parse_time_ms: float = 0.0 

66 task_count: int = 0 

67 task_ids: List[str] = field(default_factory=list) 

68 max_task_depth: int = 0 

69 has_cycles: bool = False 

70 dynamic_dag_detected: bool = False 

71 expensive_imports: List[str] = field(default_factory=list) 

72 external_dependencies: List[str] = field(default_factory=list) 

73 

74 

75@dataclass 

76class DbtModel: 

77 """dbt model definition.""" 

78 

79 name: str 

80 materialization: str = "table" 

81 description: str = "" 

82 columns: List[str] = field(default_factory=list) 

83 tests: List[str] = field(default_factory=list) 

84 downstream_models: List[str] = field(default_factory=list) 

85 upstream_models: List[str] = field(default_factory=list) 

86 source_file: str = "" 

87 

88 

89@dataclass 

90class DbtTest: 

91 """dbt test definition.""" 

92 

93 name: str 

94 test_type: str 

95 model_name: str 

96 source_name: Optional[str] = None 

97 column_name: Optional[str] = None 

98 properties: Dict[str, str] = field(default_factory=dict) 

99 

100 

101@dataclass 

102class RiskScorecard: 

103 """Risk scorecard for a resource.""" 

104 

105 resource_id: str 

106 resource_type: str 

107 timestamp: datetime = field(default_factory=datetime.utcnow) 

108 reliability_score: float = 75.0 

109 performance_score: float = 75.0 

110 maintainability_score: float = 75.0 

111 security_score: float = 75.0 

112 cost_efficiency_score: float = 75.0 

113 overall_health_score: float = 75.0 

114 risk_level: str = "medium" 

115 

116 def to_dict(self) -> Dict[str, Any]: 

117 """Convert to dictionary.""" 

118 return { 

119 "resource_id": self.resource_id, 

120 "resource_type": self.resource_type, 

121 "timestamp": self.timestamp.isoformat(), 

122 "reliability_score": self.reliability_score, 

123 "performance_score": self.performance_score, 

124 "maintainability_score": self.maintainability_score, 

125 "security_score": self.security_score, 

126 "cost_efficiency_score": self.cost_efficiency_score, 

127 "overall_health_score": self.overall_health_score, 

128 "risk_level": self.risk_level, 

129 } 

130 

131 

132@dataclass 

133class AnalysisContext: 

134 """Context for analysis.""" 

135 

136 dag_id: Optional[str] = None 

137 task_id: Optional[str] = None 

138 model_name: Optional[str] = None 

139 test_id: Optional[str] = None 

140 metadata: Dict[str, str] = field(default_factory=dict)