[BUG] Discover validation registry via a generated index module

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.

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]],

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"
         )

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.

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,28 +35,33 @@ 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:
             continue
         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

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)}"
+    )