96 lines
3.6 KiB
Python
96 lines
3.6 KiB
Python
# custom
|
|
from core.errors.logger import logger
|
|
from core.models.DatabaseOperation import DatabaseOperation, DBErrorType
|
|
|
|
# generic
|
|
from pathlib import Path
|
|
from sqlalchemy import Column, Integer, String, create_engine, inspect, text
|
|
from sqlalchemy.orm import declarative_base, Session
|
|
from sqlalchemy.exc import OperationalError, DBAPIError
|
|
from sqlalchemy import inspect
|
|
|
|
"""
|
|
Purpose:
|
|
Sync SQL database schema to current Python model definitions.
|
|
|
|
Situation Used:
|
|
Python Models have new values missing from the same SQL database
|
|
|
|
Metaphor:
|
|
This replaces something like alembic, by doing it directly
|
|
|
|
Confusion:
|
|
This may be confusing because we're using an ORM in general,
|
|
but doing migrations manually.
|
|
|
|
Why:
|
|
The reason is because alembic is good for servers, but not clients,
|
|
due to a large amount of boilerplate for migrations.
|
|
"""
|
|
|
|
def migrate_sql() -> DatabaseOperation:
|
|
from core.models.manage.session_management import engine
|
|
from core.models.orm_models.Base import Base, BaseModel
|
|
|
|
try:
|
|
logger.info("[MIGRATION] Starting Migrations..")
|
|
inspector = inspect(engine)
|
|
existing_tables = inspector.get_table_names()
|
|
changes_made = []
|
|
|
|
for table in Base.metadata.tables.values():
|
|
table_name = table.name
|
|
|
|
if table_name not in existing_tables:
|
|
table.create(engine)
|
|
added_table = f"Created table: {table_name}"
|
|
changes_made.append(added_table)
|
|
logger.info(f"[MIGRATION] {added_table}")
|
|
continue
|
|
|
|
# Check for MISSING columns (add them)
|
|
existing_columns = {col['name'] for col in inspector.get_columns(table_name)}
|
|
model_columns = {col.name for col in table.columns}
|
|
missing_columns = model_columns - existing_columns
|
|
|
|
if missing_columns:
|
|
with engine.begin() as conn:
|
|
for col_name in missing_columns:
|
|
col = table.columns[col_name]
|
|
col_type = str(col.type.compile(engine.dialect))
|
|
alter_sql = f"ALTER TABLE {table_name} ADD COLUMN {col_name} {col_type}"
|
|
conn.execute(text(alter_sql))
|
|
|
|
changes_made.append(f"Added columns to {table_name}: {missing_columns}")
|
|
|
|
# Check for EXTRA columns (drop them)
|
|
extra_columns = existing_columns - model_columns
|
|
if extra_columns:
|
|
with engine.begin() as conn:
|
|
for col_name in extra_columns:
|
|
alter_sql = f"ALTER TABLE {table_name} DROP COLUMN {col_name}"
|
|
logger.info(f"[MIGRATION] Dropping extra column: {alter_sql}")
|
|
conn.execute(text(alter_sql))
|
|
|
|
changes_made.append(f"Dropped extra columns from {table_name}: {extra_columns}")
|
|
|
|
if changes_made:
|
|
logger.info(f"[MIGRATION] Completed! Schema synced successfully")
|
|
return DatabaseOperation(
|
|
valid=True,
|
|
data=changes_made,
|
|
message=f"Schema synced: {'; '.join(changes_made)}"
|
|
)
|
|
else:
|
|
return DatabaseOperation(
|
|
valid=True,
|
|
message="Schema already matches models"
|
|
)
|
|
|
|
except (OperationalError, DBAPIError) as e:
|
|
logger.error(f"[MIGRATION] Failed: {str(e)}")
|
|
return DatabaseOperation(
|
|
valid=False,
|
|
error_type=DBErrorType.UNKNOWN,
|
|
message=f"Schema sync failed: {str(e)}"
|
|
)
|