45 lines
1.2 KiB
Python
45 lines
1.2 KiB
Python
import re
|
|
from pathlib import Path
|
|
|
|
from alembic import op
|
|
|
|
|
|
_DOLLAR_TAG_RE = re.compile(r"\$\w+\$")
|
|
_SQL_DIRECTORY = Path(__file__).parent / "versions" / "sql"
|
|
|
|
|
|
def _split_statements(sql: str) -> list[str]:
|
|
statements: list[str] = []
|
|
current: list[str] = []
|
|
dollar_tag: str | None = None
|
|
|
|
for line in sql.splitlines():
|
|
stripped = line.strip()
|
|
if stripped.startswith("--"):
|
|
continue
|
|
|
|
if dollar_tag is None:
|
|
match = _DOLLAR_TAG_RE.search(stripped)
|
|
if match:
|
|
dollar_tag = match.group(0)
|
|
elif stripped.startswith(dollar_tag):
|
|
dollar_tag = None
|
|
|
|
current.append(line)
|
|
if dollar_tag is None and stripped.endswith(";"):
|
|
statements.append("\n".join(current))
|
|
current = []
|
|
|
|
if remaining := "\n".join(current).strip():
|
|
statements.append(remaining)
|
|
|
|
return statements
|
|
|
|
|
|
def run_sql_file_migration(filename: str) -> None:
|
|
"""Execute statements from a SQL file in alembic/versions/sql."""
|
|
sql = (_SQL_DIRECTORY / filename).read_text(encoding="utf-8")
|
|
for statement in _split_statements(sql):
|
|
if stripped := statement.rstrip(";").strip():
|
|
op.execute(stripped)
|