wehub-resource-sync
276 lines
10 KiB
Python
276 lines
10 KiB
Python
import asyncio
|
|
from pathlib import Path
|
|
|
|
from alembic import command
|
|
from alembic.config import Config
|
|
from alembic.script import ScriptDirectory
|
|
from sqlalchemy import create_engine, inspect, text
|
|
|
|
from utils.db_utils import get_database_url_and_connect_args, to_sync_sqlalchemy_url
|
|
from utils.get_env import get_migrate_database_on_startup_env
|
|
|
|
|
|
LEGACY_BASELINE_REVISION = "00b3c27a13bc"
|
|
# Revision before 95b5127e93cd (template_create_infos); used when DB has theme but not that table.
|
|
REVISION_BEFORE_TEMPLATE_CREATE_INFO = "82abdbc476a7"
|
|
REVISION_TEMPLATE_CREATE_INFO = "95b5127e93cd"
|
|
REVISION_CHAT_HISTORY = "c7b70d0f31b1"
|
|
REVISION_TEMPLATE_V2 = "6e4a1b2c3d5f"
|
|
REVISION_SLIDE_UI = "7f5b2c3d4e6a"
|
|
REVISION_MERGED_TEMPLATE_V2 = "8a6c4d2e1f30"
|
|
REVISION_PRESENTATION_FONTS = "9b2d1c4e5f6a"
|
|
REVISION_TEMPLATE_V2_CHAT_SCOPE = "1d9a4c7b8e2f"
|
|
REVISION_TEMPLATE_V2_LAYOUTS_OPTIONAL = "2c8f4a1b9d7e"
|
|
REVISION_FONT_UPLOADS = "5d7e9a1b2c3f"
|
|
REVISION_TEMPLATE_V2_ID_STRINGS = "3f2a1b4c5d6e"
|
|
REVISION_TEMPLATE_V2_IS_DEFAULT = "4b7c9d0e1f2a"
|
|
REVISION_ASYNC_TASKS = "a7d4c9e2f1b3"
|
|
|
|
|
|
async def migrate_database_on_startup() -> None:
|
|
if get_migrate_database_on_startup_env() not in ["true", "True"]:
|
|
return
|
|
|
|
try:
|
|
await asyncio.to_thread(_run_migrations)
|
|
print("Migrations run successfully", flush=True)
|
|
except Exception as exc:
|
|
print(f"Error running migrations: {exc}", flush=True)
|
|
raise
|
|
|
|
|
|
def _run_migrations() -> None:
|
|
# migrations.py lives at servers/fastapi/migrations.py
|
|
# so parents[0] = servers/fastapi/, where alembic/ lives alongside it.
|
|
base_dir = Path(__file__).resolve().parents[0]
|
|
config = Config()
|
|
config.set_main_option("script_location", str(base_dir / "alembic"))
|
|
|
|
database_url, _ = get_database_url_and_connect_args()
|
|
|
|
# Alembic uses synchronous engines; strip async driver prefixes.
|
|
database_url = to_sync_sqlalchemy_url(database_url)
|
|
|
|
config.set_main_option("sqlalchemy.url", database_url)
|
|
_repair_orphan_alembic_revision(config, database_url)
|
|
_stamp_legacy_database_if_needed(config, database_url)
|
|
|
|
try:
|
|
command.upgrade(config, "head")
|
|
except Exception:
|
|
# Safety net for edge cases; legacy DBs are stamped proactively above.
|
|
if _is_unversioned_populated_database(database_url):
|
|
_stamp_legacy_database_if_needed(config, database_url)
|
|
command.upgrade(config, "head")
|
|
return
|
|
raise
|
|
|
|
|
|
def _repair_orphan_alembic_revision(config: Config, database_url: str) -> None:
|
|
"""
|
|
If alembic_version points at a revision id that no longer exists in alembic/versions
|
|
(removed branch, old image, etc.), re-stamp from the live schema so upgrade can run.
|
|
"""
|
|
script = ScriptDirectory.from_config(config)
|
|
known = {rev.revision for rev in script.walk_revisions()}
|
|
heads = script.get_heads()
|
|
if len(heads) != 1:
|
|
return
|
|
head = heads[0]
|
|
|
|
engine = create_engine(database_url)
|
|
try:
|
|
with engine.begin() as connection:
|
|
inspector = inspect(connection)
|
|
tables = set(inspector.get_table_names())
|
|
if "alembic_version" not in tables:
|
|
return
|
|
version_num = connection.execute(
|
|
text("SELECT version_num FROM alembic_version LIMIT 1")
|
|
).scalar_one_or_none()
|
|
if not version_num or version_num in known:
|
|
return
|
|
print(
|
|
f"Alembic revision {version_num!r} is missing from the codebase; "
|
|
"inferring applied migrations from schema and re-stamping.",
|
|
flush=True,
|
|
)
|
|
target = _infer_revision_from_schema(inspector, tables, head)
|
|
connection.execute(
|
|
text("UPDATE alembic_version SET version_num = :revision"),
|
|
{"revision": target},
|
|
)
|
|
finally:
|
|
engine.dispose()
|
|
|
|
|
|
def _infer_revision_from_schema(inspector, tables: set[str], head_revision: str) -> str:
|
|
"""Best-effort: map existing SQLite/Postgres schema to our linear migration chain."""
|
|
if "template_v2" in tables:
|
|
cols = {c["name"] for c in inspector.get_columns("template_v2")}
|
|
final_template_columns = {
|
|
"description",
|
|
"raw_layouts",
|
|
"components",
|
|
"layouts",
|
|
"assets",
|
|
}
|
|
presentation_version_ready = (
|
|
"presentations" not in tables
|
|
or _has_presentation_version_column(inspector, tables)
|
|
)
|
|
slide_ui_ready = "slides" not in tables or _has_column(
|
|
inspector, "slides", "ui"
|
|
)
|
|
presentation_fonts_ready = "presentations" not in tables or _has_column(
|
|
inspector, "presentations", "fonts"
|
|
)
|
|
template_v2_chat_scope_ready = (
|
|
"chat_history_messages" not in tables
|
|
or _has_column(inspector, "chat_history_messages", "template_v2_id")
|
|
)
|
|
font_uploads_ready = "font_uploads" in tables
|
|
template_v2_id_strings_ready = _has_template_v2_id_string_columns(
|
|
inspector,
|
|
tables,
|
|
)
|
|
template_v2_is_default_ready = _has_column(
|
|
inspector, "template_v2", "is_default"
|
|
)
|
|
async_tasks_ready = "async_tasks" in tables
|
|
if (
|
|
final_template_columns.issubset(cols)
|
|
and not {"cluster_candidates", "clusters"}.intersection(cols)
|
|
and presentation_version_ready
|
|
):
|
|
if slide_ui_ready and presentation_fonts_ready and template_v2_chat_scope_ready:
|
|
if not font_uploads_ready:
|
|
return REVISION_TEMPLATE_V2_LAYOUTS_OPTIONAL
|
|
if not template_v2_id_strings_ready:
|
|
return REVISION_FONT_UPLOADS
|
|
if not template_v2_is_default_ready:
|
|
return REVISION_TEMPLATE_V2_ID_STRINGS
|
|
return (
|
|
head_revision
|
|
if async_tasks_ready
|
|
else REVISION_TEMPLATE_V2_IS_DEFAULT
|
|
)
|
|
if slide_ui_ready and presentation_fonts_ready:
|
|
return REVISION_PRESENTATION_FONTS
|
|
return REVISION_MERGED_TEMPLATE_V2 if slide_ui_ready else REVISION_TEMPLATE_V2
|
|
return REVISION_CHAT_HISTORY
|
|
if "chat_history_messages" in tables:
|
|
return REVISION_CHAT_HISTORY
|
|
if "template_create_infos" in tables:
|
|
return REVISION_TEMPLATE_CREATE_INFO
|
|
if "presentations" in tables:
|
|
cols = {c["name"] for c in inspector.get_columns("presentations")}
|
|
if "theme" in cols:
|
|
return REVISION_BEFORE_TEMPLATE_CREATE_INFO
|
|
return LEGACY_BASELINE_REVISION
|
|
|
|
|
|
def _has_presentation_version_column(inspector, tables: set[str]) -> bool:
|
|
if "presentations" not in tables:
|
|
return False
|
|
|
|
cols = {c["name"] for c in inspector.get_columns("presentations")}
|
|
return "version" in cols
|
|
|
|
|
|
def _has_column(inspector, table_name: str, column_name: str) -> bool:
|
|
columns = {column["name"] for column in inspector.get_columns(table_name)}
|
|
return column_name in columns
|
|
|
|
|
|
def _has_template_v2_id_string_columns(inspector, tables: set[str]) -> bool:
|
|
if "template_v2" not in tables or not _has_column(inspector, "template_v2", "id"):
|
|
return False
|
|
if _is_uuid_storage_column(inspector, "template_v2", "id"):
|
|
return False
|
|
if "chat_history_messages" in tables and _has_column(
|
|
inspector,
|
|
"chat_history_messages",
|
|
"template_v2_id",
|
|
):
|
|
return not _is_uuid_storage_column(
|
|
inspector,
|
|
"chat_history_messages",
|
|
"template_v2_id",
|
|
)
|
|
return True
|
|
|
|
|
|
def _is_uuid_storage_column(inspector, table_name: str, column_name: str) -> bool:
|
|
for column in inspector.get_columns(table_name):
|
|
if column["name"] != column_name:
|
|
continue
|
|
column_type = column["type"]
|
|
type_class = column_type.__class__.__name__.lower()
|
|
rendered_type = str(column_type).lower().replace(" ", "")
|
|
return type_class == "uuid" or rendered_type in {"uuid", "char(32)"}
|
|
return False
|
|
|
|
|
|
def _stamp_legacy_database_if_needed(config: Config, database_url: str) -> None:
|
|
"""
|
|
If the DB has app tables but no migration reference in alembic_version,
|
|
treat it as a legacy DB and stamp the latest revision already reflected by
|
|
the live schema before upgrading.
|
|
"""
|
|
if not _is_unversioned_populated_database(database_url):
|
|
return
|
|
|
|
script = ScriptDirectory.from_config(config)
|
|
heads = script.get_heads()
|
|
head = heads[0] if len(heads) == 1 else script.get_base()
|
|
engine = create_engine(database_url)
|
|
try:
|
|
with engine.connect() as connection:
|
|
inspector = inspect(connection)
|
|
target_revision = _infer_revision_from_schema(
|
|
inspector, set(inspector.get_table_names()), head
|
|
)
|
|
finally:
|
|
engine.dispose()
|
|
|
|
print(
|
|
"Detected legacy database without migration reference. "
|
|
f"Stamping revision to {target_revision} before upgrading.",
|
|
flush=True,
|
|
)
|
|
command.stamp(config, target_revision)
|
|
|
|
|
|
def _is_unversioned_populated_database(database_url: str) -> bool:
|
|
known_app_tables = {
|
|
"presentations",
|
|
"slides",
|
|
"templates",
|
|
"keyvaluesqlmodel",
|
|
"imageasset",
|
|
"presentation_layout_codes",
|
|
"async_presentation_generation_tasks",
|
|
"webhook_subscriptions",
|
|
"template_create_infos",
|
|
"chat_history_messages",
|
|
"template_v2",
|
|
"font_uploads",
|
|
}
|
|
engine = create_engine(database_url)
|
|
try:
|
|
with engine.connect() as connection:
|
|
inspector = inspect(connection)
|
|
table_names = set(inspector.get_table_names())
|
|
has_alembic_version_table = "alembic_version" in table_names
|
|
has_applied_revision = False
|
|
if has_alembic_version_table:
|
|
revision_count = connection.execute(
|
|
text("SELECT COUNT(*) FROM alembic_version")
|
|
).scalar_one()
|
|
has_applied_revision = revision_count > 0
|
|
has_known_app_tables = len(table_names.intersection(known_app_tables)) > 0
|
|
return has_known_app_tables and not has_applied_revision
|
|
finally:
|
|
engine.dispose()
|