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

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

2

3 from contextlib import contextmanager

4 from typing import Any, Dict

5

6 from dataknobs_data.factory import DatabaseFactory

7 from dataknobs_data.database import SyncDatabase, AsyncDatabase

8 from dataknobs_data.records import Record

9 from dataknobs_data.query import Query

10

11 from dataknobs_fsm.functions.base import ResourceError

12 from dataknobs_fsm.resources.base import (

13 BaseResourceProvider,

14 ResourceHealth,

15 ResourceStatus,

16)

17

18

19 class 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}