148 lines
5.6 KiB
Python
148 lines
5.6 KiB
Python
# custom
|
|
from core.errors.logger import logger
|
|
from core.models.DatabaseOperation import DatabaseOperation, DBErrorType
|
|
from core.Constants import Constants
|
|
|
|
from core.models.manage.session_management import create_ALL_tables
|
|
# from core.models.orm_models.Location import Location
|
|
# from core.models.orm_models.Operator import Operator
|
|
# from core.models.orm_models.CachedSync import CachedSync
|
|
# from core.models.orm_models.EncryptedProxy import EncryptedProxy
|
|
|
|
# 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
|
|
from sqlalchemy import text
|
|
|
|
"""
|
|
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 get_model_columns(model_class):
|
|
"""
|
|
Extract all columns directly from a model class.
|
|
Bypasses metadata, caching, and all SQLAlchemy indirection.
|
|
"""
|
|
try:
|
|
mapper = inspect(model_class)
|
|
columns = {}
|
|
|
|
for column in mapper.columns:
|
|
columns[column.name] = {
|
|
'type': str(column.type),
|
|
'nullable': column.nullable,
|
|
'default': column.default,
|
|
'primary_key': column.primary_key,
|
|
}
|
|
|
|
return columns
|
|
except Exception as e:
|
|
raise ValueError(f"Failed to inspect {model_class}: {e}")
|
|
|
|
|
|
def migrate_sql() -> DatabaseOperation:
|
|
from core.models.manage.session_management import engine
|
|
from core.models.orm_models.Base import BaseModel, Base
|
|
try:
|
|
from core.models.orm_models.Location import Location
|
|
from core.models.orm_models.Operator import Operator
|
|
from core.models.orm_models.CachedSync import CachedSync
|
|
from core.models.orm_models.EncryptedProxy import EncryptedProxy
|
|
MODELS = [Location, Operator, CachedSync, EncryptedProxy]
|
|
logger.info(f"[MIGRATION] Tables loaded")
|
|
except:
|
|
logger.error(f"[MIGRATION] Could not load models. Critical failure.")
|
|
|
|
changes_made = []
|
|
|
|
# Step 1: Extract column definitions directly from each model
|
|
model_schema = {}
|
|
for model_class in MODELS:
|
|
table_name = model_class.__tablename__
|
|
try:
|
|
model_schema[table_name] = get_model_columns(model_class)
|
|
except ValueError as e:
|
|
return DatabaseOperation(
|
|
valid=False,
|
|
error_type=DBErrorType.PYTHON_MODEL_STRUCTURE,
|
|
message=e
|
|
)
|
|
logger.info(f"[MIGRATION] Model '{table_name}': {list(model_schema[table_name].keys())}")
|
|
|
|
# Step 2: Compare to database schema
|
|
inspector = inspect(engine)
|
|
db_tables = set(inspector.get_table_names())
|
|
|
|
for table_name, expected_columns in model_schema.items():
|
|
if table_name not in db_tables:
|
|
changes_made.append(table_name)
|
|
logger.info(f"[MIGRATION] Table '{table_name}' is MISSING from database")
|
|
else:
|
|
db_columns = {col['name']: col for col in inspector.get_columns(table_name)}
|
|
|
|
for col_name, col_info in expected_columns.items():
|
|
if col_name not in db_columns:
|
|
changes_made.append(col_name)
|
|
logger.info(f"[MIGRATION] Column '{col_name}' is MISSING from table '{table_name}'")
|
|
|
|
# Step 3: Perform the actual migration
|
|
if changes_made:
|
|
logger.info("[MIGRATION] Changes detected. Running schema update...")
|
|
Base.metadata.create_all(engine)
|
|
|
|
try:
|
|
with engine.begin() as conn:
|
|
for table_name, expected_columns in model_schema.items():
|
|
inspector = inspect(engine)
|
|
|
|
if table_name in inspector.get_table_names():
|
|
db_columns = {col['name'] for col in inspector.get_columns(table_name)}
|
|
|
|
for col_name, col_info in expected_columns.items():
|
|
if col_name not in db_columns:
|
|
col_type = col_info['type']
|
|
nullable = "NULL" if col_info['nullable'] else "NOT NULL"
|
|
print(f"[MIGRATION] Adding column '{col_name}' to '{table_name}'")
|
|
conn.execute(text(f"ALTER TABLE {table_name} ADD COLUMN {col_name} {col_type} {nullable}"))
|
|
|
|
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)}"
|
|
)
|
|
|
|
logger.info(f"[MIGRATION] Completed! Schema synced successfully with these changes: {changes_made}")
|
|
return DatabaseOperation(
|
|
valid=True,
|
|
data=changes_made,
|
|
message=f"Synced these tables and columns: {'; '.join(changes_made)}"
|
|
)
|
|
|
|
else:
|
|
no_migration_needed = "Schema already matches models. No migration needed."
|
|
logger.info(f"[MIGRATION] {no_migration_needed}")
|
|
return DatabaseOperation(
|
|
valid=True,
|
|
message=no_migration_needed
|
|
)
|
|
|