Добавлена ветка prod_prep

2026-07-31 12:58:23 +04:00
parent bf26370da9
commit f289631645
36 changed files with 893 additions and 290 deletions
@@ -1,13 +1,17 @@
+"""Fallback JSON-based storage (sync).
+Used only when DB_HOST is not set in .env.
+"""
 import json
-from app.logger import logger
 from pathlib import Path
+from app.logger import logger

 DB_DIR = Path(__file__).resolve().parent.parent.parent / "data"
 DB_DIR.mkdir(exist_ok=True)
-
 USERS_FILE = DB_DIR / "users.json"
 BOOKINGS_FILE = DB_DIR / "bookings.json"

+from app.config import USE_POSTGRES
+

 def _load_json(path, default):
    if path.exists():
@@ -24,6 +28,7 @@ def _save_json(path, data):
        json.dump(data, fh, ensure_ascii=False, indent=2)


+# --- Users ---
 def load_users():
    return _load_json(USERS_FILE, {})

@@ -43,19 +48,19 @@ def set_user(user_id, data):
    save_users(users)


+# --- Bookings ---
 def load_bookings():
    return _load_json(BOOKINGS_FILE, [])


 def save_bookings(bookings):
-    from app.logger import logger
    logger.debug("save_bookings: saving {n} bookings", n=len(bookings))
    _save_json(BOOKINGS_FILE, bookings)


 def create_booking(user_id, service, master, date, time, user_name=""):
-    from app.logger import logger
-    logger.debug("create_booking called: uid={uid}, svc={svc}, master={mst}, date={dt}, time={tm}", uid=user_id, svc=service, mst=master, dt=date, tm=time)
+    logger.debug("create_booking called: uid={uid}, svc={svc}, mst={mst}, dt={dt}, tm={tm}",
+                 uid=user_id, svc=service, mst=master, dt=date, tm=time)
    bookings = load_bookings()
    booking = {
        "id": len(bookings) + 1,
@@ -0,0 +1,49 @@
+"""Migrate data from JSON files to PostgreSQL."""
+import asyncio
+import json
+from pathlib import Path
+
+from app.db.postgres import PostgresDB
+
+DB_DIR = Path(__file__).resolve().parent.parent.parent / "data"
+USERS_FILE = DB_DIR / "users.json"
+BOOKINGS_FILE = DB_DIR / "bookings.json"
+
+
+async def migrate():
+    print("Starting migration from JSON to PostgreSQL...")
+    await PostgresDB.init_pool()
+
+    # Migrate users
+    if USERS_FILE.exists():
+        with open(USERS_FILE) as f:
+            users = json.load(f)
+        for uid, data in users.items():
+            name = data.get("name", "")
+            phone = data.get("phone", "")
+            await PostgresDB.set_user(int(uid), name, phone)
+        print(f"Migrated {len(users)} users")
+
+    # Migrate bookings
+    if BOOKINGS_FILE.exists():
+        with open(BOOKINGS_FILE) as f:
+            bookings = json.load(f)
+        for b in bookings:
+            # Re-create via raw INSERT to preserve IDs
+            pool = await PostgresDB.get_pool()
+            async with pool.acquire() as conn:
+                await conn.execute("""
+                    INSERT INTO bookings (id, user_id, user_name, service, master, date, time, status)
+                    VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
+                    ON CONFLICT (id) DO NOTHING
+                """, b.get("id"), b.get("user_id"), b.get("user_name", ""),
+                    b.get("service"), b.get("master"), b.get("date"),
+                    b.get("time"), b.get("status", "active"))
+        print(f"Migrated {len(bookings)} bookings")
+
+    print("Migration complete!")
+    await PostgresDB.close()
+
+
+if __name__ == "__main__":
+    asyncio.run(migrate())
@@ -0,0 +1,181 @@
+"""PostgreSQL storage for SPA Telegram Bot."""
+import json
+import os
+from datetime import datetime
+from typing import Optional
+import asyncpg
+from app.logger import logger
+
+DB_HOST = os.getenv("DB_HOST", "10.0.77.9")
+DB_PORT = int(os.getenv("DB_PORT", "5433"))
+DB_USER = os.getenv("DB_USER", "spabot")
+DB_PASSWORD = os.getenv("DB_PASSWORD")
+DB_NAME = os.getenv("DB_NAME", "spabot")
+
+if not DB_PASSWORD:
+    raise ValueError(
+        "DB_PASSWORD not set in environment. "
+        "Set it in .env or export DB_PASSWORD=your_password"
+    )
+
+DB_CONFIG = {
+    "host": DB_HOST,
+    "port": DB_PORT,
+    "user": DB_USER,
+    "password": DB_PASSWORD,
+    "database": DB_NAME,
+}
+
+
+class PostgresDB:
+    _pool: Optional[asyncpg.Pool] = None
+
+    @classmethod
+    async def init_pool(cls):
+        """Initialize connection pool."""
+        if cls._pool is None:
+            cls._pool = await asyncpg.create_pool(**DB_CONFIG, min_size=2, max_size=10)
+            await cls._migrate()
+            logger.info("Postgres pool initialized (host={host}, db={db})",
+                        host=DB_CONFIG["host"], db=DB_CONFIG["database"])
+        return cls._pool
+
+    @classmethod
+    async def _migrate(cls):
+        """Create tables if not exist."""
+        async with cls._pool.acquire() as conn:
+            await conn.execute("""
+                CREATE TABLE IF NOT EXISTS users (
+                    user_id BIGINT PRIMARY KEY,
+                    name TEXT NOT NULL DEFAULT '',
+                    phone TEXT DEFAULT '',
+                    created_at TIMESTAMP DEFAULT NOW(),
+                    updated_at TIMESTAMP DEFAULT NOW()
+                );
+                CREATE TABLE IF NOT EXISTS bookings (
+                    id SERIAL PRIMARY KEY,
+                    user_id BIGINT NOT NULL REFERENCES users(user_id),
+                    user_name TEXT DEFAULT '',
+                    service TEXT NOT NULL,
+                    master TEXT NOT NULL,
+                    date TEXT NOT NULL,
+                    time TEXT NOT NULL,
+                    status TEXT NOT NULL DEFAULT 'active',
+                    created_at TIMESTAMP DEFAULT NOW(),
+                    updated_at TIMESTAMP DEFAULT NOW()
+                );
+                CREATE INDEX IF NOT EXISTS idx_bookings_user_id ON bookings(user_id);
+                CREATE INDEX IF NOT EXISTS idx_bookings_master_date ON bookings(master, date);
+                CREATE INDEX IF NOT EXISTS idx_bookings_status ON bookings(status);
+            """)
+            logger.info("Postgres migrations applied")
+
+    @classmethod
+    async def close(cls):
+        if cls._pool:
+            await cls._pool.close()
+            cls._pool = None
+            logger.info("Postgres pool closed")
+
+    @classmethod
+    async def get_pool(cls):
+        if cls._pool is None:
+            await cls.init_pool()
+        return cls._pool
+
+    # --- Users ---
+    @classmethod
+    async def get_user(cls, user_id: int) -> Optional[dict]:
+        pool = await cls.get_pool()
+        async with pool.acquire() as conn:
+            row = await conn.fetchrow("SELECT * FROM users WHERE user_id = $1", user_id)
+            if row:
+                return dict(row)
+            return None
+
+    @classmethod
+    async def set_user(cls, user_id: int, name: str, phone: str = ""):
+        pool = await cls.get_pool()
+        async with pool.acquire() as conn:
+            await conn.execute("""
+                INSERT INTO users (user_id, name, phone, updated_at)
+                VALUES ($1, $2, $3, NOW())
+                ON CONFLICT (user_id)
+                DO UPDATE SET name = $2, phone = $3, updated_at = NOW()
+            """, user_id, name, phone)
+
+    # --- Bookings ---
+    @classmethod
+    async def create_booking(cls, user_id: int, service: str, master: str,
+                             date: str, time: str, user_name: str = "") -> dict:
+        pool = await cls.get_pool()
+        async with pool.acquire() as conn:
+            row = await conn.fetchrow("""
+                INSERT INTO bookings (user_id, user_name, service, master, date, time)
+                VALUES ($1, $2, $3, $4, $5, $6)
+                RETURNING *
+            """, user_id, user_name, service, master, date, time)
+            return dict(row)
+
+    @classmethod
+    async def get_busy_times(cls, master: str, date: str) -> list:
+        pool = await cls.get_pool()
+        async with pool.acquire() as conn:
+            rows = await conn.fetch(
+                "SELECT time FROM bookings WHERE master = $1 AND date = $2 AND status = 'active'",
+                master, date
+            )
+            return [r["time"] for r in rows]
+
+    @classmethod
+    async def get_user_bookings(cls, user_id: int, status: str = "active") -> list:
+        pool = await cls.get_pool()
+        async with pool.acquire() as conn:
+            rows = await conn.fetch(
+                "SELECT * FROM bookings WHERE user_id = $1 AND status = $2 ORDER BY id DESC",
+                user_id, status
+            )
+            return [dict(r) for r in rows]
+
+    @classmethod
+    async def cancel_booking(cls, booking_id: int) -> bool:
+        pool = await cls.get_pool()
+        async with pool.acquire() as conn:
+            result = await conn.execute(
+                "UPDATE bookings SET status = 'cancelled', updated_at = NOW() WHERE id = $1 AND status = 'active'",
+                booking_id
+            )
+            return "UPDATE 1" in result
+
+    @classmethod
+    async def update_booking(cls, booking_id: int, service: str, master: str,
+                             date: str, time: str) -> Optional[dict]:
+        pool = await cls.get_pool()
+        async with pool.acquire() as conn:
+            row = await conn.fetchrow("""
+                UPDATE bookings
+                SET service = $1, master = $2, date = $3, time = $4, updated_at = NOW()
+                WHERE id = $5 AND status = 'active'
+                RETURNING *
+            """, service, master, date, time, booking_id)
+            return dict(row) if row else None
+
+    @classmethod
+    async def get_all_active_bookings(cls) -> list:
+        pool = await cls.get_pool()
+        async with pool.acquire() as conn:
+            rows = await conn.fetch(
+                "SELECT * FROM bookings WHERE status = 'active' ORDER BY id DESC"
+            )
+            return [dict(r) for r in rows]
+
+    @classmethod
+    async def get_booking_by_id(cls, booking_id: int) -> Optional[dict]:
+        pool = await cls.get_pool()
+        async with pool.acquire() as conn:
+            row = await conn.fetchrow("SELECT * FROM bookings WHERE id = $1", booking_id)
+            return dict(row) if row else None
+
+    @classmethod
+    async def get_user_by_id(cls, user_id: int) -> Optional[dict]:
+        return await cls.get_user(user_id)