update

seed_rooms.py

@@ -0,0 +1,57 @@
+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())