wehub-resource-sync
83 lines
3.2 KiB
Python
83 lines
3.2 KiB
Python
import ast
|
|
from pathlib import Path
|
|
|
|
# Explicitly import all dbmodel modules so that every ORM model is registered
|
|
# with Base.registry before we inspect mappers.
|
|
import mlflow.store.model_registry.dbmodels.models # noqa: F401
|
|
import mlflow.store.tracking.dbmodels.models # noqa: F401
|
|
import mlflow.store.workspace.dbmodels.models # noqa: F401
|
|
from mlflow.store.db.base_sql_model import Base
|
|
|
|
# Locate the repository root relative to this test file so that workspace
|
|
# store paths resolve correctly regardless of the pytest working directory.
|
|
_REPO_ROOT = Path(__file__).resolve().parent.parent.parent
|
|
|
|
# Every store that provides a ``_get_query`` with workspace filtering.
|
|
WORKSPACE_STORE_PATHS = [
|
|
"mlflow/store/tracking/sqlalchemy_workspace_store.py",
|
|
"mlflow/store/model_registry/sqlalchemy_store.py",
|
|
"mlflow/store/jobs/sqlalchemy_workspace_store.py",
|
|
]
|
|
|
|
|
|
def _models_handled_by_get_query(ws_path: str) -> set[str]:
|
|
"""Parse a workspace store file and return the ``Sql*`` model names
|
|
referenced in comparisons inside its ``_get_query`` method.
|
|
"""
|
|
tree = ast.parse((_REPO_ROOT / ws_path).read_text())
|
|
|
|
# Collect module-level variables holding Sql* names (e.g. tuples of models)
|
|
var_models: dict[str, set[str]] = {}
|
|
for node in ast.iter_child_nodes(tree):
|
|
if isinstance(node, ast.Assign):
|
|
for target in node.targets:
|
|
if isinstance(target, ast.Name):
|
|
names = {
|
|
n.id
|
|
for n in ast.walk(node.value)
|
|
if isinstance(n, ast.Name) and n.id.startswith("Sql")
|
|
}
|
|
if names:
|
|
var_models[target.id] = names
|
|
|
|
# Extract model names from _get_query comparisons
|
|
models: set[str] = set()
|
|
for node in ast.walk(tree):
|
|
if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
|
|
continue
|
|
if node.name != "_get_query":
|
|
continue
|
|
for inner in ast.walk(node):
|
|
if not isinstance(inner, ast.Compare):
|
|
continue
|
|
for comp in inner.comparators:
|
|
models |= {
|
|
n.id
|
|
for n in ast.walk(comp)
|
|
if isinstance(n, ast.Name) and n.id.startswith("Sql")
|
|
}
|
|
if isinstance(comp, ast.Name) and comp.id in var_models:
|
|
models |= var_models[comp.id]
|
|
return models
|
|
|
|
|
|
def test_all_workspace_models_handled_in_get_query():
|
|
"""Every model with a workspace column must appear in at least one
|
|
workspace store's ``_get_query``.
|
|
"""
|
|
handled: set[str] = set()
|
|
for ws_path in WORKSPACE_STORE_PATHS:
|
|
handled |= _models_handled_by_get_query(ws_path)
|
|
|
|
models_with_column: set[str] = set()
|
|
for mapper in Base.registry.mappers:
|
|
if "workspace" in {col.key for col in mapper.columns}:
|
|
models_with_column.add(mapper.class_.__name__)
|
|
|
|
missing = models_with_column - handled
|
|
assert not missing, (
|
|
f"These models have a `workspace` column but are not handled by any "
|
|
f"workspace store's _get_query: {sorted(missing)}. "
|
|
f"Add handling in the appropriate workspace store's _get_query method."
|
|
)
|