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

1"""Transaction management for FSM data processing.

2

3 This 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

9 from abc import ABC, abstractmethod

10 from contextlib import contextmanager

11 from enum import Enum

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

13 import logging

14

15 logger = logging.getLogger(__name__)

16

17 T = TypeVar("T")

18

19

20 class 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

28 class 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

38 class 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

91 class 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

206 class 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

279 class 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

382 class 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

465 def 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}")