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

2026-07-31 12:58:23 +04:00
parent bf26370da9
commit f289631645
36 changed files with 893 additions and 290 deletions
@@ -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)