Coverage for src/dataknobs_fsm/core/transactions.py: 25%

1"""Transaction management for FSM data processing. 

2 

3This module provides different transaction strategies: 

4- SINGLE: Each state transition is a transaction 

5- BATCH: Multiple transitions batched into one transaction 

6- MANUAL: Explicit transaction control 

7""" 

8 

9from abc import ABC, abstractmethod 

10from contextlib import contextmanager 

11from enum import Enum 

12from typing import Any, Callable, Dict, List, TypeVar 

13import logging 

14 

15logger = logging.getLogger(__name__) 

16 

17T = TypeVar("T") 

18 

19 

20class TransactionStrategy(Enum): 

21 """Transaction management strategies.""" 

22 

23 SINGLE = "single" # One transaction per state transition 

24 BATCH = "batch" # Batch multiple transitions 

25 MANUAL = "manual" # Explicit transaction control 

26 

27 

28class TransactionState(Enum): 

29 """State of a transaction.""" 

30 

31 PENDING = "pending" 

32 ACTIVE = "active" 

33 COMMITTED = "committed" 

34 ROLLED_BACK = "rolled_back" 

35 FAILED = "failed" 

36 

37 

38class Transaction: 

39 """Represents a single transaction.""" 

40 

41 def __init__(self, transaction_id: str, strategy: TransactionStrategy): 

42 """Initialize a transaction. 

43  

44 Args: 

45 transaction_id: Unique identifier for the transaction. 

46 strategy: The transaction strategy being used. 

47 """ 

48 self.id = transaction_id 

49 self.strategy = strategy 

50 self.state = TransactionState.PENDING 

51 self.operations: List[Dict[str, Any]] = [] 

52 self.metadata: Dict[str, Any] = {} 

53 self._rollback_handlers: List[Callable[[], None]] = [] 

54 

55 def add_operation(self, operation: Dict[str, Any]) -> None: 

56 """Add an operation to the transaction. 

57  

58 Args: 

59 operation: Operation details to add. 

60 """ 

61 if self.state != TransactionState.ACTIVE: 

62 raise RuntimeError(f"Cannot add operation to transaction in state {self.state}") 

63 self.operations.append(operation) 

64 

65 def add_rollback_handler(self, handler: Callable[[], None]) -> None: 

66 """Add a rollback handler. 

67  

68 Args: 

69 handler: Function to call on rollback. 

70 """ 

71 self._rollback_handlers.append(handler) 

72 

73 def rollback(self) -> None: 

74 """Execute rollback handlers in reverse order.""" 

75 for handler in reversed(self._rollback_handlers): 

76 try: 

77 handler() 

78 except Exception as e: 

79 logger.error(f"Error in rollback handler: {e}") 

80 self.state = TransactionState.ROLLED_BACK 

81 

82 def commit(self) -> None: 

83 """Mark transaction as committed.""" 

84 self.state = TransactionState.COMMITTED 

85 

86 def __repr__(self) -> str: 

87 """String representation of the transaction.""" 

88 return f"Transaction(id={self.id}, strategy={self.strategy}, state={self.state})" 

89 

90 

91class TransactionManager(ABC): 

92 """Abstract base class for transaction managers.""" 

93 

94 def __init__(self, strategy: TransactionStrategy): 

95 """Initialize the transaction manager. 

96  

97 Args: 

98 strategy: The transaction strategy to use. 

99 """ 

100 self.strategy = strategy 

101 self._transactions: Dict[str, Transaction] = {} 

102 self._active_transaction: Transaction | None = None 

103 

104 @classmethod 

105 def create(cls, strategy: TransactionStrategy, **config) -> "TransactionManager": # noqa: ARG003 

106 """Factory method to create appropriate transaction manager. 

107  

108 Args: 

109 strategy: Transaction strategy to use 

110 **config: Strategy-specific configuration (currently unused) 

111  

112 Returns: 

113 Appropriate TransactionManager subclass instance 

114  

115 Raises: 

116 ValueError: If strategy is unknown 

117 """ 

118 # Import here to avoid circular dependencies 

119 if strategy == TransactionStrategy.SINGLE: 

120 return SingleTransactionManager() 

121 elif strategy == TransactionStrategy.BATCH: 

122 return BatchTransactionManager() 

123 elif strategy == TransactionStrategy.MANUAL: 

124 return ManualTransactionManager() 

125 else: 

126 raise ValueError(f"Unknown transaction strategy: {strategy}") 

127 

128 @abstractmethod 

129 def begin_transaction(self, transaction_id: str | None = None) -> Transaction: 

130 """Begin a new transaction. 

131  

132 Args: 

133 transaction_id: Optional ID for the transaction. 

134  

135 Returns: 

136 The created transaction. 

137 """ 

138 pass 

139 

140 @abstractmethod 

141 def commit_transaction(self, transaction_id: str | None = None) -> None: 

142 """Commit a transaction. 

143  

144 Args: 

145 transaction_id: ID of transaction to commit, or None for active. 

146 """ 

147 pass 

148 

149 @abstractmethod 

150 def rollback_transaction(self, transaction_id: str | None = None) -> None: 

151 """Rollback a transaction. 

152  

153 Args: 

154 transaction_id: ID of transaction to rollback, or None for active. 

155 """ 

156 pass 

157 

158 @abstractmethod 

159 def should_commit(self) -> bool: 

160 """Determine if a commit should happen now. 

161  

162 Returns: 

163 True if transaction should be committed. 

164 """ 

165 pass 

166 

167 @contextmanager 

168 def transaction(self, transaction_id: str | None = None): 

169 """Context manager for transactions. 

170  

171 Args: 

172 transaction_id: Optional ID for the transaction. 

173  

174 Yields: 

175 The active transaction. 

176 """ 

177 txn = self.begin_transaction(transaction_id) 

178 try: 

179 yield txn 

180 if self.should_commit(): 

181 self.commit_transaction(txn.id) 

182 except Exception as e: 

183 self.rollback_transaction(txn.id) 

184 raise e 

185 

186 def get_transaction(self, transaction_id: str) -> Transaction | None: 

187 """Get a transaction by ID. 

188  

189 Args: 

190 transaction_id: ID of the transaction. 

191  

192 Returns: 

193 The transaction if found, None otherwise. 

194 """ 

195 return self._transactions.get(transaction_id) 

196 

197 def get_active_transaction(self) -> Transaction | None: 

198 """Get the currently active transaction. 

199  

200 Returns: 

201 The active transaction if any. 

202 """ 

203 return self._active_transaction 

204 

205 

206class SingleTransactionManager(TransactionManager): 

207 """Manager for SINGLE transaction strategy - one per transition.""" 

208 

209 def __init__(self): 

210 """Initialize the single transaction manager.""" 

211 super().__init__(TransactionStrategy.SINGLE) 

212 self._transaction_counter = 0 

213 

214 def begin_transaction(self, transaction_id: str | None = None) -> Transaction: 

215 """Begin a new single transaction. 

216  

217 Args: 

218 transaction_id: Optional ID for the transaction. 

219  

220 Returns: 

221 The created transaction. 

222 """ 

223 if self._active_transaction is not None: 

224 self.commit_transaction(self._active_transaction.id) 

225 

226 if transaction_id is None: 

227 self._transaction_counter += 1 

228 transaction_id = f"single_txn_{self._transaction_counter}" 

229 

230 txn = Transaction(transaction_id, self.strategy) 

231 txn.state = TransactionState.ACTIVE 

232 self._transactions[transaction_id] = txn 

233 self._active_transaction = txn 

234 return txn 

235 

236 def commit_transaction(self, transaction_id: str | None = None) -> None: 

237 """Commit a single transaction. 

238  

239 Args: 

240 transaction_id: ID of transaction to commit. 

241 """ 

242 if transaction_id is None: 

243 if self._active_transaction is None: 

244 return 

245 transaction_id = self._active_transaction.id 

246 

247 txn = self._transactions.get(transaction_id) 

248 if txn: 

249 txn.commit() 

250 if self._active_transaction == txn: 

251 self._active_transaction = None 

252 

253 def rollback_transaction(self, transaction_id: str | None = None) -> None: 

254 """Rollback a single transaction. 

255  

256 Args: 

257 transaction_id: ID of transaction to rollback. 

258 """ 

259 if transaction_id is None: 

260 if self._active_transaction is None: 

261 return 

262 transaction_id = self._active_transaction.id 

263 

264 txn = self._transactions.get(transaction_id) 

265 if txn: 

266 txn.rollback() 

267 if self._active_transaction == txn: 

268 self._active_transaction = None 

269 

270 def should_commit(self) -> bool: 

271 """Single transactions always commit immediately. 

272  

273 Returns: 

274 True, always commit after each operation. 

275 """ 

276 return True 

277 

278 

279class BatchTransactionManager(TransactionManager): 

280 """Manager for BATCH transaction strategy.""" 

281 

282 def __init__(self, batch_size: int = 100, auto_commit: bool = True): 

283 """Initialize the batch transaction manager. 

284  

285 Args: 

286 batch_size: Number of operations before auto-commit. 

287 auto_commit: Whether to auto-commit when batch is full. 

288 """ 

289 super().__init__(TransactionStrategy.BATCH) 

290 self.batch_size = batch_size 

291 self.auto_commit = auto_commit 

292 self._batch_counter = 0 

293 self._operation_count = 0 

294 

295 def begin_transaction(self, transaction_id: str | None = None) -> Transaction: 

296 """Begin or get the current batch transaction. 

297  

298 Args: 

299 transaction_id: Optional ID for the transaction. 

300  

301 Returns: 

302 The active batch transaction. 

303 """ 

304 if self._active_transaction is None: 

305 if transaction_id is None: 

306 self._batch_counter += 1 

307 transaction_id = f"batch_txn_{self._batch_counter}" 

308 

309 txn = Transaction(transaction_id, self.strategy) 

310 txn.state = TransactionState.ACTIVE 

311 self._transactions[transaction_id] = txn 

312 self._active_transaction = txn 

313 self._operation_count = 0 

314 

315 return self._active_transaction 

316 

317 def add_to_batch(self, operation: Dict[str, Any]) -> None: 

318 """Add an operation to the current batch. 

319  

320 Args: 

321 operation: Operation to add to the batch. 

322 """ 

323 if self._active_transaction is None: 

324 self.begin_transaction() 

325 

326 self._active_transaction.add_operation(operation) 

327 self._operation_count += 1 

328 

329 if self.auto_commit and self._operation_count >= self.batch_size: 

330 self.commit_transaction() 

331 

332 def commit_transaction(self, transaction_id: str | None = None) -> None: 

333 """Commit the batch transaction. 

334  

335 Args: 

336 transaction_id: ID of transaction to commit. 

337 """ 

338 if transaction_id is None: 

339 if self._active_transaction is None: 

340 return 

341 transaction_id = self._active_transaction.id 

342 

343 txn = self._transactions.get(transaction_id) 

344 if txn: 

345 txn.commit() 

346 if self._active_transaction == txn: 

347 self._active_transaction = None 

348 self._operation_count = 0 

349 

350 def rollback_transaction(self, transaction_id: str | None = None) -> None: 

351 """Rollback the batch transaction. 

352  

353 Args: 

354 transaction_id: ID of transaction to rollback. 

355 """ 

356 if transaction_id is None: 

357 if self._active_transaction is None: 

358 return 

359 transaction_id = self._active_transaction.id 

360 

361 txn = self._transactions.get(transaction_id) 

362 if txn: 

363 txn.rollback() 

364 if self._active_transaction == txn: 

365 self._active_transaction = None 

366 self._operation_count = 0 

367 

368 def should_commit(self) -> bool: 

369 """Check if batch should be committed. 

370  

371 Returns: 

372 True if batch is full or explicitly requested. 

373 """ 

374 return self.auto_commit and self._operation_count >= self.batch_size 

375 

376 def flush(self) -> None: 

377 """Force commit of the current batch.""" 

378 if self._active_transaction is not None: 

379 self.commit_transaction() 

380 

381 

382class ManualTransactionManager(TransactionManager): 

383 """Manager for MANUAL transaction strategy - explicit control.""" 

384 

385 def __init__(self): 

386 """Initialize the manual transaction manager.""" 

387 super().__init__(TransactionStrategy.MANUAL) 

388 self._transaction_counter = 0 

389 

390 def begin_transaction(self, transaction_id: str | None = None) -> Transaction: 

391 """Begin a new manual transaction. 

392  

393 Args: 

394 transaction_id: Optional ID for the transaction. 

395  

396 Returns: 

397 The created transaction. 

398 """ 

399 if transaction_id is None: 

400 self._transaction_counter += 1 

401 transaction_id = f"manual_txn_{self._transaction_counter}" 

402 

403 if transaction_id in self._transactions: 

404 raise ValueError(f"Transaction {transaction_id} already exists") 

405 

406 txn = Transaction(transaction_id, self.strategy) 

407 txn.state = TransactionState.ACTIVE 

408 self._transactions[transaction_id] = txn 

409 self._active_transaction = txn 

410 return txn 

411 

412 def commit_transaction(self, transaction_id: str | None = None) -> None: 

413 """Manually commit a transaction. 

414  

415 Args: 

416 transaction_id: ID of transaction to commit. 

417 """ 

418 if transaction_id is None: 

419 if self._active_transaction is None: 

420 raise ValueError("No active transaction to commit") 

421 transaction_id = self._active_transaction.id 

422 

423 txn = self._transactions.get(transaction_id) 

424 if not txn: 

425 raise ValueError(f"Transaction {transaction_id} not found") 

426 

427 if txn.state != TransactionState.ACTIVE: 

428 raise ValueError(f"Transaction {transaction_id} is not active") 

429 

430 txn.commit() 

431 if self._active_transaction == txn: 

432 self._active_transaction = None 

433 

434 def rollback_transaction(self, transaction_id: str | None = None) -> None: 

435 """Manually rollback a transaction. 

436  

437 Args: 

438 transaction_id: ID of transaction to rollback. 

439 """ 

440 if transaction_id is None: 

441 if self._active_transaction is None: 

442 raise ValueError("No active transaction to rollback") 

443 transaction_id = self._active_transaction.id 

444 

445 txn = self._transactions.get(transaction_id) 

446 if not txn: 

447 raise ValueError(f"Transaction {transaction_id} not found") 

448 

449 if txn.state != TransactionState.ACTIVE: 

450 raise ValueError(f"Transaction {transaction_id} is not active") 

451 

452 txn.rollback() 

453 if self._active_transaction == txn: 

454 self._active_transaction = None 

455 

456 def should_commit(self) -> bool: 

457 """Manual transactions never auto-commit. 

458  

459 Returns: 

460 False, manual control required. 

461 """ 

462 return False 

463 

464 

465def create_transaction_manager( 

466 strategy: TransactionStrategy, 

467 **kwargs 

468) -> TransactionManager: 

469 """Factory function to create transaction managers. 

470  

471 Args: 

472 strategy: The transaction strategy to use. 

473 **kwargs: Additional arguments for specific managers. 

474  

475 Returns: 

476 The appropriate transaction manager. 

477 """ 

478 if strategy == TransactionStrategy.SINGLE: 

479 return SingleTransactionManager() 

480 elif strategy == TransactionStrategy.BATCH: 

481 batch_size = kwargs.get("batch_size", 100) 

482 auto_commit = kwargs.get("auto_commit", True) 

483 return BatchTransactionManager(batch_size, auto_commit) 

484 elif strategy == TransactionStrategy.MANUAL: 

485 return ManualTransactionManager() 

486 else: 

487 raise ValueError(f"Unknown transaction strategy: {strategy}")