database.py - main - mcmare/zp

database.py

56 строк2 KB

mcmare

fix: квотировать имя таблицы в ALTER TABLE — user зарезервированное слово в Postgres

21 июл 2026, 16:46

21 июл 2026, 16:465c9d5dd

import os
from sqlmodel import create_engine, Session, SQLModel
from sqlalchemy import text, inspect
from dotenv import load_dotenv

from logger_config import logger

load_dotenv()

DATABASE_URL = os.getenv("DATABASE_URL", "sqlite:///orders.db")
# Fix for heroku/digitalocean style "postgres://" vs "postgresql://"
if DATABASE_URL.startswith("postgres://"):
    DATABASE_URL = DATABASE_URL.replace("postgres://", "postgresql://", 1)

engine = create_engine(DATABASE_URL)

# Columns added after the initial schema. Applied by init_db() on every startup.
_COLUMN_MIGRATIONS = [
    ("user", "verification_token", "VARCHAR"),
    ("user", "business_trip_mode", "BOOLEAN DEFAULT FALSE"),
    ("order", "is_business_trip", "BOOLEAN DEFAULT FALSE"),
]

def _column_exists(table_name, column_name):
    return column_name in {c["name"] for c in inspect(engine).get_columns(table_name)}

def _ensure_column(conn, table_name, column_name, ddl_type):
    if _column_exists(table_name, column_name):
        return
    quoted_table = f'"{table_name}"'
    try:
        conn.execute(text(f"ALTER TABLE {quoted_table} ADD COLUMN {column_name} {ddl_type}"))
        conn.commit()
        logger.info(f"Migration: added column {table_name}.{column_name}")
    except Exception:
        conn.rollback()
        # This app runs with multiple uvicorn workers (--workers 4): another worker
        # may have added the column concurrently between the check above and this ALTER.
        if _column_exists(table_name, column_name):
            logger.info(f"Migration: column {table_name}.{column_name} added concurrently by another worker")
            return
        logger.exception(
            f"Migration failed: could not add column {table_name}.{column_name} — "
            "check DB user permissions (see migrate_to_postgres.py)"
        )
        raise

def init_db():
    SQLModel.metadata.create_all(engine)
    with engine.connect() as conn:
        for table_name, column_name, ddl_type in _COLUMN_MIGRATIONS:
            _ensure_column(conn, table_name, column_name, ddl_type)

def get_session():
    with Session(engine) as session:
        yield session