diff --git a/changes/1691.feature.md b/changes/1691.feature.md new file mode 100644 index 0000000000..eec0042c52 --- /dev/null +++ b/changes/1691.feature.md @@ -0,0 +1 @@ +Add the --output option to the function that outputs gql and openapi and unify the output format. diff --git a/src/ai/backend/manager/cli/api.py b/src/ai/backend/manager/cli/api.py index 1ce81df995..8d013b6a07 100644 --- a/src/ai/backend/manager/cli/api.py +++ b/src/ai/backend/manager/cli/api.py @@ -1,8 +1,11 @@ from __future__ import annotations +import asyncio import logging +from pathlib import Path from typing import TYPE_CHECKING +import aiofiles import click import graphene @@ -19,9 +22,24 @@ def cli(args) -> None: pass +async def generate_gql_schema(output_path: Path) -> None: + schema = graphene.Schema(query=Queries, mutation=Mutations, auto_camelcase=False) + if output_path == "-": + log.info("======== GraphQL API Schema ========") + print(str(schema)) + else: + async with aiofiles.open(output_path, "w") as fw: + await fw.write(str(schema)) + + @cli.command() @click.pass_obj -def dump_gql_schema(cli_ctx: CLIContext) -> None: - schema = graphene.Schema(query=Queries, mutation=Mutations, auto_camelcase=False) - log.info("======== GraphQL API Schema ========") - print(str(schema)) +@click.option( + "--output", + "-o", + default="-", + type=click.Path(dir_okay=False, writable=True), + help="Output file path (default: stdout)", +) +def dump_gql_schema(cli_ctx: CLIContext, output: Path) -> None: + asyncio.run(generate_gql_schema(output)) diff --git a/src/ai/backend/manager/openapi.py b/src/ai/backend/manager/openapi.py index bcb9f66bba..5d1e10c993 100644 --- a/src/ai/backend/manager/openapi.py +++ b/src/ai/backend/manager/openapi.py @@ -292,18 +292,26 @@ async def generate_openapi(output_path: Path) -> None: route_def["parameters"] = parameters route_def["description"] = "\n".join(description) openapi["paths"][path][method.lower()] = route_def - - async with aiofiles.open(output_path, mode="w") as fw: - await fw.write(json.dumps(openapi, ensure_ascii=False, indent=2)) + if output_path == "-" or output_path is None: + print(json.dumps(openapi, ensure_ascii=False, indent=2)) + else: + async with aiofiles.open(output_path, mode="w") as fw: + await fw.write(json.dumps(openapi, ensure_ascii=False, indent=2)) @click.command() -@click.argument("OUTPUT_PATH") -def main(output_path: Path) -> None: +@click.option( + "--output", + "-o", + default="-", + type=click.Path(dir_okay=False, writable=True), + help="Output file path (default: stdout)", +) +def main(output: Path) -> None: """ Generates OpenAPI specification of Backend.AI API. """ - asyncio.run(generate_openapi(output_path)) + asyncio.run(generate_openapi(output)) if __name__ == "__main__":