REFACTOR - Database logic
This commit is contained in:
Urtzi Alfaro
306
shared/database/transactions.py
Normal file
306
shared/database/transactions.py
Normal file
@@ -0,0 +1,306 @@
|
||||
"""
|
||||
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")
|
||||
Reference in New Issue
Block a user