Consolidate scripts to generate openapi specs

.pre-commit-config.yaml

@@ -1514,14 +1514,14 @@ repos:
       - id: generate-openapi-spec-fab
         name: Generate the FastAPI API spec for FAB
         language: python
-        entry: ./scripts/ci/prek/generate_openapi_spec_fab.py
+        entry: ./scripts/ci/prek/generate_openapi_spec_providers.py fab
         pass_filenames: false
         files: ^providers/fab/src/airflow/providers/fab/auth_manager/api_fastapi/.*\.py$
         additional_dependencies: ['rich>=12.4.4', 'openapi-spec-validator>=0.7.1']
       - id: generate-openapi-spec-keycloak
         name: Generate the FastAPI API spec for Keycloak
         language: python
-        entry: ./scripts/ci/prek/generate_openapi_spec_keycloak.py
+        entry: ./scripts/ci/prek/generate_openapi_spec_providers.py keycloak
         pass_filenames: false
         files: ^providers/keycloak/src/airflow/providers/keycloak/auth_manager/.*\.py$
         additional_dependencies: [ 'rich>=12.4.4', 'openapi-spec-validator>=0.7.1' ]

scripts/ci/prek/generate_openapi_spec_fab.py → scripts/ci/prek/generate_openapi_spec_providers.py

@@ -30,7 +30,7 @@
 initialize_breeze_prek(__name__, __file__)
 
 cmd_result = run_command_via_breeze_shell(
-    ["python3", "/opt/airflow/scripts/in_container/run_generate_openapi_spec_fab.py"],
+    ["python3", "/opt/airflow/scripts/in_container/run_generate_openapi_spec_providers.py", sys.argv[1]],
     backend="postgres",
     skip_environment_initialization=False,
 )

scripts/in_container/run_generate_openapi_spec_providers.py

@@ -0,0 +1,92 @@
+#!/usr/bin/env python
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+from typing import TYPE_CHECKING, NamedTuple
+
+from in_container_utils import console, generate_openapi_file, validate_openapi_file
+
+from airflow.providers.fab.auth_manager.api_fastapi import __file__ as FAB_AUTHMGR_API_PATH
+from airflow.providers.fab.auth_manager.fab_auth_manager import FabAuthManager
+from airflow.providers.keycloak.auth_manager import __file__ as KEYCLOAK_AUTHMGR_PATH
+from airflow.providers.keycloak.auth_manager.keycloak_auth_manager import KeycloakAuthManager
+from airflow.providers_manager import ProvidersManager
+
+if TYPE_CHECKING:
+    from fastapi import FastAPI
+
+
+class ProviderDef(NamedTuple):
+    openapi_spec_file: Path
+    app: FastAPI | None
+    prefix: str
+
+
+sys.path.insert(0, str(Path(__file__).parent.resolve()))
+ProvidersManager().initialize_providers_configuration()
+
+PROVIDERS_DEFS = {
+    "fab": ProviderDef(
+        openapi_spec_file=Path(FAB_AUTHMGR_API_PATH).parent
+        / "openapi"
+        / "v2-fab-auth-manager-generated.yaml",
+        app=FabAuthManager().get_fastapi_app(),
+        prefix="/auth",
+    ),
+    "keycloak": ProviderDef(
+        openapi_spec_file=Path(KEYCLOAK_AUTHMGR_PATH).parent
+        / "openapi"
+        / "v2-keycloak-auth-manager-generated.yaml",
+        app=KeycloakAuthManager().get_fastapi_app(),
+        prefix="/auth",
+    ),
+}
+
+
+# Generate FAB auth manager openapi spec
+def generate_openapi_specs(provider_name: str):
+    provider_def = PROVIDERS_DEFS.get(provider_name)
+    if provider_def is None:
+        console.print(f"[red]Provider '{provider_name}' not found. Skipping OpenAPI spec generation.[/]")
+        sys.exit(1)
+    app = provider_def.app
+    openapi_spec_file = provider_def.openapi_spec_file
+
+    if app:
+        generate_openapi_file(app=app, file_path=openapi_spec_file, prefix=provider_def.prefix)
+        validate_openapi_file(openapi_spec_file)
+    else:
+        console.print(
+            f"[red]Provider '{provider_name}' has no FastAPI app. Skipping OpenAPI spec generation.[/]"
+        )
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Generate openapi-spec for the specified provider.")
+    parser.add_argument(
+        "provider",
+        type=str,
+        help="The name of the provider whose openapi-spec should be compiled.",
+        choices=list(PROVIDERS_DEFS.keys()),
+    )
+    args = parser.parse_args()
+    generate_openapi_specs(args.provider)