Coverage for src/dataknobs_fsm/resources/base.py: 61%

1"""Base interfaces and classes for resource management.""" 

2 

3from abc import ABC, abstractmethod 

4from dataclasses import dataclass 

5from datetime import datetime 

6from enum import Enum 

7from typing import Any, Dict, List, Protocol, runtime_checkable 

8from contextlib import contextmanager 

9 

10 

11class ResourceStatus(Enum): 

12 """Status of a resource.""" 

13 

14 IDLE = "idle" 

15 ACTIVE = "active" 

16 BUSY = "busy" 

17 ERROR = "error" 

18 CLOSED = "closed" 

19 MAINTENANCE = "maintenance" 

20 

21 

22class ResourceHealth(Enum): 

23 """Health status of a resource.""" 

24 

25 HEALTHY = "healthy" 

26 DEGRADED = "degraded" 

27 UNHEALTHY = "unhealthy" 

28 UNKNOWN = "unknown" 

29 

30 

31@dataclass 

32class ResourceMetrics: 

33 """Metrics for resource usage.""" 

34 

35 total_acquisitions: int = 0 

36 active_connections: int = 0 

37 failed_acquisitions: int = 0 

38 average_hold_time: float = 0.0 

39 last_acquisition_time: datetime | None = None 

40 last_release_time: datetime | None = None 

41 health_check_failures: int = 0 

42 last_health_check: datetime | None = None 

43 average_acquisition_time: float = 0.0 

44 total_timeout_events: int = 0 

45 last_timeout_time: datetime | None = None 

46 

47 def record_acquisition(self, acquisition_time: float | None = None) -> None: 

48 """Record a resource acquisition. 

49  

50 Args: 

51 acquisition_time: Time taken to acquire the resource in seconds. 

52 """ 

53 self.total_acquisitions += 1 

54 self.active_connections += 1 

55 self.last_acquisition_time = datetime.now() 

56 

57 # Update average acquisition time if provided 

58 if acquisition_time is not None: 

59 if self.average_acquisition_time == 0: 

60 self.average_acquisition_time = acquisition_time 

61 else: 

62 # Rolling average 

63 self.average_acquisition_time = (self.average_acquisition_time * 0.9) + (acquisition_time * 0.1) 

64 

65 def record_release(self, hold_time: float) -> None: 

66 """Record a resource release. 

67  

68 Args: 

69 hold_time: How long the resource was held in seconds. 

70 """ 

71 self.active_connections = max(0, self.active_connections - 1) 

72 self.last_release_time = datetime.now() 

73 

74 # Update average hold time 

75 if self.average_hold_time == 0: 

76 self.average_hold_time = hold_time 

77 else: 

78 # Rolling average 

79 self.average_hold_time = (self.average_hold_time * 0.9) + (hold_time * 0.1) 

80 

81 def record_failure(self) -> None: 

82 """Record a failed acquisition.""" 

83 self.failed_acquisitions += 1 

84 

85 def record_health_check(self, success: bool) -> None: 

86 """Record a health check result. 

87  

88 Args: 

89 success: Whether the health check passed. 

90 """ 

91 self.last_health_check = datetime.now() 

92 if not success: 

93 self.health_check_failures += 1 

94 

95 def record_timeout(self) -> None: 

96 """Record a timeout event.""" 

97 self.total_timeout_events += 1 

98 self.last_timeout_time = datetime.now() 

99 self.failed_acquisitions += 1 

100 

101 

102@runtime_checkable 

103class IResourceProvider(Protocol): 

104 """Interface for resource providers.""" 

105 

106 def acquire(self, **kwargs) -> Any: 

107 """Acquire a resource. 

108  

109 Args: 

110 **kwargs: Provider-specific parameters. 

111  

112 Returns: 

113 The acquired resource. 

114  

115 Raises: 

116 ResourceError: If acquisition fails. 

117 """ 

118 ... 

119 

120 def release(self, resource: Any) -> None: 

121 """Release a resource. 

122  

123 Args: 

124 resource: The resource to release. 

125 """ 

126 ... 

127 

128 def validate(self, resource: Any) -> bool: 

129 """Validate that a resource is still valid. 

130  

131 Args: 

132 resource: The resource to validate. 

133  

134 Returns: 

135 True if the resource is valid. 

136 """ 

137 ... 

138 

139 def health_check(self) -> ResourceHealth: 

140 """Check the health of the resource provider. 

141  

142 Returns: 

143 The health status. 

144 """ 

145 ... 

146 

147 def get_metrics(self) -> ResourceMetrics: 

148 """Get resource metrics. 

149  

150 Returns: 

151 Current metrics. 

152 """ 

153 ... 

154 

155 

156@runtime_checkable 

157class IResourcePool(Protocol): 

158 """Interface for resource pools.""" 

159 

160 def acquire(self, timeout: float | None = None) -> Any: 

161 """Acquire a resource from the pool. 

162  

163 Args: 

164 timeout: Acquisition timeout in seconds. 

165  

166 Returns: 

167 The acquired resource. 

168  

169 Raises: 

170 ResourceError: If acquisition fails. 

171 """ 

172 ... 

173 

174 def release(self, resource: Any) -> None: 

175 """Return a resource to the pool. 

176  

177 Args: 

178 resource: The resource to return. 

179 """ 

180 ... 

181 

182 def size(self) -> int: 

183 """Get the current pool size. 

184  

185 Returns: 

186 Number of resources in the pool. 

187 """ 

188 ... 

189 

190 def available(self) -> int: 

191 """Get the number of available resources. 

192  

193 Returns: 

194 Number of available resources. 

195 """ 

196 ... 

197 

198 def close(self) -> None: 

199 """Close the pool and release all resources.""" 

200 ... 

201 

202 

203class BaseResourceProvider(ABC): 

204 """Base class for resource providers.""" 

205 

206 def __init__(self, name: str, config: Dict[str, Any] | None = None): 

207 """Initialize the provider. 

208  

209 Args: 

210 name: Provider name. 

211 config: Provider configuration. 

212 """ 

213 self.name = name 

214 self.config = config or {} 

215 self.status = ResourceStatus.IDLE 

216 self.metrics = ResourceMetrics() 

217 self._resources: List[Any] = [] 

218 

219 @abstractmethod 

220 def acquire(self, **kwargs) -> Any: 

221 """Acquire a resource. 

222  

223 Args: 

224 **kwargs: Provider-specific parameters. 

225  

226 Returns: 

227 The acquired resource. 

228 """ 

229 pass 

230 

231 @abstractmethod 

232 def release(self, resource: Any) -> None: 

233 """Release a resource. 

234  

235 Args: 

236 resource: The resource to release. 

237 """ 

238 pass 

239 

240 def validate(self, resource: Any) -> bool: 

241 """Validate a resource. 

242  

243 Args: 

244 resource: The resource to validate. 

245  

246 Returns: 

247 True if valid. 

248 """ 

249 return resource is not None 

250 

251 def health_check(self) -> ResourceHealth: 

252 """Check provider health. 

253  

254 Returns: 

255 Health status. 

256 """ 

257 if self.status == ResourceStatus.ERROR: 

258 return ResourceHealth.UNHEALTHY 

259 elif self.status == ResourceStatus.MAINTENANCE: 

260 return ResourceHealth.DEGRADED 

261 else: 

262 return ResourceHealth.HEALTHY 

263 

264 def get_metrics(self) -> ResourceMetrics: 

265 """Get provider metrics. 

266  

267 Returns: 

268 Current metrics. 

269 """ 

270 return self.metrics 

271 

272 @contextmanager 

273 def resource_context(self, **kwargs): 

274 """Context manager for resource acquisition. 

275  

276 Args: 

277 **kwargs: Acquisition parameters. 

278  

279 Yields: 

280 The acquired resource. 

281 """ 

282 resource = None 

283 start_time = datetime.now() 

284 try: 

285 resource = self.acquire(**kwargs) 

286 self.metrics.record_acquisition() 

287 yield resource 

288 except Exception: 

289 self.metrics.record_failure() 

290 raise 

291 finally: 

292 if resource is not None: 

293 hold_time = (datetime.now() - start_time).total_seconds() 

294 self.release(resource) 

295 self.metrics.record_release(hold_time) 

296 

297 def close(self) -> None: 

298 """Close the provider and release all resources.""" 

299 for resource in self._resources[:]: 

300 try: 

301 self.release(resource) 

302 except Exception: 

303 pass # Best effort cleanup 

304 self._resources.clear() 

305 self.status = ResourceStatus.CLOSED