from typing import List, Optional, Dict, Any from uuid import UUID from sqlalchemy import select, update, delete from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.exc import IntegrityError import structlog from app.models.distribution import DeliverySchedule logger = structlog.get_logger() class DeliveryScheduleRepository: def __init__(self, session: AsyncSession): self.session = session async def create_schedule(self, schedule_data: Dict[str, Any]) -> DeliverySchedule: """Create a new delivery schedule""" try: schedule = DeliverySchedule(**schedule_data) self.session.add(schedule) await self.session.commit() await self.session.refresh(schedule) return schedule except IntegrityError as e: await self.session.rollback() logger.error("Error creating delivery schedule", error=str(e)) raise ValueError(f"Failed to create delivery schedule: {e}") except Exception as e: await self.session.rollback() logger.error("Unexpected error creating delivery schedule", error=str(e)) raise async def get_schedule_by_id(self, schedule_id: UUID) -> Optional[DeliverySchedule]: """Get a delivery schedule by ID""" result = await self.session.execute( select(DeliverySchedule).where(DeliverySchedule.id == schedule_id) ) return result.scalar_one_or_none() async def get_schedules_by_tenant(self, tenant_id: UUID) -> List[DeliverySchedule]: """Get all delivery schedules for a tenant""" result = await self.session.execute( select(DeliverySchedule).where(DeliverySchedule.tenant_id == tenant_id) ) return result.scalars().all() async def update_schedule(self, schedule_id: UUID, update_data: Dict[str, Any]) -> Optional[DeliverySchedule]: """Update a delivery schedule""" try: stmt = ( update(DeliverySchedule) .where(DeliverySchedule.id == schedule_id) .values(**update_data) .returning(DeliverySchedule) ) result = await self.session.execute(stmt) await self.session.commit() return result.scalar_one_or_none() except Exception as e: await self.session.rollback() logger.error("Error updating delivery schedule", error=str(e), schedule_id=schedule_id) raise async def delete_schedule(self, schedule_id: UUID) -> bool: """Delete a delivery schedule""" try: stmt = delete(DeliverySchedule).where(DeliverySchedule.id == schedule_id) result = await self.session.execute(stmt) await self.session.commit() return result.rowcount > 0 except Exception as e: await self.session.rollback() logger.error("Error deleting delivery schedule", error=str(e), schedule_id=schedule_id) raise