diff --git a/packages/overture-schema-codegen/changelog.d/661.misc.md b/packages/overture-schema-codegen/changelog.d/661.misc.md new file mode 100644 index 000000000..f373987a8 --- /dev/null +++ b/packages/overture-schema-codegen/changelog.d/661.misc.md @@ -0,0 +1 @@ +The PySpark generator now emits an `_index` module listing the generated validation modules, which the runtime registry imports to discover them without walking the generated tree on disk. diff --git a/packages/overture-schema-codegen/src/overture/schema/codegen/pyspark/pipeline.py b/packages/overture-schema-codegen/src/overture/schema/codegen/pyspark/pipeline.py index 3b83f5448..65fbaea2e 100644 --- a/packages/overture-schema-codegen/src/overture/schema/codegen/pyspark/pipeline.py +++ b/packages/overture-schema-codegen/src/overture/schema/codegen/pyspark/pipeline.py @@ -60,6 +60,16 @@ class PipelineOutput: _OUTPUT_PACKAGE = "overture.schema.pyspark.expressions.generated" +_INDEX_PATH = PurePosixPath("_index.py") + +_INDEX_DOCSTRING = '''"""Index of the generated validation modules. + +The runtime registry imports this module and reads ``MODULES``. Because that is +an ordinary import, discovery works even when the package is loaded from a wheel +on ``sys.path`` (zipimport), as AWS Glue does via ``--extra-py-files``, where the +generated modules are reachable by import but not as files on disk. +"""''' + def _require_entry_point(spec: ModelSpec) -> str: """Return *spec*'s entry point or raise if it's missing.""" @@ -130,17 +140,53 @@ def generate_pyspark_modules( Returns ------- PipelineOutput - Source-tree model modules and test-tree modules. The generated - tree is PEP 420, so no `__init__.py` files are emitted. + Source-tree model modules plus an `_index` module listing them, + and test-tree modules. The index is an ordinary `_index.py`, so the + generated tree stays PEP 420 and its namespace packages are untouched. """ items = [(spec, build_checks(spec)) for spec in model_specs] source = [_render_module(spec, checks) for spec, checks in items] + if source: + source.append(_render_index(model_specs)) test: list[GeneratedModule] = [] for spec, checks in items: test.extend(_render_test_modules(spec, checks)) return PipelineOutput(source=source, test=test) +def _render_index(model_specs: Sequence[ModelSpec]) -> GeneratedModule: + """Render the `_index` module listing every generated validation module. + + The runtime registry imports the generated modules through this index, so + the codegen owns which modules exist. Each is aliased by its full dotted + path so two feature types with the same leaf name cannot collide. + """ + entries = sorted( + ( + ".".join([_OUTPUT_PACKAGE, *directory.parts]), + model_name, + "_".join([*directory.parts, model_name]), + ) + for directory, model_name in ( + _directory_and_model_name(spec) for spec in model_specs + ) + ) + lines = [ + "# This file is auto-generated by overture-schema-codegen. Do not edit.", + _INDEX_DOCSTRING, + "", + "from __future__ import annotations", + "", + *(f"from {parent} import {leaf} as {alias}" for parent, leaf, alias in entries), + "", + "MODULES = (", + *(f" {alias}," for _, _, alias in entries), + ")", + "", + ] + return GeneratedModule(content="\n".join(lines), path=_INDEX_PATH) + + def _render_module( spec: ModelSpec, checks: tuple[list[Check], list[ModelCheck]], diff --git a/packages/overture-schema-codegen/tests/test_pyspark_pipeline.py b/packages/overture-schema-codegen/tests/test_pyspark_pipeline.py index 1c04c94a4..772afb7db 100644 --- a/packages/overture-schema-codegen/tests/test_pyspark_pipeline.py +++ b/packages/overture-schema-codegen/tests/test_pyspark_pipeline.py @@ -94,7 +94,15 @@ def test_empty_specs_returns_no_modules(self) -> None: assert result.test == [] def test_one_module_per_spec(self, two_spec_modules: PipelineOutput) -> None: - assert len(two_spec_modules.source) == 2 + model_modules = [ + m for m in two_spec_modules.source if m.path.name != "_index.py" + ] + assert len(model_modules) == 2 + + def test_emits_index_module(self, two_spec_modules: PipelineOutput) -> None: + index = [m for m in two_spec_modules.source if m.path.name == "_index.py"] + assert len(index) == 1 + assert index[0].path == PurePosixPath("_index.py") def test_paths_unique_per_tree(self, two_spec_modules: PipelineOutput) -> None: # source and test trees mirror the same dirs; uniqueness is @@ -236,7 +244,7 @@ def test_module_path_mirrors_entry_point(self) -> None: SimpleModel, entry_point="overture.schema.simple:SimpleModel" ) modules = generate_pyspark_modules([spec]) - features = modules.source + features = [m for m in modules.source if m.path.name != "_index.py"] assert len(features) == 1 assert features[0].path == PurePosixPath( "overture/schema/simple/simple_model.py" @@ -264,8 +272,8 @@ def test_neither_tree_has_init_modules(self) -> None: class TestNoRegistryEmitted: def test_registry_module_is_no_longer_generated(self) -> None: - # The runtime builds the registry via entry-point discovery; codegen - # must not emit `_registry.py`. + # The runtime builds the registry from the generated `_index` module; + # codegen emits that index but never a hand-written `_registry.py`. spec = extract_model( SimpleModel, entry_point="overture.schema.simple:SimpleModel" ) diff --git a/packages/overture-schema-pyspark/changelog.d/661.bugfix.md b/packages/overture-schema-pyspark/changelog.d/661.bugfix.md new file mode 100644 index 000000000..07a043ed6 --- /dev/null +++ b/packages/overture-schema-pyspark/changelog.d/661.bugfix.md @@ -0,0 +1 @@ +Fixed the validation registry coming back empty when `overture-schema-pyspark` is loaded from a wheel on `sys.path` (zipimport) rather than installed to a real directory, as happens on AWS Glue via `--extra-py-files`. Codegen now emits an `_index` module listing the generated validation modules, and the registry imports it instead of walking the generated tree, so discovery works whether the package lives on disk or inside a zip archive. diff --git a/packages/overture-schema-pyspark/src/overture/schema/pyspark/_registry.py b/packages/overture-schema-pyspark/src/overture/schema/pyspark/_registry.py index d5c20ad5d..948277059 100644 --- a/packages/overture-schema-pyspark/src/overture/schema/pyspark/_registry.py +++ b/packages/overture-schema-pyspark/src/overture/schema/pyspark/_registry.py @@ -1,50 +1,31 @@ """Runtime registry of feature validations. -Built at import time by walking the generated `expressions.generated` -namespace and collecting every module that exposes the -codegen-emitted `ENTRY_POINT` and `MODEL_VALIDATION` constants. - -The generated tree on disk is the runtime source of truth: the -registry contains exactly what was generated, regardless of which -theme packages are installed alongside the pyspark package. A missing -`expressions/generated/` subtree simply yields an empty registry -- -the package still imports cleanly. +Built at import time from the generated `_index` module, which imports every +generated validation module and exposes them as `MODULES`. That import is +ordinary, so the registry populates whether the package is installed to a real +directory or loaded straight from a wheel on `sys.path` (as on Glue, via +`--extra-py-files`); discovery never walks the generated tree as files. + +A build without the generated tree has no `_index` module, so the registry is +empty and the package still imports cleanly. A module that is present but fails +to import (a missing dependency, a codegen bug) raises, so a real breakage is +loud and a validation is never dropped without notice. """ from __future__ import annotations import importlib import logging -from pathlib import Path from .check import ModelValidation logger = logging.getLogger(__name__) -_GENERATED_ROOT = "overture.schema.pyspark.expressions.generated" - - -def _iter_generated_module_names(root_paths: list[str]) -> list[str]: - """Return the dotted names of every generated module on disk. - - The generated tree is PEP 420 (no `__init__.py`), so its subdirectories - are namespace packages. `pkgutil.walk_packages` skips those, so the tree - is walked as files instead: every `.py` under the namespace roots, keyed - to a dotted name relative to `_GENERATED_ROOT`. - """ - names: list[str] = [] - for root_path in root_paths: - base = Path(root_path) - for path in sorted(base.rglob("*.py")): - if path.name == "__init__.py": - continue - relative = path.relative_to(base).with_suffix("") - names.append(".".join([_GENERATED_ROOT, *relative.parts])) - return names +_INDEX_MODULE = "overture.schema.pyspark.expressions.generated._index" def _walk() -> tuple[dict[str, ModelValidation], dict[str, dict[str, str]]]: - """Walk the generated tree and collect registry + partition map. + """Collect registry + partition map from the generated index. Returns a `(registry, partition_map)` pair: @@ -54,20 +35,25 @@ def _walk() -> tuple[dict[str, ModelValidation], dict[str, dict[str, str]]]: "place"}`) for path construction. Features with no `PARTITIONS` data (empty dict) are omitted; the codegen only sets `PARTITIONS` when the data lake organizes the feature by Hive partitions. - `type` is appended here from the module file name so consumers - get a complete partition path without the codegen having to - duplicate the type value. + `type` comes from the module name so consumers get a complete + partition path without the codegen having to duplicate the value. """ registry: dict[str, ModelValidation] = {} partition_map: dict[str, dict[str, str]] = {} try: - root = importlib.import_module(_GENERATED_ROOT) - except ImportError: - return registry, partition_map - - for name in _iter_generated_module_names(list(root.__path__)): - module = importlib.import_module(name) + index = importlib.import_module(_INDEX_MODULE) + except ModuleNotFoundError as e: + missing = e.name or "" + # A missing name equal to (or an ancestor of) the index module means + # the generated tree was never built, which is a legitimately empty + # registry. Any other missing name is a real dependency failure while + # importing a module the index references, so let it propagate. + if missing == _INDEX_MODULE or _INDEX_MODULE.startswith(f"{missing}."): + return registry, partition_map + raise + + for module in index.MODULES: entry_point = getattr(module, "ENTRY_POINT", None) validation = getattr(module, "MODEL_VALIDATION", None) if entry_point is None or validation is None: @@ -75,7 +61,7 @@ def _walk() -> tuple[dict[str, ModelValidation], dict[str, dict[str, str]]]: registry[entry_point] = validation partitions = getattr(module, "PARTITIONS", None) or {} if partitions: - feature_type = name.rsplit(".", 1)[-1] + feature_type = module.__name__.rsplit(".", 1)[-1] partition_map[entry_point] = {**partitions, "type": feature_type} return registry, partition_map diff --git a/packages/overture-schema-pyspark/tests/test_registry.py b/packages/overture-schema-pyspark/tests/test_registry.py index fcf261207..af35732d9 100644 --- a/packages/overture-schema-pyspark/tests/test_registry.py +++ b/packages/overture-schema-pyspark/tests/test_registry.py @@ -1,10 +1,9 @@ """Tests for the runtime registry's discovery of the generated tree. -The generated expression tree is PEP 420 (no `__init__.py`), so the -registry must walk it as a namespace package. No other test exercises the -real on-disk walk -- conformance tests import expression modules directly -and `test_validate.py` registers models through a test shim -- so an empty -registry would otherwise pass the suite unnoticed. +The registry is built from the generated `_index` module, which codegen emits +alongside the validation modules. One test asserts the registry populates; the +other asserts the index lists exactly the modules on disk, so a codegen bug +that drops a module from the index fails here and never reaches runtime. """ from __future__ import annotations @@ -17,32 +16,56 @@ from overture.schema.pyspark._registry import REGISTRY _GENERATED_ROOT = "overture.schema.pyspark.expressions.generated" +_INDEX_MODULE = f"{_GENERATED_ROOT}._index" -def _generated_leaf_count() -> int: - """Count generated model modules on disk (excludes namespace dirs). +def _generated_module_names() -> set[str]: + """Dotted names of every generated model module on disk. - Returns 0 when the generated tree is absent -- mirroring the registry's - own `ImportError` handling -- so the test skips rather than errors. + Empty when the generated tree is absent, matching the registry's own + handling, so the tests skip cleanly. The `_index` module is excluded + because it is discovery machinery; only validation modules are counted. """ try: root = importlib.import_module(_GENERATED_ROOT) except ImportError: - return 0 - return sum( - 1 - for base in root.__path__ - for path in Path(base).rglob("*.py") - if path.name != "__init__.py" - ) + return set() + names: set[str] = set() + for base in root.__path__: + for path in Path(base).rglob("*.py"): + if path.name in ("__init__.py", "_index.py"): + continue + relative = path.relative_to(base).with_suffix("") + names.add(".".join([_GENERATED_ROOT, *relative.parts])) + return names def test_registry_discovers_generated_models() -> None: - """The registry finds generated modules under the PEP 420 namespace tree.""" - if _generated_leaf_count() == 0: + """The registry finds generated modules through the index.""" + if not _generated_module_names(): pytest.skip("generated tree not present; run `make generate-pyspark`") generated_entries = [ key for key in REGISTRY if ":" in key and key.startswith("overture.schema.") ] - assert generated_entries, "registry found no generated feature modules on disk" + assert generated_entries, "registry found no generated feature modules" + + +def test_index_lists_every_generated_module() -> None: + """The generated index covers exactly the modules on disk. + + Guards the codegen step that emits `_index`: a module generated but left + out of the index (or an index entry with no module) would otherwise drop + that feature type from the registry silently. + """ + on_disk = _generated_module_names() + if not on_disk: + pytest.skip("generated tree not present; run `make generate-pyspark`") + + index = importlib.import_module(_INDEX_MODULE) + indexed = {module.__name__ for module in index.MODULES} + assert indexed == on_disk, ( + "generated _index is out of sync with the modules on disk.\n" + f" indexed but not on disk: {sorted(indexed - on_disk)}\n" + f" on disk but not indexed: {sorted(on_disk - indexed)}" + )