diff --git a/client/python/apache_polaris/cli/polaris_cli.py b/client/python/apache_polaris/cli/polaris_cli.py index 8bd49d63dd..b06286df04 100644 --- a/client/python/apache_polaris/cli/polaris_cli.py +++ b/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 diff --git a/client/python/generate_clients.py b/client/python/generate_clients.py index 9aaa07fd87..0c7c1b5af0 100644 --- a/client/python/generate_clients.py +++ b/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( diff --git a/client/python/integration_tests/conftest.py b/client/python/integration_tests/conftest.py index d448fb50e1..12f72cbd60 100644 --- a/client/python/integration_tests/conftest.py +++ b/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)