Coverage for src/dataknobs_fsm/storage/base.py: 68%

1"""Base interfaces and classes for history storage."""

2

3 from abc import ABC, abstractmethod

4 from enum import Enum

5 from typing import Any, Dict, List, Type

6

7 from dataknobs_fsm.core.data_modes import DataHandlingMode

8 from dataknobs_fsm.execution.history import ExecutionHistory, ExecutionStep

9

10

11 class StorageBackend(Enum):

12 """Available storage backends."""

13 MEMORY = "memory"

14 FILE = "file"

15 SQLITE = "sqlite"

16 POSTGRES = "postgres"

17 MONGODB = "mongodb"

18 S3 = "s3"

19 ELASTICSEARCH = "elasticsearch"

20

21

22 class StorageConfig:

23 """Configuration for history storage."""

24

25 def __init__(

26 self,

27 backend: StorageBackend = StorageBackend.MEMORY,

28 connection_params: Dict[str, Any] | None = None,

29 retention_policy: Dict[str, Any] | None = None,

30 compression: bool = False,

31 batch_size: int = 100,

32 mode_specific_config: Dict[DataHandlingMode, Dict[str, Any]] | None = None

33 ):

34 """Initialize storage configuration.

35

36 Args:

37 backend: Storage backend to use.

38 connection_params: Backend-specific connection parameters.

39 retention_policy: Policy for data retention.

40 compression: Whether to compress stored data.

41 batch_size: Batch size for bulk operations.

42 mode_specific_config: Configuration per data mode.

43 """

44 self.backend = backend

45 self.connection_params = connection_params or {}

46 self.retention_policy = retention_policy or {}

47 self.compression = compression

48 self.batch_size = batch_size

49 self.mode_specific_config = mode_specific_config or {}

50

51 def get_mode_config(self, mode: DataHandlingMode) -> Dict[str, Any]:

52 """Get configuration for a specific data mode.

53

54 Args:

55 mode: Data mode.

56

57 Returns:

58 Configuration for that mode.

59 """

60 return self.mode_specific_config.get(mode, {})

61

62

63 class IHistoryStorage(ABC):

64 """Interface for history storage backends."""

65

66 @abstractmethod

67 async def save_history(

68 self,

69 history: ExecutionHistory,

70 metadata: Dict[str, Any] | None = None

71 ) -> str:

72 """Save execution history.

73

74 Args:

75 history: Execution history to save.

76 metadata: Optional metadata.

77

78 Returns:

79 Storage ID for the saved history.

80 """

81 pass

82

83 @abstractmethod

84 async def load_history(

85 self,

86 history_id: str

87 ) -> ExecutionHistory | None:

88 """Load execution history by ID.

89

90 Args:

91 history_id: ID of the history to load.

92

93 Returns:

94 ExecutionHistory if found, None otherwise.

95 """

96 pass

97

98 @abstractmethod

99 async def save_step(

100 self,

101 execution_id: str,

102 step: ExecutionStep,

103 parent_id: str | None = None

104 ) -> str:

105 """Save a single execution step.

106

107 Args:

108 execution_id: Execution ID this step belongs to.

109 step: Execution step to save.

110 parent_id: Parent step ID if branching.

111

112 Returns:

113 Storage ID for the saved step.

114 """

115 pass

116

117 @abstractmethod

118 async def load_steps(

119 self,

120 execution_id: str,

121 filters: Dict[str, Any] | None = None

122 ) -> List[ExecutionStep]:

123 """Load execution steps.

124

125 Args:

126 execution_id: Execution ID to load steps for.

127 filters: Optional filters (e.g., state_name, status).

128

129 Returns:

130 List of execution steps.

131 """

132 pass

133

134 @abstractmethod

135 async def query_histories(

136 self,

137 filters: Dict[str, Any],

138 limit: int = 100,

139 offset: int = 0

140 ) -> List[Dict[str, Any]]:

141 """Query execution histories.

142

143 Args:

144 filters: Query filters.

145 limit: Maximum results to return.

146 offset: Result offset for pagination.

147

148 Returns:

149 List of history summaries.

150 """

151 pass

152

153 @abstractmethod

154 async def delete_history(

155 self,

156 history_id: str

157 ) -> bool:

158 """Delete execution history.

159

160 Args:

161 history_id: ID of history to delete.

162

163 Returns:

164 True if deleted successfully.

165 """

166 pass

167

168 @abstractmethod

169 async def get_statistics(

170 self,

171 execution_id: str | None = None

172 ) -> Dict[str, Any]:

173 """Get storage statistics.

174

175 Args:

176 execution_id: Optional execution ID for specific stats.

177

178 Returns:

179 Storage statistics.

180 """

181 pass

182

183 @abstractmethod

184 async def cleanup(

185 self,

186 before_timestamp: float | None = None,

187 keep_failed: bool = True

188 ) -> int:

189 """Clean up old histories.

190

191 Args:

192 before_timestamp: Delete histories before this timestamp.

193 keep_failed: Whether to keep failed executions.

194

195 Returns:

196 Number of histories deleted.

197 """

198 pass

199

200

201 class BaseHistoryStorage(IHistoryStorage):

202 """Base class for history storage implementations."""

203

204 def __init__(self, config: StorageConfig):

205 """Initialize storage with configuration.

206

207 Args:

208 config: Storage configuration.

209 """

210 self.config = config

211 self._initialized = False

212

213 async def initialize(self) -> None:

214 """Initialize the storage backend."""

215 if not self._initialized:

216 await self._setup_backend()

217 self._initialized = True

218

219 @abstractmethod

220 async def _setup_backend(self) -> None:

221 """Set up the backend storage."""

222 pass

223

224 def _serialize_history(self, history: ExecutionHistory) -> Dict[str, Any]:

225 """Serialize execution history for storage.

226

227 Args:

228 history: Execution history to serialize.

229

230 Returns:

231 Serialized history.

232 """

233 data = history.to_dict()

234

235 # Apply compression if configured

236 if self.config.compression:

237 import zlib

238 import json

239 import base64

240

241 json_str = json.dumps(data)

242 compressed = zlib.compress(json_str.encode('utf-8'))

243 data = {

244 'compressed': True,

245 'data': base64.b64encode(compressed).decode('utf-8')

246 }

247

248 return data

249

250 def _deserialize_history(

251 self,

252 data: Dict[str, Any],

253 fsm_name: str,

254 execution_id: str

255 ) -> ExecutionHistory:

256 """Deserialize execution history from storage.

257

258 Args:

259 data: Serialized history data.

260 fsm_name: FSM name.

261 execution_id: Execution ID.

262

263 Returns:

264 ExecutionHistory instance.

265 """

266 # Decompress if needed

267 if data.get('compressed'):

268 import zlib

269 import json

270 import base64

271

272 compressed = base64.b64decode(data['data'])

273 json_str = zlib.decompress(compressed).decode('utf-8')

274 data = json.loads(json_str)

275

276 # Use ExecutionHistory.from_dict which properly reconstructs the tree

277 return ExecutionHistory.from_dict(data)

278

279 def _apply_retention_policy(self, histories: List[Dict[str, Any]]) -> List[Dict[str, Any]]:

280 """Apply retention policy to histories.

281

282 Args:

283 histories: List of history records.

284

285 Returns:

286 Filtered list based on retention policy.

287 """

288 if not self.config.retention_policy:

289 return histories

290

291 import time

292

293 max_age = self.config.retention_policy.get('max_age_days')

294 max_count = self.config.retention_policy.get('max_count')

295

296 if max_age:

297 cutoff = time.time() - (max_age * 86400)

298 histories = [h for h in histories if h.get('timestamp', 0) > cutoff]

299

300 if max_count and len(histories) > max_count:

301 # Keep most recent

302 histories = sorted(histories, key=lambda x: x.get('timestamp', 0), reverse=True)

303 histories = histories[:max_count]

304

305 return histories

306

307

308 class StorageFactory:

309 """Factory for creating history storage instances."""

310

311 _registry: Dict[StorageBackend, Type[IHistoryStorage]] = {}

312

313 @classmethod

314 def register(

315 cls,

316 backend: StorageBackend,

317 storage_class: Type[IHistoryStorage]

318 ) -> None:

319 """Register a storage backend.

320

321 Args:

322 backend: Backend type.

323 storage_class: Storage class to register.

324 """

325 cls._registry[backend] = storage_class

326

327 @classmethod

328 def create(

329 cls,

330 config: StorageConfig

331 ) -> IHistoryStorage:

332 """Create a storage instance.

333

334 Args:

335 config: Storage configuration.

336

337 Returns:

338 Storage instance.

339

340 Raises:

341 ValueError: If backend not registered.

342 """

343 storage_class = cls._registry.get(config.backend)

344 if not storage_class:

345 raise ValueError(f"Unknown storage backend: {config.backend}")

346

347 return storage_class(config) # type: ignore

348

349 @classmethod

350 def get_available_backends(cls) -> List[StorageBackend]:

351 """Get list of available backends.

352

353 Returns:

354 List of registered backends.

355 """

356 return list(cls._registry.keys())