"""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)