fix(alembic): auto-discover models for fresh databases

This commit is contained in:
GuoQing Zhang
2026-07-17 15:39:25 +08:00
parent bf05265623
commit 0cbbe37974
3 changed files with 141 additions and 0 deletions
@@ -393,8 +393,10 @@ def run_migrations_online() -> None:
"""
from bisheng.core.database.manager import sync_get_database_connection
from bisheng.core.database.model_discovery import import_all_sqlmodel_models
database_conn_manager = sync_get_database_connection()
import_all_sqlmodel_models()
with database_conn_manager.engine.connect() as connection:
ensure_alembic_version_table(connection)
@@ -0,0 +1,63 @@
"""Convention-based SQLModel discovery for schema bootstrap tooling."""
import ast
import importlib
from pathlib import Path
_BISHENG_PACKAGE_ROOT = Path(__file__).resolve().parents[2]
_SQLMODEL_DIRECTORY_PATTERNS = (
"database/models",
"common/models",
"*/domain/models",
)
def _declares_sqlmodel_table(module_path: Path) -> bool:
"""Return whether a module declares a literal ``table=True`` class."""
tree = ast.parse(module_path.read_text(encoding="utf-8"), filename=str(module_path))
return any(
isinstance(node, ast.ClassDef)
and any(
keyword.arg == "table" and isinstance(keyword.value, ast.Constant) and keyword.value.value is True
for keyword in node.keywords
)
for node in ast.walk(tree)
)
def _module_name(module_path: Path) -> str:
relative_path = module_path.relative_to(_BISHENG_PACKAGE_ROOT.parent).with_suffix("")
module_parts = list(relative_path.parts)
if module_parts[-1] == "__init__":
module_parts.pop()
return ".".join(module_parts)
def discover_sqlmodel_module_names() -> tuple[str, ...]:
"""Find table modules under the repository's model-directory conventions."""
model_directories = {
model_directory
for pattern in _SQLMODEL_DIRECTORY_PATTERNS
for model_directory in _BISHENG_PACKAGE_ROOT.glob(pattern)
if model_directory.is_dir()
}
module_names = {
_module_name(module_path)
for model_directory in model_directories
for module_path in model_directory.rglob("*.py")
if _declares_sqlmodel_table(module_path)
}
return tuple(sorted(module_names))
def import_all_sqlmodel_models() -> None:
"""Strictly import every discovered model so metadata is complete."""
module_names = discover_sqlmodel_module_names()
if not module_names:
raise RuntimeError("No SQLModel table modules found under the model-directory conventions")
for module_name in module_names:
try:
importlib.import_module(module_name)
except Exception as exc:
raise RuntimeError(f"Failed to import SQLModel module {module_name}") from exc
@@ -0,0 +1,76 @@
"""Guards for convention-based SQLModel discovery used by DB bootstrap."""
import ast
import subprocess
import sys
from pathlib import Path
import pytest
from bisheng.core.database import model_discovery
from bisheng.core.database.model_discovery import discover_sqlmodel_module_names, import_all_sqlmodel_models
_BISHENG_ROOT = Path(__file__).resolve().parents[2] / "bisheng"
def _declares_sqlmodel_table(path: Path) -> bool:
tree = ast.parse(path.read_text(encoding="utf-8"))
return any(
isinstance(node, ast.ClassDef)
and any(
keyword.arg == "table" and isinstance(keyword.value, ast.Constant) and keyword.value.value is True
for keyword in node.keywords
)
for node in ast.walk(tree)
)
def _module_name(path: Path) -> str:
module_parts = list(path.relative_to(_BISHENG_ROOT.parent).with_suffix("").parts)
if module_parts[-1] == "__init__":
module_parts.pop()
return ".".join(module_parts)
def test_discovery_covers_every_literal_sqlmodel_table_module():
discovered_modules = {_module_name(path) for path in _BISHENG_ROOT.rglob("*.py") if _declares_sqlmodel_table(path)}
convention_modules = set(discover_sqlmodel_module_names())
missing_modules = discovered_modules - convention_modules
assert not missing_modules, (
f"SQLModel table modules outside the model-directory conventions: {sorted(missing_modules)}"
)
def test_strict_model_loading_registers_space_channel_member():
result = subprocess.run(
[
sys.executable,
"-c",
"from sqlmodel import SQLModel; "
"from bisheng.core.database.model_discovery import import_all_sqlmodel_models; "
"import_all_sqlmodel_models(); "
"assert 'space_channel_member' in SQLModel.metadata.tables",
],
check=False,
capture_output=True,
text=True,
)
assert result.returncode == 0, result.stderr
def test_strict_model_loading_raises_when_a_discovered_module_cannot_import(monkeypatch):
missing_module = "bisheng.common.models.does_not_exist"
monkeypatch.setattr(model_discovery, "discover_sqlmodel_module_names", lambda: (missing_module,))
with pytest.raises(RuntimeError, match=missing_module):
import_all_sqlmodel_models()
def test_strict_model_loading_rejects_empty_discovery(monkeypatch):
monkeypatch.setattr(model_discovery, "discover_sqlmodel_module_names", tuple)
with pytest.raises(RuntimeError, match="No SQLModel table modules"):
import_all_sqlmodel_models()