306 lines
11 KiB
Python
306 lines
11 KiB
Python
|
"""
|
||
|
|
Transaction Decorators and Context Managers
|
||
|
|
Provides convenient transaction handling for service methods
|
||
|
|
"""
|
||
|
|
|
||
|
|
from functools import wraps
|
||
|
|
from typing import Callable, Any, Optional
|
||
|
|
from contextlib import asynccontextmanager
|
||
|
|
import structlog
|
||
|
|
|
||
|
|
from .base import DatabaseManager
|
||
|
|
from .unit_of_work import UnitOfWork
|
||
|
|
from .exceptions import TransactionError
|
||
|
|
|
||
|
|
logger = structlog.get_logger()
|
||
|
|
|
||
|
|
|
||
|
|
def transactional(database_manager: DatabaseManager, auto_commit: bool = True):
|
||
|
|
"""
|
||
|
|
Decorator that wraps a method in a database transaction
|
||
|
|
|
||
|
|
Args:
|
||
|
|
database_manager: DatabaseManager instance
|
||
|
|
auto_commit: Whether to auto-commit on success
|
||
|
|
|
||
|
|
Usage:
|
||
|
|
@transactional(database_manager)
|
||
|
|
async def create_user_with_profile(self, user_data, profile_data):
|
||
|
|
# Your business logic here
|
||
|
|
# Transaction is automatically managed
|
||
|
|
pass
|
||
|
|
"""
|
||
|
|
def decorator(func: Callable) -> Callable:
|
||
|
|
@wraps(func)
|
||
|
|
async def wrapper(*args, **kwargs):
|
||
|
|
async with database_manager.get_background_session() as session:
|
||
|
|
try:
|
||
|
|
# Inject session into kwargs if not present
|
||
|
|
if 'session' not in kwargs:
|
||
|
|
kwargs['session'] = session
|
||
|
|
|
||
|
|
result = await func(*args, **kwargs)
|
||
|
|
|
||
|
|
# Session is auto-committed by get_background_session
|
||
|
|
logger.debug(f"Transaction completed successfully for {func.__name__}")
|
||
|
|
return result
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
# Session is auto-rolled back by get_background_session
|
||
|
|
logger.error(f"Transaction failed for {func.__name__}", error=str(e))
|
||
|
|
raise TransactionError(f"Transaction failed: {str(e)}")
|
||
|
|
|
||
|
|
return wrapper
|
||
|
|
return decorator
|
||
|
|
|
||
|
|
|
||
|
|
def unit_of_work_transactional(database_manager: DatabaseManager):
|
||
|
|
"""
|
||
|
|
Decorator that provides Unit of Work pattern for complex operations
|
||
|
|
|
||
|
|
Usage:
|
||
|
|
@unit_of_work_transactional(database_manager)
|
||
|
|
async def complex_business_operation(self, data, uow: UnitOfWork):
|
||
|
|
user_repo = uow.register_repository("users", UserRepository, User)
|
||
|
|
sales_repo = uow.register_repository("sales", SalesRepository, SalesData)
|
||
|
|
|
||
|
|
user = await user_repo.create(data.user)
|
||
|
|
sale = await sales_repo.create(data.sale)
|
||
|
|
|
||
|
|
# UnitOfWork automatically commits
|
||
|
|
return {"user": user, "sale": sale}
|
||
|
|
"""
|
||
|
|
def decorator(func: Callable) -> Callable:
|
||
|
|
@wraps(func)
|
||
|
|
async def wrapper(*args, **kwargs):
|
||
|
|
async with database_manager.get_background_session() as session:
|
||
|
|
async with UnitOfWork(session, auto_commit=True) as uow:
|
||
|
|
try:
|
||
|
|
# Inject UnitOfWork into kwargs
|
||
|
|
kwargs['uow'] = uow
|
||
|
|
|
||
|
|
result = await func(*args, **kwargs)
|
||
|
|
|
||
|
|
logger.debug(f"Unit of Work transaction completed for {func.__name__}")
|
||
|
|
return result
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
logger.error(f"Unit of Work transaction failed for {func.__name__}",
|
||
|
|
error=str(e))
|
||
|
|
raise TransactionError(f"Transaction failed: {str(e)}")
|
||
|
|
|
||
|
|
return wrapper
|
||
|
|
return decorator
|
||
|
|
|
||
|
|
|
||
|
|
@asynccontextmanager
|
||
|
|
async def managed_transaction(database_manager: DatabaseManager):
|
||
|
|
"""
|
||
|
|
Context manager for explicit transaction control
|
||
|
|
|
||
|
|
Usage:
|
||
|
|
async with managed_transaction(database_manager) as session:
|
||
|
|
# Your database operations here
|
||
|
|
user = User(name="John")
|
||
|
|
session.add(user)
|
||
|
|
# Auto-commits on exit, rolls back on exception
|
||
|
|
"""
|
||
|
|
async with database_manager.get_background_session() as session:
|
||
|
|
try:
|
||
|
|
logger.debug("Starting managed transaction")
|
||
|
|
yield session
|
||
|
|
logger.debug("Managed transaction completed successfully")
|
||
|
|
except Exception as e:
|
||
|
|
logger.error("Managed transaction failed", error=str(e))
|
||
|
|
raise
|
||
|
|
|
||
|
|
|
||
|
|
@asynccontextmanager
|
||
|
|
async def managed_unit_of_work(database_manager: DatabaseManager, event_publisher=None):
|
||
|
|
"""
|
||
|
|
Context manager for explicit Unit of Work control
|
||
|
|
|
||
|
|
Usage:
|
||
|
|
async with managed_unit_of_work(database_manager) as uow:
|
||
|
|
user_repo = uow.register_repository("users", UserRepository, User)
|
||
|
|
user = await user_repo.create(user_data)
|
||
|
|
await uow.commit()
|
||
|
|
"""
|
||
|
|
async with database_manager.get_background_session() as session:
|
||
|
|
uow = UnitOfWork(session)
|
||
|
|
try:
|
||
|
|
logger.debug("Starting managed Unit of Work")
|
||
|
|
yield uow
|
||
|
|
|
||
|
|
if not uow._committed:
|
||
|
|
await uow.commit()
|
||
|
|
|
||
|
|
logger.debug("Managed Unit of Work completed successfully")
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
if not uow._rolled_back:
|
||
|
|
await uow.rollback()
|
||
|
|
logger.error("Managed Unit of Work failed", error=str(e))
|
||
|
|
raise
|
||
|
|
|
||
|
|
|
||
|
|
class TransactionManager:
|
||
|
|
"""
|
||
|
|
Advanced transaction manager for complex scenarios
|
||
|
|
|
||
|
|
Usage:
|
||
|
|
tx_manager = TransactionManager(database_manager)
|
||
|
|
|
||
|
|
async with tx_manager.create_transaction() as tx:
|
||
|
|
await tx.execute_in_transaction(my_business_logic, data)
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, database_manager: DatabaseManager):
|
||
|
|
self.database_manager = database_manager
|
||
|
|
|
||
|
|
@asynccontextmanager
|
||
|
|
async def create_transaction(self, isolation_level: Optional[str] = None):
|
||
|
|
"""Create a transaction with optional isolation level"""
|
||
|
|
async with self.database_manager.get_background_session() as session:
|
||
|
|
transaction_context = TransactionContext(session, isolation_level)
|
||
|
|
try:
|
||
|
|
yield transaction_context
|
||
|
|
except Exception as e:
|
||
|
|
logger.error("Transaction manager failed", error=str(e))
|
||
|
|
raise
|
||
|
|
|
||
|
|
async def execute_with_retry(
|
||
|
|
self,
|
||
|
|
func: Callable,
|
||
|
|
max_retries: int = 3,
|
||
|
|
*args,
|
||
|
|
**kwargs
|
||
|
|
):
|
||
|
|
"""Execute function with transaction retry on failure"""
|
||
|
|
last_error = None
|
||
|
|
|
||
|
|
for attempt in range(max_retries):
|
||
|
|
try:
|
||
|
|
async with managed_transaction(self.database_manager) as session:
|
||
|
|
kwargs['session'] = session
|
||
|
|
result = await func(*args, **kwargs)
|
||
|
|
logger.debug(f"Transaction succeeded on attempt {attempt + 1}")
|
||
|
|
return result
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
last_error = e
|
||
|
|
logger.warning(f"Transaction attempt {attempt + 1} failed",
|
||
|
|
error=str(e), remaining_attempts=max_retries - attempt - 1)
|
||
|
|
|
||
|
|
if attempt == max_retries - 1:
|
||
|
|
break
|
||
|
|
|
||
|
|
logger.error(f"All transaction attempts failed after {max_retries} tries")
|
||
|
|
raise TransactionError(f"Transaction failed after {max_retries} retries: {str(last_error)}")
|
||
|
|
|
||
|
|
|
||
|
|
class TransactionContext:
|
||
|
|
"""Context for managing individual transactions"""
|
||
|
|
|
||
|
|
def __init__(self, session, isolation_level: Optional[str] = None):
|
||
|
|
self.session = session
|
||
|
|
self.isolation_level = isolation_level
|
||
|
|
|
||
|
|
async def execute_in_transaction(self, func: Callable, *args, **kwargs):
|
||
|
|
"""Execute function within the transaction context"""
|
||
|
|
try:
|
||
|
|
kwargs['session'] = self.session
|
||
|
|
result = await func(*args, **kwargs)
|
||
|
|
return result
|
||
|
|
except Exception as e:
|
||
|
|
logger.error("Function execution failed in transaction context", error=str(e))
|
||
|
|
raise
|
||
|
|
|
||
|
|
|
||
|
|
# ===== UTILITY FUNCTIONS =====
|
||
|
|
|
||
|
|
async def run_in_transaction(database_manager: DatabaseManager, func: Callable, *args, **kwargs):
|
||
|
|
"""
|
||
|
|
Utility function to run any async function in a transaction
|
||
|
|
|
||
|
|
Usage:
|
||
|
|
result = await run_in_transaction(
|
||
|
|
database_manager,
|
||
|
|
my_async_function,
|
||
|
|
arg1, arg2,
|
||
|
|
kwarg1="value"
|
||
|
|
)
|
||
|
|
"""
|
||
|
|
async with managed_transaction(database_manager) as session:
|
||
|
|
kwargs['session'] = session
|
||
|
|
return await func(*args, **kwargs)
|
||
|
|
|
||
|
|
|
||
|
|
async def run_with_unit_of_work(
|
||
|
|
database_manager: DatabaseManager,
|
||
|
|
func: Callable,
|
||
|
|
*args,
|
||
|
|
**kwargs
|
||
|
|
):
|
||
|
|
"""
|
||
|
|
Utility function to run any async function with Unit of Work
|
||
|
|
|
||
|
|
Usage:
|
||
|
|
result = await run_with_unit_of_work(
|
||
|
|
database_manager,
|
||
|
|
my_complex_function,
|
||
|
|
arg1, arg2
|
||
|
|
)
|
||
|
|
"""
|
||
|
|
async with managed_unit_of_work(database_manager) as uow:
|
||
|
|
kwargs['uow'] = uow
|
||
|
|
return await func(*args, **kwargs)
|
||
|
|
|
||
|
|
|
||
|
|
# ===== BATCH OPERATIONS =====
|
||
|
|
|
||
|
|
@asynccontextmanager
|
||
|
|
async def batch_operation(database_manager: DatabaseManager, batch_size: int = 1000):
|
||
|
|
"""
|
||
|
|
Context manager for batch operations with automatic commit batching
|
||
|
|
|
||
|
|
Usage:
|
||
|
|
async with batch_operation(database_manager, batch_size=500) as batch:
|
||
|
|
for item in large_dataset:
|
||
|
|
await batch.add_operation(create_record, item)
|
||
|
|
"""
|
||
|
|
async with database_manager.get_background_session() as session:
|
||
|
|
batch_context = BatchOperationContext(session, batch_size)
|
||
|
|
try:
|
||
|
|
yield batch_context
|
||
|
|
await batch_context.flush_remaining()
|
||
|
|
except Exception as e:
|
||
|
|
logger.error("Batch operation failed", error=str(e))
|
||
|
|
raise
|
||
|
|
|
||
|
|
|
||
|
|
class BatchOperationContext:
|
||
|
|
"""Context for managing batch database operations"""
|
||
|
|
|
||
|
|
def __init__(self, session, batch_size: int):
|
||
|
|
self.session = session
|
||
|
|
self.batch_size = batch_size
|
||
|
|
self.operation_count = 0
|
||
|
|
|
||
|
|
async def add_operation(self, func: Callable, *args, **kwargs):
|
||
|
|
"""Add operation to batch"""
|
||
|
|
kwargs['session'] = self.session
|
||
|
|
await func(*args, **kwargs)
|
||
|
|
|
||
|
|
self.operation_count += 1
|
||
|
|
|
||
|
|
if self.operation_count >= self.batch_size:
|
||
|
|
await self.session.commit()
|
||
|
|
self.operation_count = 0
|
||
|
|
logger.debug(f"Batch committed at {self.batch_size} operations")
|
||
|
|
|
||
|
|
async def flush_remaining(self):
|
||
|
|
"""Commit any remaining operations"""
|
||
|
|
if self.operation_count > 0:
|
||
|
|
await self.session.commit()
|
||
|
|
logger.debug(f"Final batch committed with {self.operation_count} operations")
|