Coverage for cli/cost_analytics.py: 95.76%

1"""Athena-backed cost analytics for the GCO CLI. 

2 

3Queries the cost allocation data the per-region cost-monitor services write 

4to the central cost report bucket (Hive-partitioned Parquet), through the 

5Glue table + Athena workgroup the monitoring stack provisions. Query results 

6land in the workgroup-enforced ``athena-results/`` prefix, KMS-encrypted and 

7lifecycle-expired. 

8 

9All SQL is assembled from validated identifiers and integer bounds only — 

10user-supplied strings (namespace filters) are passed through Athena 

11``ExecutionParameters``, never interpolated. 

12""" 

13 

14from __future__ import annotations 

15 

16import logging 

17import re 

18import time 

19from dataclasses import dataclass 

20from typing import Any 

21 

22import boto3 

23 

24from .config import GCOConfig, get_config 

25 

26logger = logging.getLogger(__name__) 

27 

28_DEFAULT_POLL_SECONDS = 1.0 

29_DEFAULT_TIMEOUT_SECONDS = 120.0 

30_MAX_RESULT_ROWS = 1_000 

31 

32_IDENTIFIER_PATTERN = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") 

33 

34_GRANULARITY_EXPRESSIONS = { 

35 "daily": "date", 

36 "hourly": "date_format(window_start, '%Y-%m-%dT%H:00Z')", 

37} 

38 

39_TOP_GROUPINGS = { 

40 "namespace": "namespace", 

41 "region": "region", 

42 "cluster": "cluster", 

43} 

44 

45 

46class AthenaQueryError(RuntimeError): 

47 """Raised when an Athena query fails, is cancelled, or times out.""" 

48 

49 

50def _database_name(project_name: str) -> str: 

51 """Glue database name — must mirror gco.stacks.constants.cost_glue_database_name.""" 

52 return f"{project_name.replace('-', '_')}_cost" 

53 

54 

55def _workgroup_name(project_name: str) -> str: 

56 """Athena workgroup — must mirror gco.stacks.constants.cost_athena_workgroup_name.""" 

57 return f"{project_name}-cost" 

58 

59 

60@dataclass 

61class QueryResult: 

62 """Column names plus row dictionaries for one executed query.""" 

63 

64 columns: list[str] 

65 rows: list[dict[str, Any]] 

66 query_execution_id: str 

67 

68 

69class CostAnalytics: 

70 """Runs canned cost aggregation queries through Athena.""" 

71 

72 def __init__( 

73 self, 

74 config: GCOConfig | None = None, 

75 *, 

76 athena_client: Any | None = None, 

77 region: str | None = None, 

78 poll_seconds: float = _DEFAULT_POLL_SECONDS, 

79 timeout_seconds: float = _DEFAULT_TIMEOUT_SECONDS, 

80 ) -> None: 

81 self._config = config or get_config() 

82 self.project_name = self._config.project_name 

83 # Athena + Glue live in the monitoring region. 

84 self.region = region or self._monitoring_region() 

85 self.database = _database_name(self.project_name) 

86 self.workgroup = _workgroup_name(self.project_name) 

87 self.table = "allocation_reports" 

88 self.poll_seconds = poll_seconds 

89 self.timeout_seconds = timeout_seconds 

90 self._athena = athena_client or boto3.client("athena", region_name=self.region) 

91 

92 def _monitoring_region(self) -> str: 

93 from .config import _load_cdk_json 

94 

95 regions = _load_cdk_json() or {} 

96 monitoring = regions.get("monitoring") 

97 if isinstance(monitoring, str) and monitoring: 

98 return monitoring 

99 return str(self._config.default_region) 

100 

101 # ------------------------------------------------------------------ 

102 # Query execution plumbing 

103 # ------------------------------------------------------------------ 

104 

105 def run_query(self, sql: str, parameters: list[str] | None = None) -> QueryResult: 

106 """Execute one query in the cost workgroup and return decoded rows.""" 

107 start_kwargs: dict[str, Any] = { 

108 "QueryString": sql, 

109 "QueryExecutionContext": {"Database": self.database}, 

110 "WorkGroup": self.workgroup, 

111 } 

112 if parameters: 

113 start_kwargs["ExecutionParameters"] = parameters 

114 try: 

115 start = self._athena.start_query_execution(**start_kwargs) 

116 except Exception as exc: # noqa: BLE001 - boto error shapes vary 

117 raise AthenaQueryError( 

118 f"Failed to start Athena query (workgroup {self.workgroup}, " 

119 f"database {self.database}): {exc}" 

120 ) from exc 

121 

122 execution_id = str(start["QueryExecutionId"]) 

123 deadline = time.monotonic() + self.timeout_seconds 

124 while True: 

125 execution = self._athena.get_query_execution(QueryExecutionId=execution_id) 

126 state = execution["QueryExecution"]["Status"]["State"] 

127 if state == "SUCCEEDED": 

128 break 

129 if state in {"FAILED", "CANCELLED"}: 

130 reason = execution["QueryExecution"]["Status"].get( 

131 "StateChangeReason", "no reason provided" 

132 ) 

133 raise AthenaQueryError(f"Athena query {state.lower()}: {reason}") 

134 if time.monotonic() >= deadline: 

135 self._athena.stop_query_execution(QueryExecutionId=execution_id) 

136 raise AthenaQueryError(f"Athena query timed out after {self.timeout_seconds:.0f}s") 

137 time.sleep(self.poll_seconds) 

138 

139 return self._collect_results(execution_id) 

140 

141 def _collect_results(self, execution_id: str) -> QueryResult: 

142 paginator = self._athena.get_paginator("get_query_results") 

143 columns: list[str] = [] 

144 rows: list[dict[str, Any]] = [] 

145 header_skipped = False 

146 for page in paginator.paginate(QueryExecutionId=execution_id): 

147 metadata = page.get("ResultSet", {}).get("ResultSetMetadata", {}) 

148 if not columns: 148 ↛ 150line 148 didn't jump to line 150 because the condition on line 148 was always true

149 columns = [str(column["Name"]) for column in metadata.get("ColumnInfo", [])] 

150 for record in page.get("ResultSet", {}).get("Rows", []): 

151 values = [entry.get("VarCharValue") for entry in record.get("Data", [])] 

152 if not header_skipped: 

153 # The first row of the first page repeats the header. 

154 header_skipped = True 

155 if values == columns: 155 ↛ 157line 155 didn't jump to line 157 because the condition on line 155 was always true

156 continue 

157 rows.append(dict(zip(columns, values, strict=False))) 

158 if len(rows) >= _MAX_RESULT_ROWS: 158 ↛ 159line 158 didn't jump to line 159 because the condition on line 158 was never true

159 return QueryResult(columns, rows, execution_id) 

160 return QueryResult(columns, rows, execution_id) 

161 

162 def _qualified_table(self) -> str: 

163 for identifier in (self.database, self.table): 

164 if not _IDENTIFIER_PATTERN.fullmatch(identifier): 164 ↛ 165line 164 didn't jump to line 165 because the condition on line 164 was never true

165 raise AthenaQueryError(f"Invalid Athena identifier: {identifier!r}") 

166 return f'"{self.database}"."{self.table}"' 

167 

168 @staticmethod 

169 def _days_clause(days: int) -> str: 

170 bounded = min(max(int(days), 1), 3_650) 

171 return f"date >= date_format(date_add('day', -{bounded}, current_date), '%Y-%m-%d')" 

172 

173 # ------------------------------------------------------------------ 

174 # Canned aggregations 

175 # ------------------------------------------------------------------ 

176 

177 def cost_by_namespace(self, days: int = 7, region: str | None = None) -> QueryResult: 

178 """Total cost per namespace across all regions (optionally one region).""" 

179 table = self._qualified_table() 

180 clauses = [self._days_clause(days)] 

181 parameters: list[str] = [] 

182 if region: 

183 clauses.append("region = ?") 

184 parameters.append(region) 

185 where = " AND ".join(clauses) 

186 sql = ( 

187 "SELECT namespace, " 

188 "round(sum(total_cost), 4) AS total_cost, " 

189 "round(sum(cpu_cost), 4) AS cpu_cost, " 

190 "round(sum(ram_cost), 4) AS ram_cost, " 

191 "round(sum(gpu_cost), 4) AS gpu_cost, " 

192 "round(sum(pv_cost), 4) AS pv_cost " 

193 f"FROM {table} WHERE {where} " # nosec B608 - identifiers regex-validated, values via ExecutionParameters 

194 "GROUP BY namespace ORDER BY total_cost DESC" 

195 ) 

196 return self.run_query(sql, parameters or None) 

197 

198 def cost_by_region(self, days: int = 7) -> QueryResult: 

199 """Total Kubernetes allocation cost per deployment region.""" 

200 table = self._qualified_table() 

201 sql = ( 

202 "SELECT region, " 

203 "round(sum(total_cost), 4) AS total_cost, " 

204 "round(sum(cpu_cost), 4) AS cpu_cost, " 

205 "round(sum(ram_cost), 4) AS ram_cost, " 

206 "round(sum(gpu_cost), 4) AS gpu_cost " 

207 f"FROM {table} WHERE {self._days_clause(days)} " # nosec B608 - identifiers regex-validated, days bounded int 

208 "GROUP BY region ORDER BY total_cost DESC" 

209 ) 

210 return self.run_query(sql) 

211 

212 def cost_over_time( 

213 self, 

214 days: int = 14, 

215 granularity: str = "daily", 

216 namespace: str | None = None, 

217 ) -> QueryResult: 

218 """Cost trend bucketed by day or hour, optionally for one namespace.""" 

219 bucket = _GRANULARITY_EXPRESSIONS.get(granularity) 

220 if bucket is None: 

221 raise AthenaQueryError(f"granularity must be one of {sorted(_GRANULARITY_EXPRESSIONS)}") 

222 table = self._qualified_table() 

223 clauses = [self._days_clause(days)] 

224 parameters: list[str] = [] 

225 if namespace: 

226 clauses.append("namespace = ?") 

227 parameters.append(namespace) 

228 where = " AND ".join(clauses) 

229 sql = ( 

230 f"SELECT {bucket} AS period, " 

231 "round(sum(total_cost), 4) AS total_cost " 

232 f"FROM {table} WHERE {where} " # nosec B608 - bucket from a fixed mapping, identifiers regex-validated, values via ExecutionParameters 

233 "GROUP BY 1 ORDER BY 1" 

234 ) 

235 return self.run_query(sql, parameters or None) 

236 

237 def top_spenders(self, n: int = 10, by: str = "namespace", days: int = 7) -> QueryResult: 

238 """Top-N spenders grouped by namespace, region, or cluster.""" 

239 grouping = _TOP_GROUPINGS.get(by) 

240 if grouping is None: 

241 raise AthenaQueryError(f"by must be one of {sorted(_TOP_GROUPINGS)}") 

242 bounded_n = min(max(int(n), 1), 100) 

243 table = self._qualified_table() 

244 sql = ( 

245 f"SELECT {grouping}, " 

246 "round(sum(total_cost), 4) AS total_cost " 

247 f"FROM {table} WHERE {self._days_clause(days)} " # nosec B608 - grouping from a fixed mapping, identifiers regex-validated, n/days bounded ints 

248 f"GROUP BY {grouping} ORDER BY total_cost DESC LIMIT {bounded_n}" 

249 ) 

250 return self.run_query(sql) 

251 

252 

253def get_cost_analytics(config: GCOConfig | None = None) -> CostAnalytics: 

254 """Factory function for CostAnalytics.""" 

255 return CostAnalytics(config=config)