Coverage for src/edwh/health.py: 21%

1import enum 

2import json 

3import sys 

4import typing as t 

5from dataclasses import dataclass 

6 

7from ewok import Context 

8from termcolor import colored, cprint, termcolor 

9 

10from .constants import DOCKER_COMPOSE, AnyDict 

11 

12StatusOptions = t.Literal[ 

13 "created", "restarting", "running", "removing", "paused", "exited", "exited ok", "dead", "unknown" 

14] 

15HealthOptions = t.Literal["starting", "unhealthy", "healthy"] | None 

16 

17 

18def find_container_ids(ctx: Context, container: str) -> list[str]: 

19 """Retrieve container IDs from Docker Compose. 

20 

21 Args: 

22 ctx (Context): The context in which to run the Docker command. 

23 container (str): The name of the container for which to retrieve IDs. 

24 

25 Returns: 

26 list[str]: A list of container IDs associated with the specified container. 

27 """ 

28 if result := ctx.run(f"{DOCKER_COMPOSE} ps -aq {container}", hide=True, warn=True): 

29 return result.stdout.strip().split("\n") 

30 else: 

31 return [] 

32 

33 

34def find_containers_ids(ctx: Context, *containers: str) -> dict[str, list[str]]: 

35 """Finds the IDs of the specified containers. 

36 

37 Args: 

38 ctx (Context): The context in which to find container IDs. 

39 *containers (str): Names of the containers to find IDs for. 

40 

41 Returns: 

42 dict[str, list[str]]: A dictionary where the keys are container names 

43 and the values are lists of corresponding container IDs. 

44 """ 

45 return {container: find_container_ids(ctx, container) for container in containers} 

46 

47 

48class HealthLevel(enum.IntEnum): 

49 # int-enum makes ordering possible 

50 

51 HEALTHY = enum.auto() # running and healthy 

52 RUNNING = enum.auto() # running but health unknown 

53 DEGRADED = enum.auto() # running and unhealthy 

54 STARTING = enum.auto() # starting, restarting 

55 UNKNOWN = enum.auto() # created 

56 DYING = enum.auto() # paused or removing 

57 STOPPED = enum.auto() # dead (in a good way) 

58 CRITICAL = enum.auto() # dead 

59 

60 @property 

61 def ok(self) -> bool: 

62 match self: 

63 case self.HEALTHY | self.RUNNING: 

64 return True 

65 case _: 

66 return False 

67 

68 @property 

69 def color(self) -> termcolor.Color: 

70 match self: 

71 case self.HEALTHY: 

72 return "green" 

73 case self.RUNNING: 

74 return "cyan" 

75 case self.DEGRADED: 

76 return "yellow" 

77 case self.STARTING: 

78 return "light_yellow" 

79 case self.STOPPED: 

80 return "blue" 

81 case self.DYING: 

82 return "light_red" 

83 case self.CRITICAL: 

84 return "red" 

85 case _: 

86 # unknown 

87 return "grey" 

88 

89 

90@dataclass 

91class HealthStatus: 

92 container_id: str 

93 container: str 

94 status: StatusOptions 

95 health: HealthOptions 

96 

97 @property 

98 def level(self) -> HealthLevel: 

99 """ 

100 Return health level (lower is better). 

101 """ 

102 if self.status == "exited ok": 

103 return HealthLevel.STOPPED 

104 elif self.status in {"exited", "dead"}: 

105 return HealthLevel.CRITICAL 

106 elif self.health == "healthy": 

107 return HealthLevel.HEALTHY 

108 elif self.health == "unhealthy": 

109 return HealthLevel.DEGRADED 

110 elif self.health == "starting": 

111 return HealthLevel.STARTING 

112 elif self.status == "running": 

113 # running but health unknown 

114 return HealthLevel.RUNNING 

115 elif self.status in {"restarting", "removing", "paused"}: 

116 return HealthLevel.DYING 

117 

118 else: 

119 return HealthLevel.UNKNOWN 

120 

121 @property 

122 def ok(self): 

123 return self.level.ok 

124 

125 @property 

126 def color(self) -> termcolor.Color: 

127 return self.level.color 

128 

129 def __repr__(self): 

130 return colored(f"Health({self})", self.color) 

131 

132 def __str__(self): 

133 status = self.status 

134 if self.health: 

135 status = f"{status} & {self.health}" 

136 return colored(f"{self.container}: {status}", self.color) 

137 

138 

139def docker_inspect(ctx: Context, container_id: str, *args: str) -> AnyDict | list[AnyDict]: 

140 """ 

141 Docker inspect a container by ID and get the first result. 

142 

143 Args: 

144 container_id: may be multiple (space separated) 

145 

146 :raise EnvironmentError if docker inspect failed. 

147 """ 

148 command = f"docker inspect {container_id}" 

149 if args: 

150 command = f"{command} {' '.join(args)}" 

151 

152 # note: this assumes bash is installed and available at /usr/bin/bash, 

153 # which should be fine in debian-based Linuces. 

154 # this allows you to pass "`docker compose ps -aq`" as container_id 

155 ran = ctx.run(command, hide=True, warn=True, shell="/usr/bin/bash") 

156 

157 if not ran.ok: 

158 cprint(f"docker inspect says: {ran.stderr}", file=sys.stderr, color="yellow") 

159 

160 try: 

161 # even if 'ran' is falsey, it could still have valid data. 

162 # e.g. `docker inspect <real> <real> <fake> 

163 return t.cast(AnyDict, json.loads(ran.stdout)) 

164 except json.decoder.JSONDecodeError: 

165 raise EnvironmentError(f"docker inspect {container_id} failed") 

166 

167 

168def get_healths(ctx: Context, *container_names: str) -> list[HealthStatus]: 

169 """ 

170 Retrieves the health statuses of specified containers. 

171 

172 Args: 

173 ctx (Context): The context in which the function is executed. 

174 *container_names (str): Variable length argument list of container names. 

175 

176 Returns: 

177 list[HealthStatus]: A list containing the health statuses of the specified containers. 

178 

179 Note: 

180 The amount of output health statuses can differ from the amount of containers if multiple replicas are used. 

181 """ 

182 

183 # {name: [ids]} 

184 container_name_to_ids = find_containers_ids(ctx, *container_names) 

185 

186 # note: use `docker inspect `docker compose ps -aq`` to prevent issues 

187 # when containers die between these two statements: 

188 # info_by_id = {_["Id"]: _["State"] for _ in inspect(ctx, " ".join(_ for _ in container_ids.values() if _))} 

189 try: 

190 docker_info = docker_inspect(ctx, "`docker compose ps -aq`") 

191 if isinstance(docker_info, list): 

192 info_by_id: dict[str, AnyDict] = {_["Id"]: _ for _ in docker_info} 

193 else: 

194 # invalid data returned 

195 info_by_id = {} 

196 

197 except EnvironmentError: 

198 # probably everything down (warning is already shown by inspect() -> this can be safely ignored) 

199 info_by_id = {} 

200 

201 def container_health(container_id: str, container_name: str, multiple: bool = False): 

202 if not (container_id and container_id in info_by_id): 

203 return HealthStatus( 

204 container_id, 

205 container_name, 

206 "dead", 

207 None, 

208 ) 

209 

210 info = info_by_id[container_id] 

211 state = info["State"] 

212 

213 config = info.get("Config", {}) 

214 labels = config.get("Labels", {}) 

215 container_number = labels.get("com.docker.compose.container-number", "1") 

216 

217 health = state.get("Health", {}) 

218 

219 container_status = state.get("Status") 

220 health_status = health.get("Status") 

221 

222 if container_status == "exited" and str(state.get("ExitCode")) == "0": 

223 # use exit code to know whether it was critical or not 

224 container_status = "exited ok" 

225 

226 return HealthStatus( 

227 container_id, 

228 f"{container_name}-{container_number}" if multiple else container_name, 

229 container_status, 

230 health_status, 

231 ) 

232 

233 result = [] 

234 for container_name in container_names: 

235 if container_ids := container_name_to_ids.get(container_name, []): 

236 for container_id in container_ids: 

237 result.append(container_health(container_id, container_name, multiple=len(container_ids) > 1)) 

238 else: 

239 # weird scenario 

240 result.append( 

241 HealthStatus( 

242 "", 

243 container_name, 

244 "unknown", 

245 None, 

246 ) 

247 ) 

248 

249 return result