DutyLog/seed_rooms.py

import argparse
import asyncio
import asyncpg


async def seed_floors_and_rooms(database_url: str):
    async with asyncpg.create_pool(database_url) as pool:
        async with pool.acquire() as conn:
            async with conn.transaction():

                # Создаем этажи (со 2-го по 5-й)
                floors = [2, 3, 4, 5]
                await conn.executemany(
                    """
                    INSERT INTO floors (number)
                    VALUES ($1)
                    ON CONFLICT (number) DO NOTHING;
                    """,
                    [(f,) for f in floors]
                )
                print(f"Добавлены или уже существуют этажи: {floors}")

                # Получаем ID созданных этажей
                floor_rows = await conn.fetch("SELECT id, number FROM floors WHERE number = ANY($1::int[]);", floors)
                floor_map = {row['number']: row['id'] for row in floor_rows}

                # Создаем комнаты (x01 .. x16 для каждого этажа)
                rooms_to_insert = []
                for floor_num in floors:
                    floor_id = floor_map[floor_num]
                    for room_suffix in range(1, 17):
                        room_number = floor_num * 100 + room_suffix  # 201, 202, ..., 516
                        rooms_to_insert.append((room_number, floor_id))

                # Выполняем вставку комнат
                await conn.executemany(
                    """
                    INSERT INTO rooms (number, on_floor)
                    VALUES ($1, $2)
                    ON CONFLICT (number) DO NOTHING;
                    """,
                    rooms_to_insert
                )
                print(f"Добавлены или уже существуют {len(rooms_to_insert)} комнат.")


async def main():
    parser = argparse.ArgumentParser(description="Инициализация этажей (2-5) и комнат (x01-x16)")
    parser.add_argument("--db-url", required=True, help="URL базы данных")
    args = parser.parse_args()

    await seed_floors_and_rooms(args.db_url)
    print("Инициализация завершена.")


if __name__ == "__main__":
    asyncio.run(main())