yuqei
121 lines
3.7 KiB
Python
121 lines
3.7 KiB
Python
from pathlib import Path
|
|
import json
|
|
import os
|
|
import sys
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
if str(ROOT) not in sys.path:
|
|
sys.path.insert(0, str(ROOT))
|
|
|
|
from sqlalchemy import text
|
|
from sqlalchemy.exc import SQLAlchemyError
|
|
|
|
from src.db.engine import engine
|
|
|
|
MIGRATIONS_DIR = ROOT / "migrations"
|
|
|
|
|
|
def iter_migration_files() -> list[Path]:
|
|
return sorted(path for path in MIGRATIONS_DIR.glob("*.sql") if path.is_file())
|
|
|
|
|
|
def split_sql_statements(sql: str) -> list[str]:
|
|
statements = []
|
|
current = []
|
|
in_single_quote = False
|
|
in_double_quote = False
|
|
index = 0
|
|
while index < len(sql):
|
|
char = sql[index]
|
|
next_char = sql[index + 1] if index + 1 < len(sql) else ""
|
|
|
|
if not in_single_quote and not in_double_quote and char == "-" and next_char == "-":
|
|
while index < len(sql) and sql[index] != "\n":
|
|
index += 1
|
|
continue
|
|
if char == "'" and not in_double_quote:
|
|
in_single_quote = not in_single_quote
|
|
elif char == '"' and not in_single_quote:
|
|
in_double_quote = not in_double_quote
|
|
if char == ";" and not in_single_quote and not in_double_quote:
|
|
statement = "".join(current).strip()
|
|
if statement:
|
|
statements.append(statement)
|
|
current = []
|
|
index += 1
|
|
continue
|
|
current.append(char)
|
|
index += 1
|
|
|
|
statement = "".join(current).strip()
|
|
if statement:
|
|
statements.append(statement)
|
|
return statements
|
|
|
|
|
|
def apply_migration_file(connection, path: Path) -> int:
|
|
sql = path.read_text(encoding="utf-8").strip()
|
|
statements = split_sql_statements(sql)
|
|
for statement in statements:
|
|
connection.execute(text(statement))
|
|
return len(statements)
|
|
|
|
|
|
def run_migrations(connection, files: list[Path]) -> list[dict]:
|
|
results = []
|
|
for path in files:
|
|
statement_count = apply_migration_file(connection, path)
|
|
results.append({"file": path.name, "statements": statement_count})
|
|
return results
|
|
|
|
|
|
def main(env=None) -> int:
|
|
env = os.environ if env is None else env
|
|
json_output = _get_bool(env, "MIGRATIONS_OUTPUT_JSON", False)
|
|
files = iter_migration_files()
|
|
if not files:
|
|
summary = {"passed": True, "migration_count": 0, "results": []}
|
|
_print_summary(summary, json_output=json_output)
|
|
return 0
|
|
|
|
try:
|
|
with engine.begin() as connection:
|
|
results = run_migrations(connection=connection, files=files)
|
|
except SQLAlchemyError as exc:
|
|
summary = {
|
|
"passed": False,
|
|
"migration_count": len(files),
|
|
"results": [],
|
|
"error_type": exc.__class__.__name__,
|
|
"error": str(exc),
|
|
}
|
|
_print_summary(summary, json_output=json_output)
|
|
return 2
|
|
summary = {"passed": True, "migration_count": len(files), "results": results}
|
|
_print_summary(summary, json_output=json_output)
|
|
return 0
|
|
|
|
|
|
def _print_summary(summary: dict, json_output: bool) -> None:
|
|
if json_output:
|
|
print(json.dumps(summary, ensure_ascii=False, sort_keys=True))
|
|
return
|
|
if not summary["results"] and summary["passed"]:
|
|
print("no migration files found")
|
|
for item in summary["results"]:
|
|
if item["statements"]:
|
|
print(f"applied {item['file']}; statements: {item['statements']}")
|
|
if not summary["passed"]:
|
|
print(f"Database unavailable for migrations: {summary.get('error_type')}")
|
|
|
|
|
|
def _get_bool(env, key: str, default: bool) -> bool:
|
|
value = env.get(key)
|
|
if value is None or value == "":
|
|
return default
|
|
return str(value).strip().lower() in {"1", "true", "yes", "on"}
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|