demo_spa_bot/app/db/postgres.py

"""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_CONFIG = {
    "host": os.getenv("DB_HOST", "10.0.77.9"),
    "port": int(os.getenv("DB_PORT", "5432")),
    "user": os.getenv("DB_USER", "spabot"),
    "password": os.getenv("DB_PASSWORD", "spabot_password"),
    "database": os.getenv("DB_NAME", "spabot"),
}


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)