app.seed

  1"""
  2Database seed module.
  3
  4Populates the database with rooms and time slots for February, March,
  5and April 2026. Each day gets a random subset of rooms, and each room
  6gets hourly time slots from 08:00 to 18:00.
  7
  8This module is idempotent — it skips seeding if rooms already exist.
  9"""
 10
 11import random
 12from datetime import date, time, timedelta
 13
 14from loguru import logger
 15from sqlmodel import select
 16from sqlmodel.ext.asyncio.session import AsyncSession
 17
 18from app.database import engine
 19from app.models.booking import TimeSlot
 20from app.models.room import Room
 21
 22ROOMS = [
 23    {"name": "A-101", "capacity": 10},
 24    {"name": "A-102", "capacity": 15},
 25    {"name": "A-201", "capacity": 20},
 26    {"name": "A-202", "capacity": 8},
 27    {"name": "A-203", "capacity": 30},
 28    {"name": "B-101", "capacity": 12},
 29    {"name": "B-102", "capacity": 25},
 30    {"name": "B-201", "capacity": 6},
 31    {"name": "B-301", "capacity": 40},
 32    {"name": "C-101", "capacity": 18},
 33]
 34
 35# Hourly slots from 08:00 to 18:00
 36SLOT_HOURS = [(time(h, 0), time(h + 1, 0)) for h in range(8, 18)]
 37
 38# Months to seed: (year, month)
 39SEED_MONTHS = [(2026, 2), (2026, 3), (2026, 4)]
 40
 41# Deterministic seed so the data is reproducible
 42RNG_SEED = 2026
 43
 44
 45# ── Helpers ────────────────────────────────────────────────────────────────
 46
 47
 48def _days_in_month(year: int, month: int) -> list[date]:
 49    """Return all dates in a given month."""
 50    first = date(year, month, 1)
 51    if month == 12:
 52        last = date(year + 1, 1, 1)
 53    else:
 54        last = date(year, month + 1, 1)
 55    return [first + timedelta(days=i) for i in range((last - first).days)]
 56
 57
 58# ── Main seed function ─────────────────────────────────────────────────────
 59
 60
 61async def seed_rooms_and_slots() -> None:
 62    """
 63    Populate the database with rooms and time slots.
 64
 65    Idempotent: if any rooms already exist, the function returns early.
 66    """
 67    async with AsyncSession(engine) as session:
 68        # Check if rooms already exist
 69        existing = (await session.exec(select(Room).limit(1))).first()
 70        if existing:
 71            logger.info("Rooms already seeded — skipping.")
 72            return
 73
 74        rng = random.Random(RNG_SEED)
 75
 76        # 1. Create all rooms
 77        room_objects: list[Room] = []
 78        for room_data in ROOMS:
 79            room = Room(**room_data)
 80            session.add(room)
 81            room_objects.append(room)
 82
 83        await session.flush()  # assigns IDs
 84
 85        # 2. For each day in each month, pick a random subset of rooms
 86        #    and create hourly time slots for them.
 87        all_slots: list[TimeSlot] = []
 88
 89        for year, month in SEED_MONTHS:
 90            for day in _days_in_month(year, month):
 91                # Pick 1–len(rooms) rooms available this day
 92                num_rooms = rng.randint(1, len(room_objects))
 93                day_rooms = rng.sample(room_objects, num_rooms)
 94
 95                for room in day_rooms:
 96                    for start, end in SLOT_HOURS:
 97                        slot = TimeSlot(
 98                            room_id=room.id,  # type: ignore[arg-type]
 99                            slot_date=day,
100                            start_time=start,
101                            end_time=end,
102                        )
103                        all_slots.append(slot)
104
105        session.add_all(all_slots)
106        await session.commit()
107
108        logger.info(
109            f"Seeded {len(room_objects)} rooms and {len(all_slots)} time slots "
110            f"across {len(SEED_MONTHS)} months."
111        )