import logging from sqlalchemy import text from sqlalchemy.ext.asyncio import ( AsyncConnection, AsyncEngine, AsyncSession, async_sessionmaker, create_async_engine, ) from app.db.base import Base from app.db import models as _models # noqa: F401 logger = logging.getLogger(__name__) PAYMENT_ORDERS_TABLE_NAME = "payment_orders" PAYMENT_ORDERS_PROVISION_USERNAME_COLUMN = "provision_username" PAYMENT_ORDERS_PROVISION_USERNAME_INDEX = "ix_payment_orders_provision_username" SUPPORT_TICKETS_TABLE_NAME = "support_tickets" REFERRAL_INVITES_TABLE_NAME = "referral_invites" TELEGRAM_USERS_TABLE_NAME = "telegram_users" def create_engine_and_session_factory( database_url: str, *, echo: bool = False, ) -> tuple[AsyncEngine, async_sessionmaker[AsyncSession]]: engine = create_async_engine( database_url, echo=echo, pool_pre_ping=True, pool_recycle=3600, ) session_factory = async_sessionmaker(engine, expire_on_commit=False, class_=AsyncSession) return engine, session_factory def _quote_mysql_identifier(identifier: str) -> str: return f"`{identifier.replace('`', '``')}`" async def _get_mysql_single_column_indexes( connection: AsyncConnection, *, table_name: str, column_name: str, non_unique: int, ) -> list[str]: result = await connection.execute( text( """ SELECT index_name FROM information_schema.statistics WHERE table_schema = DATABASE() AND table_name = :table_name GROUP BY index_name, non_unique HAVING non_unique = :non_unique AND COUNT(*) = 1 AND MAX(CASE WHEN column_name = :column_name THEN 1 ELSE 0 END) = 1 """ ), { "table_name": table_name, "column_name": column_name, "non_unique": non_unique, }, ) return [row[0] for row in result] async def _normalize_payment_order_provision_username_index( connection: AsyncConnection, ) -> None: if connection.dialect.name != "mysql": return unique_index_names = await _get_mysql_single_column_indexes( connection, table_name=PAYMENT_ORDERS_TABLE_NAME, column_name=PAYMENT_ORDERS_PROVISION_USERNAME_COLUMN, non_unique=0, ) quoted_table_name = _quote_mysql_identifier(PAYMENT_ORDERS_TABLE_NAME) for index_name in unique_index_names: logger.info( "Dropping stale UNIQUE index %s on %s.%s", index_name, PAYMENT_ORDERS_TABLE_NAME, PAYMENT_ORDERS_PROVISION_USERNAME_COLUMN, ) await connection.execute( text( "DROP INDEX " f"{_quote_mysql_identifier(index_name)} " f"ON {quoted_table_name}" ) ) non_unique_index_names = await _get_mysql_single_column_indexes( connection, table_name=PAYMENT_ORDERS_TABLE_NAME, column_name=PAYMENT_ORDERS_PROVISION_USERNAME_COLUMN, non_unique=1, ) if non_unique_index_names: return logger.info( "Creating non-unique index %s on %s.%s", PAYMENT_ORDERS_PROVISION_USERNAME_INDEX, PAYMENT_ORDERS_TABLE_NAME, PAYMENT_ORDERS_PROVISION_USERNAME_COLUMN, ) await connection.execute( text( "CREATE INDEX " f"{_quote_mysql_identifier(PAYMENT_ORDERS_PROVISION_USERNAME_INDEX)} " f"ON {quoted_table_name} " f"({_quote_mysql_identifier(PAYMENT_ORDERS_PROVISION_USERNAME_COLUMN)})" ) ) async def _get_mysql_column_names( connection: AsyncConnection, *, table_name: str, ) -> set[str]: result = await connection.execute( text( """ SELECT column_name FROM information_schema.columns WHERE table_schema = DATABASE() AND table_name = :table_name """ ), {"table_name": table_name}, ) return {str(row[0]) for row in result} async def _normalize_support_ticket_columns(connection: AsyncConnection) -> None: if connection.dialect.name != "mysql": return existing_columns = await _get_mysql_column_names( connection, table_name=SUPPORT_TICKETS_TABLE_NAME, ) quoted_table_name = _quote_mysql_identifier(SUPPORT_TICKETS_TABLE_NAME) required_columns = { "support_note": "TEXT NULL", "closed_by_telegram_id": "BIGINT NULL", "closed_at": "DATETIME NULL", } for column_name, ddl in required_columns.items(): if column_name in existing_columns: continue logger.info( "Adding missing column %s to %s", column_name, SUPPORT_TICKETS_TABLE_NAME, ) await connection.execute( text( "ALTER TABLE " f"{quoted_table_name} " "ADD COLUMN " f"{_quote_mysql_identifier(column_name)} {ddl}" ) ) async def ensure_database_exists( server_url: str, database_name: str, *, echo: bool = False, ) -> None: engine = create_async_engine( server_url, echo=echo, pool_pre_ping=True, pool_recycle=3600, ) try: async with engine.begin() as connection: await connection.execute( text( "CREATE DATABASE IF NOT EXISTS " f"{_quote_mysql_identifier(database_name)} " "CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci" ) ) finally: await engine.dispose() async def _normalize_referral_invites_columns(connection: AsyncConnection) -> None: """Add discount_used column to referral_invites if missing.""" if connection.dialect.name != "mysql": return existing_columns = await _get_mysql_column_names( connection, table_name=REFERRAL_INVITES_TABLE_NAME, ) quoted_table_name = _quote_mysql_identifier(REFERRAL_INVITES_TABLE_NAME) if "discount_used" not in existing_columns: logger.info( "Adding missing column discount_used to %s", REFERRAL_INVITES_TABLE_NAME, ) await connection.execute( text( "ALTER TABLE " f"{quoted_table_name} " "ADD COLUMN `discount_used` BOOLEAN NOT NULL DEFAULT FALSE" ) ) async def _normalize_telegram_users_columns(connection: AsyncConnection) -> None: """Add discount_ever_used column to telegram_users if missing.""" if connection.dialect.name != "mysql": return existing_columns = await _get_mysql_column_names( connection, table_name=TELEGRAM_USERS_TABLE_NAME, ) quoted_table_name = _quote_mysql_identifier(TELEGRAM_USERS_TABLE_NAME) if "discount_ever_used" not in existing_columns: logger.info( "Adding missing column discount_ever_used to %s", TELEGRAM_USERS_TABLE_NAME, ) await connection.execute( text( "ALTER TABLE " f"{quoted_table_name} " "ADD COLUMN `discount_ever_used` BOOLEAN NOT NULL DEFAULT FALSE" ) ) async def init_db(engine: AsyncEngine) -> None: async with engine.begin() as connection: await connection.run_sync(Base.metadata.create_all) await _normalize_payment_order_provision_username_index(connection) await _normalize_support_ticket_columns(connection) await _normalize_referral_invites_columns(connection) await _normalize_telegram_users_columns(connection)