scripts/in_container/run_generate_openapi_spec_providers.py

#!/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_manager import ProvidersManager

if TYPE_CHECKING:
    from fastapi import FastAPI


class ProviderDef(NamedTuple):
    openapi_spec_file: Path
    app: FastAPI | None
    prefix: str


ProvidersManager().initialize_providers_configuration()


def get_providers_defs(provider_name: str) -> ProviderDef:
    if provider_name == "fab":
        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

        return ProviderDef(
            openapi_spec_file=Path(FAB_AUTHMGR_API_PATH).parent
            / "openapi"
            / "v2-fab-auth-manager-generated.yaml",
            app=FabAuthManager().get_fastapi_app(),
            prefix="/auth",
        )
    if provider_name == "edge":
        from airflow.providers.edge3 import __file__ as EDGE_PATH
        from airflow.providers.edge3.worker_api.app import create_edge_worker_api_app

        # Ensure dist exists on aclean git...
        (Path(EDGE_PATH).parent / "plugins" / "www" / "dist").mkdir(parents=True, exist_ok=True)

        return ProviderDef(
            openapi_spec_file=Path(EDGE_PATH).parent / "worker_api" / "v2-edge-generated.yaml",
            app=create_edge_worker_api_app(),
            prefix="/edge_worker",
        )
    if provider_name == "keycloak":
        from airflow.providers.keycloak.auth_manager import __file__ as KEYCLOAK_AUTHMGR_PATH
        from airflow.providers.keycloak.auth_manager.keycloak_auth_manager import KeycloakAuthManager

        return ProviderDef(
            openapi_spec_file=Path(KEYCLOAK_AUTHMGR_PATH).parent
            / "openapi"
            / "v2-keycloak-auth-manager-generated.yaml",
            app=KeycloakAuthManager().get_fastapi_app(),
            prefix="/auth",
        )
    raise NotImplementedError(f"Provider '{provider_name}' is not supported.")


# Generate FAB auth manager openapi spec
def generate_openapi_specs(provider_name: str):
    provider_def = get_providers_defs(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=["fab", "edge", "keycloak"],
    )
    args = parser.parse_args()
    generate_openapi_specs(args.provider)