278 lines
10 KiB
Python
278 lines
10 KiB
Python
import logging
|
|
from pathlib import Path
|
|
from typing import Tuple
|
|
|
|
from alembic import command
|
|
from alembic.config import Config
|
|
from sqlalchemy import inspect, text
|
|
from sqlalchemy.engine import Connection
|
|
from sqlalchemy.ext.asyncio import AsyncEngine
|
|
|
|
from config.settings import Settings
|
|
|
|
|
|
_BASELINE_REVISION = "0001_initial_schema"
|
|
|
|
|
|
def _build_alembic_config(settings: Settings) -> Config:
|
|
project_root = Path(__file__).resolve().parents[1]
|
|
config = Config(str(project_root / "alembic.ini"))
|
|
config.set_main_option("script_location", str(project_root / "alembic"))
|
|
config.set_main_option("sqlalchemy.url", settings.DATABASE_URL)
|
|
return config
|
|
|
|
|
|
def _inspect_database_state(connection: Connection) -> Tuple[bool, bool, bool]:
|
|
db_inspector = inspect(connection)
|
|
has_alembic_version = db_inspector.has_table("alembic_version")
|
|
has_users_table = db_inspector.has_table("users")
|
|
has_legacy_migrator_table = db_inspector.has_table("schema_migrations")
|
|
return has_alembic_version, has_users_table, has_legacy_migrator_table
|
|
|
|
|
|
def _run_legacy_migrator_compatibility(connection: Connection) -> None:
|
|
db_inspector = inspect(connection)
|
|
if not db_inspector.has_table("users"):
|
|
return
|
|
|
|
users_columns = {
|
|
column["name"]
|
|
for column in db_inspector.get_columns("users")
|
|
}
|
|
user_alter_statements = []
|
|
|
|
if "channel_subscription_verified" not in users_columns:
|
|
user_alter_statements.append(
|
|
"ALTER TABLE users ADD COLUMN channel_subscription_verified BOOLEAN"
|
|
)
|
|
if "channel_subscription_checked_at" not in users_columns:
|
|
user_alter_statements.append(
|
|
"ALTER TABLE users ADD COLUMN channel_subscription_checked_at TIMESTAMPTZ"
|
|
)
|
|
if "channel_subscription_verified_for" not in users_columns:
|
|
user_alter_statements.append(
|
|
"ALTER TABLE users ADD COLUMN channel_subscription_verified_for BIGINT"
|
|
)
|
|
if "referral_code" not in users_columns:
|
|
user_alter_statements.append(
|
|
"ALTER TABLE users ADD COLUMN referral_code VARCHAR(16)"
|
|
)
|
|
|
|
for statement in user_alter_statements:
|
|
connection.execute(text(statement))
|
|
|
|
users_columns = {
|
|
column["name"]
|
|
for column in inspect(connection).get_columns("users")
|
|
}
|
|
if "referral_code" in users_columns:
|
|
connection.execute(
|
|
text(
|
|
"""
|
|
WITH generated_codes AS (
|
|
SELECT
|
|
user_id,
|
|
UPPER(
|
|
SUBSTRING(
|
|
md5(
|
|
user_id::text
|
|
|| clock_timestamp()::text
|
|
|| random()::text
|
|
)
|
|
FROM 1 FOR 9
|
|
)
|
|
) AS referral_code
|
|
FROM users
|
|
WHERE referral_code IS NULL OR referral_code = ''
|
|
)
|
|
UPDATE users AS u
|
|
SET referral_code = g.referral_code
|
|
FROM generated_codes AS g
|
|
WHERE u.user_id = g.user_id
|
|
"""
|
|
)
|
|
)
|
|
connection.execute(
|
|
text(
|
|
"""
|
|
CREATE UNIQUE INDEX IF NOT EXISTS uq_users_referral_code
|
|
ON users (referral_code)
|
|
WHERE referral_code IS NOT NULL
|
|
"""
|
|
)
|
|
)
|
|
connection.execute(
|
|
text(
|
|
"""
|
|
UPDATE users
|
|
SET referral_code = UPPER(referral_code)
|
|
WHERE referral_code IS NOT NULL
|
|
AND referral_code <> UPPER(referral_code)
|
|
"""
|
|
)
|
|
)
|
|
|
|
db_inspector = inspect(connection)
|
|
if db_inspector.has_table("payments"):
|
|
payments_columns = {
|
|
column["name"]
|
|
for column in db_inspector.get_columns("payments")
|
|
}
|
|
if "original_amount" not in payments_columns:
|
|
connection.execute(text("ALTER TABLE payments ADD COLUMN original_amount FLOAT"))
|
|
if "discount_applied" not in payments_columns:
|
|
connection.execute(text("ALTER TABLE payments ADD COLUMN discount_applied FLOAT"))
|
|
|
|
db_inspector = inspect(connection)
|
|
has_promo_codes = db_inspector.has_table("promo_codes")
|
|
if has_promo_codes:
|
|
promo_columns = {
|
|
column["name"]
|
|
for column in db_inspector.get_columns("promo_codes")
|
|
}
|
|
if "promo_type" not in promo_columns:
|
|
connection.execute(
|
|
text(
|
|
"ALTER TABLE promo_codes ADD COLUMN promo_type VARCHAR NOT NULL DEFAULT 'bonus_days'"
|
|
)
|
|
)
|
|
if "discount_percentage" not in promo_columns:
|
|
connection.execute(
|
|
text("ALTER TABLE promo_codes ADD COLUMN discount_percentage INTEGER")
|
|
)
|
|
if "bonus_days" in promo_columns:
|
|
connection.execute(
|
|
text("ALTER TABLE promo_codes ALTER COLUMN bonus_days DROP NOT NULL")
|
|
)
|
|
connection.execute(
|
|
text(
|
|
"CREATE INDEX IF NOT EXISTS idx_promo_codes_promo_type ON promo_codes (promo_type)"
|
|
)
|
|
)
|
|
|
|
db_inspector = inspect(connection)
|
|
has_active_discounts = db_inspector.has_table("active_discounts")
|
|
if not has_active_discounts and has_promo_codes:
|
|
connection.execute(
|
|
text(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS active_discounts (
|
|
user_id BIGINT PRIMARY KEY,
|
|
promo_code_id INTEGER NOT NULL,
|
|
discount_percentage INTEGER NOT NULL,
|
|
activated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
|
CONSTRAINT fk_active_discounts_user
|
|
FOREIGN KEY (user_id) REFERENCES users (user_id) ON DELETE CASCADE,
|
|
CONSTRAINT fk_active_discounts_promo_code
|
|
FOREIGN KEY (promo_code_id) REFERENCES promo_codes (promo_code_id) ON DELETE CASCADE
|
|
)
|
|
"""
|
|
)
|
|
)
|
|
has_active_discounts = True
|
|
|
|
if has_active_discounts and has_promo_codes:
|
|
connection.execute(
|
|
text(
|
|
"DELETE FROM active_discounts ad "
|
|
"WHERE NOT EXISTS (SELECT 1 FROM users u WHERE u.user_id = ad.user_id) "
|
|
"OR NOT EXISTS (SELECT 1 FROM promo_codes p WHERE p.promo_code_id = ad.promo_code_id)"
|
|
)
|
|
)
|
|
connection.execute(
|
|
text(
|
|
"ALTER TABLE active_discounts "
|
|
"DROP CONSTRAINT IF EXISTS active_discounts_user_id_fkey"
|
|
)
|
|
)
|
|
connection.execute(
|
|
text(
|
|
"ALTER TABLE active_discounts "
|
|
"DROP CONSTRAINT IF EXISTS fk_active_discounts_user"
|
|
)
|
|
)
|
|
connection.execute(
|
|
text(
|
|
"ALTER TABLE active_discounts "
|
|
"DROP CONSTRAINT IF EXISTS active_discounts_promo_code_id_fkey"
|
|
)
|
|
)
|
|
connection.execute(
|
|
text(
|
|
"ALTER TABLE active_discounts "
|
|
"DROP CONSTRAINT IF EXISTS fk_active_discounts_promo_code"
|
|
)
|
|
)
|
|
connection.execute(
|
|
text(
|
|
"ALTER TABLE active_discounts "
|
|
"ADD CONSTRAINT fk_active_discounts_user "
|
|
"FOREIGN KEY (user_id) REFERENCES users (user_id) ON DELETE CASCADE"
|
|
)
|
|
)
|
|
connection.execute(
|
|
text(
|
|
"ALTER TABLE active_discounts "
|
|
"ADD CONSTRAINT fk_active_discounts_promo_code "
|
|
"FOREIGN KEY (promo_code_id) REFERENCES promo_codes (promo_code_id) ON DELETE CASCADE"
|
|
)
|
|
)
|
|
elif has_active_discounts and not has_promo_codes:
|
|
logging.warning(
|
|
"Alembic legacy compatibility: skipped active_discounts FK repair "
|
|
"because promo_codes table is missing."
|
|
)
|
|
|
|
|
|
def _run_stamp(connection: Connection, alembic_config: Config, revision: str) -> None:
|
|
alembic_config.attributes["connection"] = connection
|
|
command.stamp(alembic_config, revision)
|
|
|
|
|
|
def _run_upgrade(connection: Connection, alembic_config: Config) -> None:
|
|
alembic_config.attributes["connection"] = connection
|
|
command.upgrade(alembic_config, "head")
|
|
|
|
|
|
async def run_alembic_migrations(settings: Settings, async_engine: AsyncEngine) -> None:
|
|
"""Apply Alembic migrations with bootstrap for existing installations."""
|
|
|
|
alembic_config = _build_alembic_config(settings)
|
|
|
|
async with async_engine.begin() as async_connection:
|
|
(
|
|
has_alembic_version,
|
|
has_users_table,
|
|
has_legacy_migrator_table,
|
|
) = await async_connection.run_sync(
|
|
_inspect_database_state
|
|
)
|
|
|
|
if not has_alembic_version and has_users_table:
|
|
if not has_legacy_migrator_table:
|
|
raise RuntimeError(
|
|
"Alembic bootstrap refused: found existing users table without "
|
|
"alembic_version and without legacy schema_migrations marker. "
|
|
"Cannot safely determine migration baseline."
|
|
)
|
|
|
|
logging.info(
|
|
"Alembic: applying legacy migrator compatibility fixes before stamp."
|
|
)
|
|
await async_connection.run_sync(_run_legacy_migrator_compatibility)
|
|
|
|
logging.info(
|
|
"Alembic: existing schema detected without alembic_version; stamping %s.",
|
|
_BASELINE_REVISION,
|
|
)
|
|
await async_connection.run_sync(
|
|
_run_stamp,
|
|
alembic_config,
|
|
_BASELINE_REVISION,
|
|
)
|
|
|
|
logging.info("Alembic: running upgrade to head...")
|
|
await async_connection.run_sync(_run_upgrade, alembic_config)
|
|
|
|
logging.info("Alembic: migrations applied successfully.")
|