opensource/WrenAI
3
0
Fork 0
mirror of https://github.com/Canner/WrenAI.git synced 2026-09-01 15:34:04 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(wren): load cubes from folder-per-entity layout (#2350)

Browse Source
This commit is contained in:
Peter
2026-06-09 15:21:35 +08:00
committed by GitHub
parent a2b58cedf3
commit 6f2542e117
2 changed files with 204 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files
core/wren/src/wren/context.py
+80 -4
View File
@@ -589,12 +589,19 @@ def _load_views_v2(project_path: Path) -> list[dict]:
def load_cubes(project_path: Path) -> list[dict]:
"""Load cubes from project_path/cubes/*.yml.
"""Load cubes — dispatches on schema_version.
Each file is one cube definition. Files use the same YAML shape on
every schema version (no v1/v2 dispatch — cubes were added after the
directory-per-entity migration).
v1 (legacy): cubes/*.yml
v2: cubes/<name>/metadata.yml
"""
sv = get_schema_version(project_path)
if sv == 1:
return _load_cubes_v1(project_path)
return _load_cubes_v2(project_path)
def _load_cubes_v1(project_path: Path) -> list[dict]:
"""Legacy: load cube YAML files from project_path/cubes/*.yml."""
cubes_dir = project_path / "cubes"
if not cubes_dir.is_dir():
return []
@@ -607,6 +614,34 @@ def load_cubes(project_path: Path) -> list[dict]:
return cubes
def _load_cubes_v2(project_path: Path) -> list[dict]:
"""v2: load cubes from project_path/cubes/<cube_name>/ directories.
Each cube directory must contain metadata.yml.
"""
cubes_dir = project_path / "cubes"
if not cubes_dir.is_dir():
return []
cubes = []
for d in sorted(cubes_dir.iterdir()):
if not d.is_dir():
continue
meta_file = d / "metadata.yml"
if not meta_file.exists():
continue
data = yaml.safe_load(meta_file.read_text())
if isinstance(data, dict):
data["_source_file"] = str(meta_file.relative_to(cubes_dir))
cubes.append(data)
return cubes
def _cube_migration_target(cube: dict, source_file: str | None) -> tuple[str, str]:
"""Return (cube_name, target_metadata_path) for a v1 cube migration."""
name = cube.get("name", Path(source_file).stem if source_file else "unknown")
return name, f"cubes/{name}/metadata.yml"
def load_relationships(project_path: Path) -> list[dict]:
"""Load relationships from project_path/relationships.yml."""
rel_file = project_path / "relationships.yml"
@@ -1169,6 +1204,23 @@ def _plan_v1_to_v2(project_path: Path) -> tuple[list[str], list[str]]:
if views_file.exists():
deleted.append("views.yml")
# Cubes: flat files → directories
seen_cube_targets: set[str] = set()
cubes = _load_cubes_v1(project_path)
for cube in cubes:
source_file = cube.pop("_source_file", None)
_, target = _cube_migration_target(cube, source_file)
if target in seen_cube_targets:
raise UpgradeError(
f"Cannot upgrade: multiple legacy cube files map to '{target}'"
)
seen_cube_targets.add(target)
created.append(target)
if source_file:
deleted.append(f"cubes/{source_file}")
return created, deleted
@@ -1241,6 +1293,30 @@ def _apply_v1_to_v2(project_path: Path) -> None:
if views_file.exists():
views_file.unlink()
# Write new cube directories
seen_cube_targets: set[str] = set()
cubes = _load_cubes_v1(project_path)
for cube in cubes:
source_file = cube.pop("_source_file", None)
name, target = _cube_migration_target(cube, source_file)
if target in seen_cube_targets:
raise UpgradeError(
f"Cannot upgrade: multiple legacy cube files map to '{target}'"
)
seen_cube_targets.add(target)
cube_dir = project_path / "cubes" / name
cube_dir.mkdir(parents=True, exist_ok=True)
(cube_dir / "metadata.yml").write_text(
yaml.dump(cube, default_flow_style=False, sort_keys=False)
)
if source_file:
old_file = project_path / "cubes" / source_file
if old_file.exists():
old_file.unlink()
# ── Semantic validation (view dry-plan + description completeness) ─────────
core/wren/tests/unit/test_context.py
+124 -36
View File
@@ -774,10 +774,10 @@ def test_validate_view_dialect_unknown(tmp_path):
# ── Cubes ───────────────────────────────────────────────────────────────────
def _make_v3_cube_project(tmp_path: Path) -> Path:
"""v3 project with an orders model, ready for cubes/*.yml files."""
def _make_v2_cube_project(tmp_path: Path) -> Path:
"""v2 project with an orders model, ready for cubes/*/metadata.yml files."""
(tmp_path / "wren_project.yml").write_text(
"schema_version: 3\nname: test\ndata_source: postgres\ncatalog: wren\nschema: public\n"
"schema_version: 2\nname: test\ndata_source: postgres\ncatalog: wren\nschema: public\n"
)
d = tmp_path / "models" / "orders"
d.mkdir(parents=True)
@@ -791,12 +791,22 @@ def _make_v3_cube_project(tmp_path: Path) -> Path:
return tmp_path
def _write_cube(tmp_path: Path, name: str, content: str) -> Path:
cube_dir = tmp_path / "cubes" / name
cube_dir.mkdir(parents=True)
cube_file = cube_dir / "metadata.yml"
cube_file.write_text(content)
return cube_file
def test_load_cubes_returns_empty_when_no_dir(tmp_path):
assert load_cubes(tmp_path) == []
def test_load_cubes_parses_yaml(tmp_path):
_make_v3_cube_project(tmp_path)
def test_load_cubes_v1_parses_flat_yaml(tmp_path):
(tmp_path / "wren_project.yml").write_text(
"schema_version: 1\nname: test\ndata_source: postgres\ncatalog: wren\nschema: public\n"
)
cubes_dir = tmp_path / "cubes"
cubes_dir.mkdir()
(cubes_dir / "order_metrics.yml").write_text(
@@ -804,6 +814,21 @@ def test_load_cubes_parses_yaml(tmp_path):
"base_object: orders\n"
"measures:\n"
" - name: revenue\n expression: SUM(o_totalprice)\n type: DOUBLE\n"
)
cubes = load_cubes(tmp_path)
assert len(cubes) == 1
assert cubes[0]["name"] == "order_metrics"
def test_load_cubes_v2_parses_metadata_yaml(tmp_path):
_make_v2_cube_project(tmp_path)
_write_cube(
tmp_path,
"order_metrics",
"name: order_metrics\n"
"base_object: orders\n"
"measures:\n"
" - name: revenue\n expression: SUM(o_totalprice)\n type: DOUBLE\n"
"dimensions:\n"
" - name: status\n expression: o_orderstatus\n type: VARCHAR\n"
)
@@ -814,8 +839,8 @@ def test_load_cubes_parses_yaml(tmp_path):
assert cubes[0]["measures"][0]["name"] == "revenue"
def test_build_manifest_includes_cubes(tmp_path):
_make_v3_cube_project(tmp_path)
def test_load_cubes_v2_ignores_flat_yaml(tmp_path):
_make_v2_cube_project(tmp_path)
cubes_dir = tmp_path / "cubes"
cubes_dir.mkdir()
(cubes_dir / "order_metrics.yml").write_text(
@@ -824,6 +849,37 @@ def test_build_manifest_includes_cubes(tmp_path):
"measures:\n"
" - name: revenue\n expression: SUM(o_totalprice)\n type: DOUBLE\n"
)
assert load_cubes(tmp_path) == []
def test_load_cubes_v3_uses_v2_layout(tmp_path):
_make_v2_cube_project(tmp_path)
(tmp_path / "wren_project.yml").write_text(
"schema_version: 3\nname: test\ndata_source: postgres\ncatalog: wren\nschema: public\n"
)
_write_cube(
tmp_path,
"order_metrics",
"name: order_metrics\n"
"base_object: orders\n"
"measures:\n"
" - name: revenue\n expression: SUM(o_totalprice)\n type: DOUBLE\n"
)
cubes = load_cubes(tmp_path)
assert len(cubes) == 1
assert cubes[0]["name"] == "order_metrics"
def test_build_manifest_includes_cubes(tmp_path):
_make_v2_cube_project(tmp_path)
_write_cube(
tmp_path,
"order_metrics",
"name: order_metrics\n"
"base_object: orders\n"
"measures:\n"
" - name: revenue\n expression: SUM(o_totalprice)\n type: DOUBLE\n"
)
manifest = build_manifest(tmp_path)
assert "cubes" in manifest
assert manifest["cubes"][0]["name"] == "order_metrics"
@@ -831,10 +887,10 @@ def test_build_manifest_includes_cubes(tmp_path):
def test_build_json_cube_camel_case(tmp_path):
_make_v3_cube_project(tmp_path)
cubes_dir = tmp_path / "cubes"
cubes_dir.mkdir()
(cubes_dir / "order_metrics.yml").write_text(
_make_v2_cube_project(tmp_path)
_write_cube(
tmp_path,
"order_metrics",
"name: order_metrics\n"
"base_object: orders\n"
"measures:\n"
@@ -849,10 +905,10 @@ def test_build_json_cube_camel_case(tmp_path):
def test_validate_cube_unknown_base_object(tmp_path):
_make_v3_cube_project(tmp_path)
cubes_dir = tmp_path / "cubes"
cubes_dir.mkdir()
(cubes_dir / "bad.yml").write_text(
_make_v2_cube_project(tmp_path)
_write_cube(
tmp_path,
"bad",
"name: bad\nbase_object: nosuch\nmeasures: [{name: c, expression: 'COUNT(*)', type: BIGINT}]\n"
)
errors = validate_project(tmp_path)
@@ -860,25 +916,23 @@ def test_validate_cube_unknown_base_object(tmp_path):
def test_validate_cube_duplicate_name(tmp_path):
_make_v3_cube_project(tmp_path)
cubes_dir = tmp_path / "cubes"
cubes_dir.mkdir()
_make_v2_cube_project(tmp_path)
body = (
"name: order_metrics\nbase_object: orders\n"
"measures: [{name: c, expression: 'COUNT(*)', type: BIGINT}]\n"
)
(cubes_dir / "a.yml").write_text(body)
(cubes_dir / "b.yml").write_text(body)
_write_cube(tmp_path, "a", body)
_write_cube(tmp_path, "b", body)
errors = validate_project(tmp_path)
assert any("duplicate cube name" in e.message for e in errors)
def test_validate_cube_missing_base_object_uses_snake_case(tmp_path):
"""Validation error should reference the YAML field name (snake_case)."""
_make_v3_cube_project(tmp_path)
cubes_dir = tmp_path / "cubes"
cubes_dir.mkdir()
(cubes_dir / "om.yml").write_text(
_make_v2_cube_project(tmp_path)
_write_cube(
tmp_path,
"om",
"name: order_metrics\n"
"measures: [{name: c, expression: 'COUNT(*)', type: BIGINT}]\n"
)
@@ -889,10 +943,10 @@ def test_validate_cube_missing_base_object_uses_snake_case(tmp_path):
def test_validate_cube_non_string_hierarchy_level(tmp_path):
"""Non-string hierarchy levels must be reported, not crash."""
_make_v3_cube_project(tmp_path)
cubes_dir = tmp_path / "cubes"
cubes_dir.mkdir()
(cubes_dir / "om.yml").write_text(
_make_v2_cube_project(tmp_path)
_write_cube(
tmp_path,
"om",
"name: order_metrics\n"
"base_object: orders\n"
"measures: [{name: c, expression: 'COUNT(*)', type: BIGINT}]\n"
@@ -907,10 +961,10 @@ def test_validate_cube_non_string_hierarchy_level(tmp_path):
def test_validate_cube_bad_hierarchy(tmp_path):
_make_v3_cube_project(tmp_path)
cubes_dir = tmp_path / "cubes"
cubes_dir.mkdir()
(cubes_dir / "om.yml").write_text(
_make_v2_cube_project(tmp_path)
_write_cube(
tmp_path,
"om",
"name: order_metrics\n"
"base_object: orders\n"
"measures: [{name: c, expression: 'COUNT(*)', type: BIGINT}]\n"
@@ -923,10 +977,10 @@ def test_validate_cube_bad_hierarchy(tmp_path):
def test_validate_cube_ok(tmp_path):
_make_v3_cube_project(tmp_path)
cubes_dir = tmp_path / "cubes"
cubes_dir.mkdir()
(cubes_dir / "om.yml").write_text(
_make_v2_cube_project(tmp_path)
_write_cube(
tmp_path,
"om",
"name: order_metrics\n"
"base_object: orders\n"
"measures: [{name: c, expression: 'COUNT(*)', type: BIGINT}]\n"
@@ -965,6 +1019,14 @@ def _make_v1_project(tmp_path: Path) -> Path:
" - name: monthly\n"
' statement: "SELECT\\n date_trunc(month, d)\\n FROM t"\n'
)
cubes_dir = tmp_path / "cubes"
cubes_dir.mkdir()
(cubes_dir / "order_metrics.yml").write_text(
"name: order_metrics\n"
"base_object: orders\n"
"measures:\n"
" - name: revenue\n expression: SUM(amount)\n type: DOUBLE\n"
)
(tmp_path / "relationships.yml").write_text("relationships: []\n")
(tmp_path / "instructions.md").write_text("## Rule 1\nAlways use UTC.\n")
return tmp_path
@@ -977,7 +1039,9 @@ def test_plan_upgrade_v1_to_v2(tmp_path):
assert result.to_version == 2
assert any("models/orders/metadata.yml" in f for f in result.files_created)
assert any("models/revenue/ref_sql.sql" in f for f in result.files_created)
assert any("cubes/order_metrics/metadata.yml" in f for f in result.files_created)
assert any("models/orders.yml" in f for f in result.files_deleted)
assert any("cubes/order_metrics.yml" in f for f in result.files_deleted)
assert any("views.yml" in f for f in result.files_deleted)
@@ -989,6 +1053,22 @@ def test_plan_upgrade_v1_to_v3(tmp_path):
assert len(result.files_created) > 0
def test_plan_upgrade_v1_to_v2_rejects_duplicate_cube_targets(tmp_path):
_make_v1_project(tmp_path)
(tmp_path / "cubes" / "other_metrics.yml").write_text(
"name: order_metrics\n"
"base_object: orders\n"
"measures:\n"
" - name: count\n expression: COUNT(*)\n type: BIGINT\n"
)
# Use fresh import to avoid stale class reference after importlib.reload in earlier tests.
from wren.context import UpgradeError as _UE # noqa: PLC0415
with pytest.raises(_UE, match="multiple legacy cube files"):
plan_upgrade(tmp_path, target_version=2)
def test_plan_upgrade_v2_to_v3(tmp_path):
_make_v2_project(tmp_path)
result = plan_upgrade(tmp_path, target_version=3)
@@ -1034,10 +1114,12 @@ def test_apply_upgrade_v1_to_v2(tmp_path):
assert (tmp_path / "models" / "revenue" / "metadata.yml").exists()
assert (tmp_path / "models" / "revenue" / "ref_sql.sql").exists()
assert (tmp_path / "views" / "summary" / "metadata.yml").exists()
assert (tmp_path / "cubes" / "order_metrics" / "metadata.yml").exists()
# Old files deleted
assert not (tmp_path / "models" / "orders.yml").exists()
assert not (tmp_path / "models" / "revenue.yml").exists()
assert not (tmp_path / "cubes" / "order_metrics.yml").exists()
assert not (tmp_path / "views.yml").exists()
# schema_version updated
@@ -1050,6 +1132,9 @@ def test_apply_upgrade_v1_to_v2(tmp_path):
assert names == {"orders", "revenue"}
revenue = next(m for m in models if m["name"] == "revenue")
assert "SELECT SUM(amount)" in revenue["ref_sql"]
cubes = load_cubes(tmp_path)
assert len(cubes) == 1
assert cubes[0]["name"] == "order_metrics"
def test_apply_upgrade_v2_to_v3(tmp_path):
@@ -1066,6 +1151,9 @@ def test_apply_upgrade_v1_to_v3(tmp_path):
assert get_schema_version(tmp_path) == 3
assert (tmp_path / "models" / "orders" / "metadata.yml").exists()
assert not (tmp_path / "models" / "orders.yml").exists()
assert (tmp_path / "cubes" / "order_metrics" / "metadata.yml").exists()
assert not (tmp_path / "cubes" / "order_metrics.yml").exists()
assert load_cubes(tmp_path)[0]["name"] == "order_metrics"
def test_upgrade_preserves_relationships(tmp_path):