Coverage for src/dataknobs_fsm/resources/database.py: 0%

1"""Database resource adapter for dataknobs_data backends.""" 

2 

3from contextlib import contextmanager 

4from typing import Any, Dict 

5 

6from dataknobs_data.factory import DatabaseFactory 

7from dataknobs_data.database import SyncDatabase, AsyncDatabase 

8from dataknobs_data.records import Record 

9from dataknobs_data.query import Query 

10 

11from dataknobs_fsm.functions.base import ResourceError 

12from dataknobs_fsm.resources.base import ( 

13 BaseResourceProvider, 

14 ResourceHealth, 

15 ResourceStatus, 

16) 

17 

18 

19class DatabaseResourceAdapter(BaseResourceProvider): 

20 """Adapter to use dataknobs_data databases as FSM resources. 

21  

22 This adapter wraps dataknobs_data database backends to provide 

23 resource management capabilities for FSM states. 

24 """ 

25 

26 def __init__( 

27 self, 

28 name: str, 

29 backend: str = "memory", 

30 **backend_config 

31 ): 

32 """Initialize database resource adapter. 

33  

34 Args: 

35 name: Resource name. 

36 backend: Database backend type (memory, file, postgres, sqlite, etc). 

37 **backend_config: Backend-specific configuration passed to DatabaseFactory. 

38 """ 

39 config = {"backend": backend, **backend_config} 

40 super().__init__(name, config) 

41 

42 self.backend = backend 

43 self.factory = DatabaseFactory() 

44 self._database: SyncDatabase | None = None 

45 self._initialize_database() 

46 

47 def _initialize_database(self) -> None: 

48 """Initialize the database backend.""" 

49 try: 

50 # Create database using factory 

51 self._database = self.factory.create(**self.config) 

52 self.status = ResourceStatus.IDLE 

53 except Exception as e: 

54 self.status = ResourceStatus.ERROR 

55 raise ResourceError( 

56 f"Failed to initialize database backend '{self.backend}': {e}", 

57 resource_name=self.name, 

58 operation="initialize" 

59 ) from e 

60 

61 def acquire(self, **kwargs) -> SyncDatabase: 

62 """Acquire database connection/instance. 

63  

64 The returned database object can be used for all database operations. 

65 For backends that support connection pooling (postgres, etc), this 

66 manages the underlying connections transparently. 

67  

68 Args: 

69 **kwargs: Additional parameters (unused, for interface compatibility). 

70  

71 Returns: 

72 Database instance for operations. 

73  

74 Raises: 

75 ResourceError: If acquisition fails. 

76 """ 

77 if self._database is None: 

78 raise ResourceError( 

79 "Database not initialized", 

80 resource_name=self.name, 

81 operation="acquire" 

82 ) 

83 

84 try: 

85 # For most backends, we return the same instance 

86 # Connection pooling is handled internally by the backend 

87 self.status = ResourceStatus.ACTIVE 

88 self._resources.append(self._database) 

89 return self._database 

90 except Exception as e: 

91 self.status = ResourceStatus.ERROR 

92 raise ResourceError( 

93 f"Failed to acquire database resource: {e}", 

94 resource_name=self.name, 

95 operation="acquire" 

96 ) from e 

97 

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

99 """Release database resource. 

100  

101 For pooled backends, this returns connections to the pool. 

102 For non-pooled backends, this is a no-op. 

103  

104 Args: 

105 resource: The database resource to release. 

106 """ 

107 if resource in self._resources: 

108 self._resources.remove(resource) 

109 

110 if not self._resources: 

111 self.status = ResourceStatus.IDLE 

112 

113 # Most backends handle connection cleanup internally 

114 # We don't need to do anything special here 

115 

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

117 """Validate database resource is still usable. 

118  

119 Args: 

120 resource: The database resource to validate. 

121  

122 Returns: 

123 True if the resource is valid and usable. 

124 """ 

125 if resource is None or not isinstance(resource, (SyncDatabase, AsyncDatabase)): 

126 return False 

127 

128 try: 

129 # Try a simple operation to validate the connection 

130 # This will vary by backend but should be lightweight 

131 if hasattr(resource, 'count'): 

132 # Try to count records (should return quickly even if 0) 

133 _ = resource.count() 

134 return True 

135 except Exception: 

136 return False 

137 

138 def health_check(self) -> ResourceHealth: 

139 """Check database health. 

140  

141 Returns: 

142 Health status of the database backend. 

143 """ 

144 if self._database is None: 

145 self.metrics.record_health_check(False) 

146 return ResourceHealth.UNKNOWN 

147 

148 try: 

149 # Perform a simple health check operation 

150 valid = self.validate(self._database) 

151 

152 if valid: 

153 self.metrics.record_health_check(True) 

154 return ResourceHealth.HEALTHY 

155 else: 

156 self.metrics.record_health_check(False) 

157 return ResourceHealth.UNHEALTHY 

158 except Exception: 

159 self.metrics.record_health_check(False) 

160 return ResourceHealth.UNHEALTHY 

161 

162 @contextmanager 

163 def transaction_context(self, database: SyncDatabase | None = None): 

164 """Context manager for database transactions. 

165  

166 Note: Transaction support depends on the backend. 

167 Some backends (memory, file) may not support true transactions. 

168  

169 Args: 

170 database: Optional database instance to use. 

171  

172 Yields: 

173 Database instance for operations within transaction. 

174 """ 

175 if database is None: 

176 database = self.acquire() 

177 should_release = True 

178 else: 

179 should_release = False 

180 

181 try: 

182 # For backends that support transactions, we could add 

183 # transaction begin/commit/rollback logic here 

184 # For now, we just ensure proper resource cleanup 

185 yield database 

186 finally: 

187 if should_release: 

188 self.release(database) 

189 

190 def close(self) -> None: 

191 """Close the database resource and clean up.""" 

192 import logging 

193 logger = logging.getLogger(__name__) 

194 

195 # Release all tracked resources first 

196 super().close() 

197 

198 # Close the database backend if it has a close method 

199 if self._database and hasattr(self._database, 'close'): 

200 try: 

201 # Attempt to flush any pending operations 

202 if hasattr(self._database, 'flush'): 

203 self._database.flush() 

204 

205 # Close the connection 

206 self._database.close() 

207 logger.debug(f"Successfully closed database connection for {self.name}") 

208 except AttributeError as e: 

209 logger.warning(f"Database {self.name} missing expected close method: {e}") 

210 except Exception as e: 

211 logger.error(f"Error closing database {self.name}: {e}") 

212 # Store error for debugging but don't re-raise 

213 if not hasattr(self, '_cleanup_errors'): 

214 self._cleanup_errors = [] 

215 self._cleanup_errors.append(f"Database close error: {e}") 

216 

217 self._database = None 

218 

219 # Convenience methods that delegate to the database 

220 

221 def create(self, record: Record, database: SyncDatabase | None = None) -> str: 

222 """Create a record in the database. 

223  

224 Args: 

225 record: Record to create. 

226 database: Optional database instance. 

227  

228 Returns: 

229 ID of the created record. 

230 """ 

231 if database is None: 

232 database = self._database 

233 

234 if database is None: 

235 raise ResourceError("No database available", resource_name=self.name, operation="create") 

236 

237 return database.create(record) 

238 

239 def read(self, record_id: str, database: SyncDatabase | None = None) -> Record | None: 

240 """Read a record from the database. 

241  

242 Args: 

243 record_id: ID of the record to read. 

244 database: Optional database instance. 

245  

246 Returns: 

247 The record if found, None otherwise. 

248 """ 

249 if database is None: 

250 database = self._database 

251 

252 if database is None: 

253 raise ResourceError("No database available", resource_name=self.name, operation="read") 

254 

255 return database.read(record_id) 

256 

257 def update(self, record_id: str, record: Record, database: SyncDatabase | None = None) -> bool: 

258 """Update a record in the database. 

259  

260 Args: 

261 record_id: ID of the record to update. 

262 record: Record with updates. 

263 database: Optional database instance. 

264  

265 Returns: 

266 True if update was successful. 

267 """ 

268 if database is None: 

269 database = self._database 

270 

271 if database is None: 

272 raise ResourceError("No database available", resource_name=self.name, operation="update") 

273 

274 return database.update(record_id, record) 

275 

276 def delete(self, record_id: str, database: SyncDatabase | None = None) -> bool: 

277 """Delete a record from the database. 

278  

279 Args: 

280 record_id: ID of the record to delete. 

281 database: Optional database instance. 

282  

283 Returns: 

284 True if deletion was successful. 

285 """ 

286 if database is None: 

287 database = self._database 

288 

289 if database is None: 

290 raise ResourceError("No database available", resource_name=self.name, operation="delete") 

291 

292 return database.delete(record_id) 

293 

294 def search(self, query: Query, database: SyncDatabase | None = None) -> list[Record]: 

295 """Search for records in the database. 

296  

297 Args: 

298 query: Search query. 

299 database: Optional database instance. 

300  

301 Returns: 

302 List of matching records. 

303 """ 

304 if database is None: 

305 database = self._database 

306 

307 if database is None: 

308 raise ResourceError("No database available", resource_name=self.name, operation="search") 

309 

310 return database.search(query) 

311 

312 def get_backend_info(self) -> Dict[str, Any]: 

313 """Get information about the database backend. 

314  

315 Returns: 

316 Backend information including capabilities. 

317 """ 

318 if self.factory: 

319 return self.factory.get_backend_info(self.backend) 

320 return {"backend": self.backend, "status": self.status.value}