nex_basse/backend/check_migration.py

110 lines
4.3 KiB
Python

import sys
from pathlib import Path
from sqlalchemy import inspect, text
# Add the parent directory to sys.path
backend_dir = Path(__file__).resolve().parent
if str(backend_dir) not in sys.path:
sys.path.insert(0, str(backend_dir))
from app.core.db import engine
def check_and_migrate():
inspector = inspect(engine)
def get_columns(table_name: str):
try:
return [c['name'] for c in inspector.get_columns(table_name)]
except Exception:
return []
columns = get_columns('biz_meeting')
migrations = [
(
"asr_model_id",
"ALTER TABLE biz_meeting ADD COLUMN asr_model_id INTEGER REFERENCES biz_ai_model(model_id)",
[
"ALTER TABLE biz_meeting ADD COLUMN asr_model_id INT",
"ALTER TABLE biz_meeting ADD CONSTRAINT fk_meeting_asr_model FOREIGN KEY (asr_model_id) REFERENCES biz_ai_model(model_id)"
],
"ALTER TABLE biz_meeting ADD COLUMN asr_model_id INTEGER REFERENCES biz_ai_model(model_id)"
),
(
"summary_model_id",
"ALTER TABLE biz_meeting ADD COLUMN summary_model_id INTEGER REFERENCES biz_ai_model(model_id)",
[
"ALTER TABLE biz_meeting ADD COLUMN summary_model_id INT",
"ALTER TABLE biz_meeting ADD CONSTRAINT fk_meeting_summary_model FOREIGN KEY (summary_model_id) REFERENCES biz_ai_model(model_id)"
],
"ALTER TABLE biz_meeting ADD COLUMN summary_model_id INTEGER REFERENCES biz_ai_model(model_id)"
),
(
"summary_prompt_id",
"ALTER TABLE biz_meeting ADD COLUMN summary_prompt_id INTEGER REFERENCES biz_prompt_template(id)",
[
"ALTER TABLE biz_meeting ADD COLUMN summary_prompt_id INT",
"ALTER TABLE biz_meeting ADD CONSTRAINT fk_meeting_summary_prompt FOREIGN KEY (summary_prompt_id) REFERENCES biz_prompt_template(id)"
],
"ALTER TABLE biz_meeting ADD COLUMN summary_prompt_id INTEGER REFERENCES biz_prompt_template(id)"
),
]
for name, pg_sql, mysql_sqls, sqlite_sql in migrations:
if name in columns:
print(f"Column '{name}' already exists.")
continue
print(f"Column '{name}' missing in 'biz_meeting'. Adding it...")
with engine.connect() as conn:
if engine.dialect.name == 'postgresql':
conn.execute(text(pg_sql))
elif engine.dialect.name == 'mysql':
for sql in mysql_sqls:
conn.execute(text(sql))
else:
conn.execute(text(sqlite_sql))
conn.commit()
columns.append(name)
print("Column added successfully.")
hotword_columns = get_columns('biz_hotword')
hotword_migrations = [
(
"scope",
"ALTER TABLE biz_hotword ADD COLUMN scope VARCHAR(50) DEFAULT 'personal'",
[
"ALTER TABLE biz_hotword ADD COLUMN scope VARCHAR(50) DEFAULT 'personal'"
],
"ALTER TABLE biz_hotword ADD COLUMN scope VARCHAR(50) DEFAULT 'personal'"
),
(
"user_id",
"ALTER TABLE biz_hotword ADD COLUMN user_id INTEGER REFERENCES sys_user(user_id)",
[
"ALTER TABLE biz_hotword ADD COLUMN user_id INT",
"ALTER TABLE biz_hotword ADD CONSTRAINT fk_hotword_user FOREIGN KEY (user_id) REFERENCES sys_user(user_id)"
],
"ALTER TABLE biz_hotword ADD COLUMN user_id INTEGER REFERENCES sys_user(user_id)"
),
]
for name, pg_sql, mysql_sqls, sqlite_sql in hotword_migrations:
if name in hotword_columns:
print(f"Column '{name}' already exists.")
continue
print(f"Column '{name}' missing in 'biz_hotword'. Adding it...")
with engine.connect() as conn:
if engine.dialect.name == 'postgresql':
conn.execute(text(pg_sql))
elif engine.dialect.name == 'mysql':
for sql in mysql_sqls:
conn.execute(text(sql))
else:
conn.execute(text(sqlite_sql))
conn.commit()
hotword_columns.append(name)
print("Column added successfully.")
if __name__ == "__main__":
check_and_migrate()