Fix monkey patching

client/python/apache_polaris/cli/polaris_cli.py

@@ -46,50 +46,8 @@ class PolarisCli:
     # Can be enabled if the client is able to authenticate directly without first fetching a token
     DIRECT_AUTHENTICATION_ENABLED = False
 
-    @staticmethod
-    def _patch_generated_models() -> None:
-        """
-        The OpenAPI generator creates an `api_client` that dynamically looks up
-        model classes from the `apache_polaris.sdk.catalog.models` module using `getattr()`.
-        For example, when a response for a `create_policy` call is received, the
-        deserializer tries to find the `LoadPolicyResponse` class by looking for
-        `apache_polaris.sdk.catalog.models.LoadPolicyResponse`.
-
-        However, the generator fails to add the necessary `import` statements
-        to the `apache_polaris/sdk/catalog/models/__init__.py` file. This means that even
-        though the model files exist (e.g., `load_policy_response.py`), the classes
-        are not part of the `apache_polaris.sdk.catalog.models` namespace.
-
-        This method works around the bug in the generated code without modifying
-        the source files. It runs once per CLI execution, before any commands, and
-        manually injects the missing response-side model classes into the
-        `apache_polaris.sdk.catalog.models` namespace, allowing the deserializer to find them.
-        """
-        import apache_polaris.sdk.catalog.models
-        from apache_polaris.sdk.catalog.models.applicable_policy import ApplicablePolicy
-        from apache_polaris.sdk.catalog.models.get_applicable_policies_response import GetApplicablePoliciesResponse
-        from apache_polaris.sdk.catalog.models.list_policies_response import ListPoliciesResponse
-        from apache_polaris.sdk.catalog.models.load_policy_response import LoadPolicyResponse
-        from apache_polaris.sdk.catalog.models.policy import Policy
-        from apache_polaris.sdk.catalog.models.policy_attachment_target import PolicyAttachmentTarget
-        from apache_polaris.sdk.catalog.models.policy_identifier import PolicyIdentifier
-
-        models_to_patch = {
-            "ApplicablePolicy": ApplicablePolicy,
-            "GetApplicablePoliciesResponse": GetApplicablePoliciesResponse,
-            "ListPoliciesResponse": ListPoliciesResponse,
-            "LoadPolicyResponse": LoadPolicyResponse,
-            "Policy": Policy,
-            "PolicyAttachmentTarget": PolicyAttachmentTarget,
-            "PolicyIdentifier": PolicyIdentifier,
-        }
-
-        for name, model_class in models_to_patch.items():
-            setattr(apache_polaris.sdk.catalog.models, name, model_class)
-
     @staticmethod
     def execute(args=None):
-        PolarisCli._patch_generated_models()
         options = Parser.parse(args)
         if options.command == Commands.PROFILES:
             from apache_polaris.cli.command import Command

client/python/generate_clients.py

@@ -32,6 +32,7 @@
 import logging
 import argparse
 import shutil
+import ast
 
 # Paths
 CLIENT_DIR = Path(__file__).parent
@@ -306,9 +307,53 @@ def build() -> None:
     generate_polaris_management_client()
     generate_polaris_catalog_client()
     generate_iceberg_catalog_client()
+    fix_catalog_models_init()
     prepend_licenses()
 
 
+def fix_catalog_models_init() -> None:
+    """
+    Regenerate the `apache_polaris.sdk.catalog.models.__init__.py` file by consolidating
+    imports for all model classes found under `apache_polaris/sdk/catalog/models`.
+
+    This ensures that rerunning the OpenAPI Generator (which overwrites `__init__.py`)
+    does not cause missing imports for earlier generated model files.
+    """
+    logger.info("Fixing catalog models __init__.py...")
+    models_dir = CLIENT_DIR / "apache_polaris" / "sdk" / "catalog" / "models"
+    init_py = models_dir / "__init__.py"
+
+    # Get all python files in the models directory except __init__.py
+    model_files = [
+        f for f in models_dir.glob("*.py") if f.is_file() and f.name != "__init__.py"
+    ]
+
+    # Generate import statements
+    imports = []
+    for model_file in sorted(model_files):
+        module_name = model_file.stem
+        with open(model_file, "r") as f:
+            tree = ast.parse(f.read())
+        class_name = None
+        for node in ast.walk(tree):
+            if isinstance(node, ast.ClassDef):
+                # Find the first class that doesn't start with an underscore
+                if not node.name.startswith("_"):
+                    class_name = node.name
+                    break
+        if class_name:
+            imports.append(
+                f"from apache_polaris.sdk.catalog.models.{module_name} import {class_name}"
+            )
+        else:
+            logger.warning(f"Could not find a suitable class in {model_file}")
+
+    # Write the new __init__.py
+    with open(init_py, "w") as f:
+        f.write("\n".join(sorted(imports)))
+    logger.info("Catalog models __init__.py fixed.")
+
+
 def main():
     parser = argparse.ArgumentParser(description="Generate Polaris Python clients.")
     parser.add_argument(

client/python/integration_tests/conftest.py

@@ -375,53 +375,3 @@ def clear_namespace(
 
 def format_namespace(namespace: List[str]) -> str:
     return codecs.decode("1F", "hex").decode("UTF-8").join(namespace)
-
-
-@pytest.fixture(scope="session", autouse=True)
-def _patch_generated_models() -> None:
-    """
-    The OpenAPI generator creates an `api_client` that dynamically looks up
-    model classes from the `apache_polaris.sdk.catalog.models` module using `getattr()`.
-    For example, when a response for a `create_policy` call is received, the
-    deserializer tries to find the `LoadPolicyResponse` class by looking for
-    `apache_polaris.sdk.catalog.models.LoadPolicyResponse`.
-
-    However, the generator fails to add the necessary `import` statements
-    to the `apache_polaris/sdk/catalog/models/__init__.py` file. This means that even
-    though the model files exist (e.g., `load_policy_response.py`), the classes
-    are not part of the `apache_polaris.sdk.catalog.models` namespace.
-
-    This fixture works around the bug in the generated code without modifying
-    the source files. It runs once per test session, before any tests, and
-    manually injects the missing response-side model classes into the
-    `apache_polaris.sdk.catalog.models` namespace, allowing the deserializer to find them.
-    """
-    import apache_polaris.sdk.catalog.models
-    from apache_polaris.sdk.catalog.models.applicable_policy import ApplicablePolicy
-    from apache_polaris.sdk.catalog.models.get_applicable_policies_response import (
-        GetApplicablePoliciesResponse,
-    )
-    from apache_polaris.sdk.catalog.models.list_policies_response import (
-        ListPoliciesResponse,
-    )
-    from apache_polaris.sdk.catalog.models.load_policy_response import (
-        LoadPolicyResponse,
-    )
-    from apache_polaris.sdk.catalog.models.policy import Policy
-    from apache_polaris.sdk.catalog.models.policy_attachment_target import (
-        PolicyAttachmentTarget,
-    )
-    from apache_polaris.sdk.catalog.models.policy_identifier import PolicyIdentifier
-
-    models_to_patch = {
-        "ApplicablePolicy": ApplicablePolicy,
-        "GetApplicablePoliciesResponse": GetApplicablePoliciesResponse,
-        "ListPoliciesResponse": ListPoliciesResponse,
-        "LoadPolicyResponse": LoadPolicyResponse,
-        "Policy": Policy,
-        "PolicyAttachmentTarget": PolicyAttachmentTarget,
-        "PolicyIdentifier": PolicyIdentifier,
-    }
-
-    for name, model_class in models_to_patch.items():
-        setattr(apache_polaris.sdk.catalog.models, name, model_class)