263 lines
7.6 KiB
Python
263 lines
7.6 KiB
Python
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)
|