pz-cli 5.0.0__py3-none-any.whl

pz_foundation/db/aerich_executor.py

@@ -0,0 +1,664 @@
+from __future__ import annotations
+
+import asyncio
+import importlib.util
+import json
+import os
+from pathlib import Path
+import re
+import shutil
+import subprocess
+import sys
+from typing import Any
+
+from pz_foundation.db.config import build_foundation_tortoise_config
+
+
+def _require_aerich() -> None:
+    if importlib.util.find_spec("aerich") is None:
+        raise RuntimeError("Aerich is not installed. Install the project dependencies that provide the migration executor.")
+
+
+def _runnable_sources(migration_sources: tuple[Any, ...], *, action: str) -> tuple[Any, ...]:
+    runnable: list[Any] = []
+    for source in migration_sources:
+        if _is_shared_source_package(source):
+            continue
+        if not str(getattr(source, "models_module", "") or "").strip():
+            continue
+        files = tuple(str(item) for item in (getattr(source, "files", ()) or ()))
+        has_real_migration = any(
+            Path(item).suffix == ".py" and Path(item).name != "__init__.py"
+            for item in files
+        )
+        if has_real_migration or action == "makemigrations":
+            runnable.append(source)
+    return tuple(runnable)
+
+
+def _shared_init_sources(migration_sources: tuple[Any, ...]) -> tuple[Any, ...]:
+    return tuple(
+        source
+        for source in migration_sources
+        if _is_shared_source_package(source)
+    )
+
+
+def _is_shared_source_package(source: Any) -> bool:
+    raw_package_root = str(getattr(source, "package_root", "") or "").strip()
+    if not raw_package_root:
+        return False
+    try:
+        package_root = Path(raw_package_root).resolve()
+    except Exception:
+        return False
+    return ".sources" in package_root.parts
+
+
+def _ensure_workspace_packages(*, workspace_dir: Path, migration_sources: tuple[Any, ...]) -> Path:
+    owners_dir = workspace_dir / "owners"
+    owners_dir.mkdir(parents=True, exist_ok=True)
+    (owners_dir / "__init__.py").touch(exist_ok=True)
+    for source in migration_sources:
+        owner_dir = owners_dir / str(getattr(source, "owner"))
+        owner_dir.mkdir(parents=True, exist_ok=True)
+        (owner_dir / "__init__.py").touch(exist_ok=True)
+    return owners_dir
+
+
+def _build_tortoise_config(*, plan: Any, migration_sources: tuple[Any, ...]) -> dict[str, Any]:
+    installed_apps = [
+        str(getattr(source, "module_path", "")).split(":", 1)[0].strip()
+        for source in migration_sources
+        if str(getattr(source, "module_path", "")).strip()
+    ]
+    return build_foundation_tortoise_config(
+        project_dir=str(plan.project.project_dir),
+        installed_apps=installed_apps,
+    )
+
+
+def _ensure_aerich_history_table(*, tortoise_config: dict[str, Any]) -> None:
+    if "aerich" not in (tortoise_config.get("apps", {}) or {}):
+        return
+    try:
+        from tortoise import Tortoise
+    except Exception:
+        return
+
+    async def _bootstrap() -> None:
+        try:
+            await Tortoise.init(
+                config={
+                    "connections": tortoise_config.get("connections", {}),
+                    "apps": {"aerich": tortoise_config["apps"]["aerich"]},
+                }
+            )
+            await Tortoise.generate_schemas(safe=True)
+        finally:
+            try:
+                await Tortoise.close_connections()
+            except Exception:
+                pass
+
+    asyncio.run(_bootstrap())
+
+
+def _migration_python_files(owner_dir: Path) -> tuple[Path, ...]:
+    if not owner_dir.is_dir():
+        return ()
+    return tuple(
+        path
+        for path in sorted(owner_dir.glob("*.py"), key=_migration_sort_key)
+        if path.name != "__init__.py"
+    )
+
+
+def _migration_sort_key(path: Path) -> tuple[int, str]:
+    match = re.match(r"^(?P<prefix>\d+)_", path.name)
+    if match is None:
+        return (sys.maxsize, path.name)
+    return (int(match.group("prefix")), path.name)
+
+
+def _migration_file_has_aerich_state(path: Path) -> bool:
+    try:
+        text = path.read_text(encoding="utf-8")
+    except Exception:
+        return False
+    return "MODELS_STATE" in text
+
+
+def _owner_uses_legacy_sql_migrations(owner_dir: Path) -> bool:
+    return bool(_legacy_sql_migration_files(owner_dir))
+
+
+def _legacy_sql_migration_files(owner_dir: Path) -> tuple[Path, ...]:
+    files = _migration_python_files(owner_dir)
+    if any(not _migration_file_has_aerich_state(path) for path in files):
+        return files
+    return tuple(path for path in files if not _migration_file_has_aerich_state(path))
+
+
+def _coerce_query_row(row: Any) -> Any:
+    if isinstance(row, dict):
+        return row
+    keys = getattr(row, "keys", None)
+    if not callable(keys):
+        return row
+    try:
+        return {key: row[key] for key in keys()}
+    except Exception:
+        return row
+
+
+def _coerce_execute_query_result(result: Any) -> Any:
+    if not isinstance(result, tuple) or len(result) < 2:
+        return result
+    rows = result[1]
+    if not isinstance(rows, list):
+        return result
+    coerced = [_coerce_query_row(row) for row in rows]
+    return (result[0], coerced, *result[2:])
+
+
+class _LegacySqlMigrationConnection:
+    def __init__(self, connection: Any) -> None:
+        self._connection = connection
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self._connection, name)
+
+    async def execute_query(self, *args: Any, **kwargs: Any) -> Any:
+        result = await self._connection.execute_query(*args, **kwargs)
+        return _coerce_execute_query_result(result)
+
+
+def _connection_dialect(connection: Any) -> str:
+    try:
+        dialect = str(getattr(getattr(connection, "capabilities", None), "dialect", "") or "").lower()
+        if dialect:
+            return dialect
+    except Exception:
+        pass
+    try:
+        module_name = str(type(connection).__module__ or "").lower()
+    except Exception:
+        module_name = ""
+    if "sqlite" in module_name:
+        return "sqlite"
+    if "asyncpg" in module_name or "postgres" in module_name or "psycopg" in module_name:
+        return "postgres"
+    if "mysql" in module_name or "aiomysql" in module_name or "mariadb" in module_name:
+        return "mysql"
+    return ""
+
+
+def _split_sql_statements(script: str) -> list[str]:
+    statements: list[str] = []
+    current: list[str] = []
+    in_single_quote = False
+    in_double_quote = False
+    idx = 0
+    while idx < len(script):
+        char = script[idx]
+        current.append(char)
+        if char == "'" and not in_double_quote:
+            next_char = script[idx + 1] if idx + 1 < len(script) else ""
+            if in_single_quote and next_char == "'":
+                current.append(next_char)
+                idx += 1
+            else:
+                in_single_quote = not in_single_quote
+        elif char == '"' and not in_single_quote:
+            in_double_quote = not in_double_quote
+        elif char == ";" and not in_single_quote and not in_double_quote:
+            statement = "".join(current).strip()
+            if statement:
+                statements.append(statement)
+            current = []
+        idx += 1
+    tail = "".join(current).strip()
+    if tail:
+        statements.append(tail)
+    return statements
+
+
+def _strip_sqlite_unsupported_comment_statements(script: str) -> str:
+    statements = [
+        statement
+        for statement in _split_sql_statements(script)
+        if not statement.lstrip().upper().startswith("COMMENT ON ")
+    ]
+    if not statements:
+        return ""
+    return "\n".join(statements)
+
+
+def _strip_sqlite_postgres_do_blocks(script: str) -> str:
+    return re.sub(
+        r"\bDO\s+\$\$.*?\$\$\s*;?",
+        "",
+        script,
+        flags=re.IGNORECASE | re.DOTALL,
+    )
+
+
+def _strip_sqlite_postgres_json_casts(script: str) -> str:
+    result: list[str] = []
+    in_single_quote = False
+    in_double_quote = False
+    idx = 0
+    while idx < len(script):
+        char = script[idx]
+        if char == "'" and not in_double_quote:
+            result.append(char)
+            next_char = script[idx + 1] if idx + 1 < len(script) else ""
+            if in_single_quote and next_char == "'":
+                result.append(next_char)
+                idx += 2
+                continue
+            in_single_quote = not in_single_quote
+            idx += 1
+            continue
+        if char == '"' and not in_single_quote:
+            in_double_quote = not in_double_quote
+            result.append(char)
+            idx += 1
+            continue
+        if not in_single_quote and not in_double_quote:
+            lower_tail = script[idx:].lower()
+            if lower_tail.startswith("::jsonb"):
+                idx += len("::jsonb")
+                continue
+            if lower_tail.startswith("::json"):
+                idx += len("::json")
+                continue
+        result.append(char)
+        idx += 1
+    return "".join(result)
+
+
+def _replace_sqlite_postgres_default_functions(script: str) -> str:
+    return re.sub(
+        r"\bDEFAULT\s+NOW\s*\(\s*\)",
+        "DEFAULT CURRENT_TIMESTAMP",
+        script,
+        flags=re.IGNORECASE,
+    )
+
+
+def _prepare_legacy_sql_script(script: str, *, connection: Any) -> str:
+    if _connection_dialect(connection) != "sqlite":
+        return script
+    script = _strip_sqlite_postgres_do_blocks(script)
+    script = _strip_sqlite_unsupported_comment_statements(script)
+    script = _strip_sqlite_postgres_json_casts(script)
+    return _replace_sqlite_postgres_default_functions(script)
+
+
+def _query_result_rows(result: Any) -> list[Any]:
+    if not isinstance(result, tuple) or len(result) < 2:
+        return []
+    rows = result[1]
+    if not isinstance(rows, list):
+        return []
+    return rows
+
+
+def _row_value(row: Any, key: str) -> Any:
+    get = getattr(row, "get", None)
+    if callable(get):
+        return get(key)
+    try:
+        return row[key]
+    except Exception:
+        return None
+
+
+async def _sqlite_column_exists(connection: Any, *, table_name: str, column_name: str) -> bool:
+    result = await connection.execute_query(f'PRAGMA table_info("{table_name}")')
+    return any(
+        str(_row_value(_coerce_query_row(row), "name") or "") == column_name
+        for row in _query_result_rows(result)
+    )
+
+
+_SQLITE_ADD_COLUMN_IF_NOT_EXISTS_RE = re.compile(
+    r'^\s*ALTER\s+TABLE\s+"?(?P<table>[^"\s]+)"?\s+ADD\s+COLUMN\s+IF\s+NOT\s+EXISTS\s+"?(?P<column>[^"\s]+)"?\s+',
+    flags=re.IGNORECASE | re.DOTALL,
+)
+_SQLITE_ALTER_COLUMN_RE = re.compile(
+    r'^\s*ALTER\s+TABLE\s+"?[^"\s]+"?\s+ALTER\s+COLUMN\s+',
+    flags=re.IGNORECASE | re.DOTALL,
+)
+_SQLITE_ALTER_CONSTRAINT_RE = re.compile(
+    r'^\s*ALTER\s+TABLE\s+"?[^"\s]+"?\s+(?:DROP|ADD)\s+CONSTRAINT\s+',
+    flags=re.IGNORECASE | re.DOTALL,
+)
+_SQLITE_UPDATE_TABLE_RE = re.compile(
+    r'^\s*UPDATE\s+"?(?P<table>[^"\s]+)"?\s+',
+    flags=re.IGNORECASE | re.DOTALL,
+)
+
+
+async def _sqlite_table_has_rows(connection: Any, *, table_name: str) -> bool:
+    result = await connection.execute_query(f'SELECT 1 FROM "{table_name}" LIMIT 1')
+    return bool(_query_result_rows(result))
+
+
+async def _execute_sqlite_legacy_statement(connection: Any, statement: str) -> None:
+    normalized = statement.strip().rstrip(";").strip()
+    if not normalized:
+        return
+    add_column_match = _SQLITE_ADD_COLUMN_IF_NOT_EXISTS_RE.match(normalized)
+    if add_column_match is not None:
+        table_name = add_column_match.group("table")
+        column_name = add_column_match.group("column")
+        if await _sqlite_column_exists(connection, table_name=table_name, column_name=column_name):
+            return
+        normalized = re.sub(
+            r"\bADD\s+COLUMN\s+IF\s+NOT\s+EXISTS\b",
+            "ADD COLUMN",
+            normalized,
+            count=1,
+            flags=re.IGNORECASE,
+        )
+    if _SQLITE_ALTER_COLUMN_RE.match(normalized) or _SQLITE_ALTER_CONSTRAINT_RE.match(normalized):
+        return
+    update_match = _SQLITE_UPDATE_TABLE_RE.match(normalized)
+    if update_match is not None and not await _sqlite_table_has_rows(connection, table_name=update_match.group("table")):
+        return
+    await connection.execute_script(normalized)
+
+
+async def _execute_legacy_sql_script(connection: Any, script: str) -> None:
+    if _connection_dialect(connection) != "sqlite":
+        await connection.execute_script(script)
+        return
+    for statement in _split_sql_statements(script):
+        await _execute_sqlite_legacy_statement(connection, statement)
+
+
+def _apply_legacy_sql_migrations(
+    *,
+    owner: str,
+    owner_dir: Path,
+    tortoise_config: dict[str, Any],
+) -> dict[str, Any]:
+    files = _legacy_sql_migration_files(owner_dir)
+    if not files:
+        return {"owner": owner, "command": "legacy-upgrade", "files": [], "returncode": 0}
+    try:
+        from tortoise import Tortoise
+    except Exception as exc:
+        raise RuntimeError(f"Tortoise is required to apply legacy SQL migrations for app '{owner}': {exc}") from exc
+
+    async def _upgrade() -> None:
+        try:
+            await Tortoise.init(config=tortoise_config)
+            connection = Tortoise.get_connection("default")
+            migration_connection = _LegacySqlMigrationConnection(connection)
+            for path in files:
+                spec = importlib.util.spec_from_file_location(f"_pz_legacy_migration_{owner}_{path.stem}", path)
+                if spec is None or spec.loader is None:
+                    raise RuntimeError(f"Cannot load legacy migration file: {path}")
+                module = importlib.util.module_from_spec(spec)
+                spec.loader.exec_module(module)
+                upgrade = getattr(module, "upgrade", None)
+                if upgrade is None:
+                    continue
+                sql = await upgrade(migration_connection)
+                sql_text = str(sql or "").strip()
+                sql_text = _prepare_legacy_sql_script(sql_text, connection=connection).strip()
+                if sql_text:
+                    try:
+                        await _execute_legacy_sql_script(connection, sql_text)
+                    except Exception as exc:
+                        raise RuntimeError(
+                            f"Legacy SQL migration failed for app '{owner}' file '{path.name}': {exc}"
+                        ) from exc
+        finally:
+            try:
+                await Tortoise.close_connections()
+            except Exception:
+                pass
+
+    asyncio.run(_upgrade())
+    return {
+        "owner": owner,
+        "command": "legacy-upgrade",
+        "files": [path.name for path in files],
+        "returncode": 0,
+    }
+
+
+def _init_shared_source_schemas(
+    *,
+    plan: Any,
+    migration_sources: tuple[Any, ...],
+    owners_dir: Path | None = None,
+) -> tuple[str, ...]:
+    sources = _shared_init_sources(migration_sources)
+    if not sources:
+        return ()
+    installed_apps = [
+        str(getattr(source, "module_path", "")).split(":", 1)[0].strip()
+        for source in sources
+        if str(getattr(source, "module_path", "")).strip()
+    ]
+    if not installed_apps:
+        return ()
+    tortoise_config = build_foundation_tortoise_config(
+        project_dir=str(plan.project.project_dir),
+        installed_apps=installed_apps,
+    )
+    try:
+        from tortoise import Tortoise
+    except Exception as exc:
+        raise RuntimeError(f"Tortoise is required to initialize shared source database schemas: {exc}") from exc
+
+    async def _bootstrap() -> None:
+        try:
+            await Tortoise.init(config=tortoise_config)
+            await Tortoise.generate_schemas(safe=True)
+        finally:
+            try:
+                await Tortoise.close_connections()
+            except Exception:
+                pass
+
+    asyncio.run(_bootstrap())
+    if owners_dir is not None:
+        for source in sources:
+            owner = str(getattr(source, "owner"))
+            owner_dir = owners_dir / owner
+            if _owner_uses_legacy_sql_migrations(owner_dir):
+                _apply_legacy_sql_migrations(
+                    owner=owner,
+                    owner_dir=owner_dir,
+                    tortoise_config=tortoise_config,
+                )
+    return tuple(str(getattr(source, "owner")) for source in sources)
+
+
+def _write_workspace_config(*, workspace_dir: Path, tortoise_config: dict[str, Any]) -> tuple[Path, Path]:
+    module_path = workspace_dir / "_pz_migration_config.py"
+    module_path.write_text(
+        "\n".join(
+            [
+                "from __future__ import annotations",
+                "",
+                f"TORTOISE_ORM = {json.dumps(tortoise_config, ensure_ascii=True, indent=2, sort_keys=True)}",
+                "TORTOISE_CONFIG = TORTOISE_ORM",
+                "",
+            ]
+        ),
+        encoding="utf-8",
+    )
+    config_path = workspace_dir / "aerich.ini"
+    config_path.write_text(
+        "\n".join(
+            [
+                "[tool.aerich]",
+                'tortoise_orm = "_pz_migration_config.TORTOISE_CONFIG"',
+                'location = "./owners"',
+                f'src_folder = "{workspace_dir}"',
+                "",
+            ]
+        ),
+        encoding="utf-8",
+    )
+    return module_path, config_path
+
+
+def _aerich_command() -> list[str]:
+    if importlib.util.find_spec("aerich.__main__") is not None:
+        return [sys.executable, "-m", "aerich"]
+    candidate = Path(sys.executable).with_name("aerich")
+    if candidate.exists():
+        return [str(candidate)]
+    found = shutil.which("aerich")
+    if found:
+        return [found]
+    return [sys.executable, "-m", "aerich"]
+
+
+def _runtime_import_roots(*, plan: Any) -> tuple[Path, ...]:
+    settings = getattr(plan, "settings", None)
+    values = tuple(getattr(settings, "import_roots", ()) or ())
+    roots: list[Path] = []
+    seen: set[Path] = set()
+    for item in values:
+        try:
+            root = Path(item).resolve()
+        except Exception:
+            continue
+        if root in seen:
+            continue
+        seen.add(root)
+        roots.append(root)
+    return tuple(roots)
+
+
+def _run_aerich(
+    *,
+    workspace_dir: Path,
+    owners_dir: Path,
+    config_path: Path,
+    owner: str,
+    command: str,
+    tortoise_config: dict[str, Any],
+    plan: Any,
+) -> dict[str, Any]:
+    env = dict(os.environ)
+    pythonpath = env.get("PYTHONPATH", "")
+    import_entries = [str(workspace_dir), str(Path(getattr(getattr(plan, "project", None), "project_dir", workspace_dir)).resolve())]
+    import_entries.extend(str(path) for path in _runtime_import_roots(plan=plan))
+    if pythonpath:
+        import_entries.extend(item for item in pythonpath.split(os.pathsep) if item)
+    deduped_entries: list[str] = []
+    seen_entries: set[str] = set()
+    for entry in import_entries:
+        value = str(entry or "").strip()
+        if not value or value in seen_entries:
+            continue
+        seen_entries.add(value)
+        deduped_entries.append(value)
+    env["PYTHONPATH"] = os.pathsep.join(deduped_entries)
+    base_cmd = _aerich_command() + ["-c", str(config_path), "--app", owner]
+    owner_dir = owners_dir / owner
+    if command == "upgrade" and _owner_uses_legacy_sql_migrations(owner_dir):
+        return _apply_legacy_sql_migrations(owner=owner, owner_dir=owner_dir, tortoise_config=tortoise_config)
+
+    def _invoke(*args: str) -> subprocess.CompletedProcess[str]:
+        return subprocess.run(
+            base_cmd + list(args),
+            cwd=str(workspace_dir),
+            env=env,
+            capture_output=True,
+            text=True,
+            check=False,
+        )
+
+    proc = _invoke(command)
+    stderr_text = proc.stderr or ""
+    stdout_text = proc.stdout or ""
+    if proc.returncode != 0 and "Old format of migration file detected" in (stderr_text or stdout_text):
+        fix_proc = _invoke("fix-migrations")
+        if fix_proc.returncode == 0:
+            proc = _invoke(command)
+            stderr_text = proc.stderr or ""
+            stdout_text = proc.stdout or ""
+    if proc.returncode != 0:
+        raise RuntimeError(
+            f"Aerich {command} failed for app '{owner}' with exit code {proc.returncode}: "
+            f"{(stderr_text or stdout_text).strip()}"
+        )
+    return {
+        "owner": owner,
+        "command": command,
+        "stdout": stdout_text,
+        "stderr": stderr_text,
+        "returncode": proc.returncode,
+    }
+
+
+def execute(
+    *,
+    action: str,
+    plan: Any,
+    workspace_dir: Path,
+    manifest_path: Path,
+    orm_config_path: Path,
+    lock_path: Path,
+    migration_sources: tuple[Any, ...],
+) -> dict[str, Any]:
+    del manifest_path, orm_config_path, lock_path
+    _require_aerich()
+    runnable = _runnable_sources(migration_sources, action=action)
+    owners_dir = _ensure_workspace_packages(workspace_dir=workspace_dir, migration_sources=migration_sources)
+    tortoise_config = _build_tortoise_config(plan=plan, migration_sources=migration_sources)
+    _ensure_aerich_history_table(tortoise_config=tortoise_config)
+    initialized_shared_sources = ()
+    if action in {"migrate", "upgrade"}:
+        initialized_shared_sources = _init_shared_source_schemas(
+            plan=plan,
+            migration_sources=migration_sources,
+            owners_dir=owners_dir,
+        )
+    _, config_path = _write_workspace_config(
+        workspace_dir=workspace_dir,
+        tortoise_config=tortoise_config,
+    )
+    command_by_action = {
+        "makemigrations": "migrate",
+        "migrate": "upgrade",
+        "upgrade": "upgrade",
+    }
+    command = command_by_action.get(action)
+    if command is None:
+        raise RuntimeError(f"Unsupported Aerich executor action: {action}")
+    command_results = [
+        _run_aerich(
+            workspace_dir=workspace_dir,
+            owners_dir=owners_dir,
+            config_path=config_path,
+            owner=str(getattr(source, "owner")),
+            command=command,
+            tortoise_config=tortoise_config,
+            plan=plan,
+        )
+        for source in runnable
+    ]
+    return {
+        "action": action,
+        "owners": [str(getattr(source, "owner")) for source in runnable],
+        "initialized_shared_source_owners": list(initialized_shared_sources),
+        "results": command_results,
+    }
+
+
+__all__ = [
+    "execute",
+]