RemnaWave-bot Update

app/db/session.py

@@ -17,6 +17,9 @@ 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(
@@ -121,6 +124,60 @@ async def _normalize_payment_order_provision_username_index(
     )
 
 
+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,
@@ -146,7 +203,60 @@ async def ensure_database_exists(
         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)