add infrahubctl graphql generate-return-types command

Author: dgarros

infrahub_sdk/ctl/cli_commands.py

@@ -26,6 +26,7 @@
 from ..ctl.client import initialize_client, initialize_client_sync
 from ..ctl.exceptions import QueryNotFoundError
 from ..ctl.generator import run as run_generator
+from ..ctl.graphql import app as graphql_app
 from ..ctl.menu import app as menu_app
 from ..ctl.object import app as object_app
 from ..ctl.render import list_jinja2_transforms, print_template_errors
@@ -63,6 +64,7 @@
 app.add_typer(repository_app, name="repository")
 app.add_typer(menu_app, name="menu")
 app.add_typer(object_app, name="object")
+app.add_typer(graphql_app, name="graphql")
 
 app.command(name="dump")(dump)
 app.command(name="load")(load)

infrahub_sdk/ctl/graphql.py

@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+from collections import defaultdict
+from pathlib import Path
+
+import typer
+from ariadne_codegen.client_generators.package import get_package_generator
+from ariadne_codegen.plugins.explorer import get_plugins_types
+from ariadne_codegen.plugins.manager import PluginManager
+from ariadne_codegen.schema import (
+    filter_operations_definitions,
+    get_graphql_queries,
+    get_graphql_schema_from_path,
+)
+from ariadne_codegen.settings import ClientSettings, CommentsStrategy
+from rich.console import Console
+
+from ..async_typer import AsyncTyper
+from ..ctl.utils import catch_exception
+from .parameters import CONFIG_PARAM
+
+app = AsyncTyper()
+console = Console()
+
+ARIADNE_PLUGINS = ["infrahub_sdk.graphql.plugin.InfrahubPlugin"]
+
+
+def find_gql_files(query_path: Path) -> list[Path]:
+    """
+    Find all files with .gql extension in the specified directory.
+
+    Args:
+        query_path: Path to the directory to search for .gql files
+
+    Returns:
+        List of Path objects for all .gql files found
+    """
+    if not query_path.exists():
+        raise FileNotFoundError(f"Directory not found: {query_path}")
+
+    if not query_path.is_dir() and query_path.is_file():
+        yield query_path
+
+    else:
+        yield from query_path.glob("*/*.gql")
+
+
+@app.callback()
+def callback() -> None:
+    """
+    Various GraphQL related commands.
+    """
+
+
+@app.command()
+@catch_exception(console=console)
+async def generate_return_types(
+    query: Path = typer.Argument(Path.cwd(), help="Location of the GraphQL query file(s)."),
+    schema: Path = typer.Option("schema.graphql", help="Path to the GraphQL schema file."),
+    _: str = CONFIG_PARAM,
+) -> None:
+    """Create Pydantic Models for GraphQL query return types"""
+
+    # Load the GraphQL schema
+    if not schema.exists():
+        raise FileNotFoundError(f"GraphQL Schema file not found: {schema}")
+    graphql_schema = get_graphql_schema_from_path(schema_path=str(schema))
+
+    # Initialize the plugin manager
+    plugin_manager = PluginManager(
+        schema=graphql_schema,
+        plugins_types=get_plugins_types(plugins_strs=ARIADNE_PLUGINS),
+    )
+
+    # Find the GraphQL files and organize them by directory
+    gql_files = find_gql_files(query)
+    gql_per_directory: dict[Path, list[Path]] = defaultdict(list)
+    for gql_file in gql_files:
+        gql_per_directory[gql_file.parent].append(gql_file)
+
+    # Generate the Pydantic Models for the GraphQL queries
+    for directory in gql_per_directory.keys():
+        package_generator = get_package_generator(
+            schema=graphql_schema,
+            fragments=[],
+            settings=ClientSettings(
+                schema_path=str(schema),
+                target_package_name=directory.name,
+                queries_path=str(directory),
+                include_comments=CommentsStrategy.NONE,
+            ),
+            plugin_manager=plugin_manager,
+        )
+
+        definitions = get_graphql_queries(queries_path=str(directory), schema=graphql_schema)
+        queries = filter_operations_definitions(definitions)
+
+        for query_operation in queries:
+            package_generator.add_operation(query_operation)
+
+        package_generator._generate_result_types()
+
+        for file_name in package_generator._result_types_files.keys():
+            print(f"Generated {file_name} in {directory}")

infrahub_sdk/graphql/__init__.py

@@ -0,0 +1,14 @@
+from .constants import VARIABLE_TYPE_MAPPING
+from .query import Mutation, Query
+from .renderers import render_input_block, render_query_block, render_variables_to_string
+from .return_type import GraphQLReturnTypeModel
+
+__all__ = [
+    "VARIABLE_TYPE_MAPPING",
+    "GraphQLReturnTypeModel",
+    "Mutation",
+    "Query",
+    "render_input_block",
+    "render_query_block",
+    "render_variables_to_string",
+]

infrahub_sdk/graphql/constants.py

@@ -0,0 +1 @@
+VARIABLE_TYPE_MAPPING = ((str, "String!"), (int, "Int!"), (float, "Float!"), (bool, "Boolean!"))

infrahub_sdk/graphql/plugin.py

@@ -0,0 +1,58 @@
+from __future__ import annotations
+
+import ast
+from typing import TYPE_CHECKING
+
+from ariadne_codegen.plugins.base import Plugin
+
+from .return_type import GraphQLReturnTypeModel
+
+if TYPE_CHECKING:
+    from graphql import ExecutableDefinitionNode
+
+
+class InfrahubPlugin(Plugin):
+    @staticmethod
+    def find_base_model_index(module: ast.Module) -> int:
+        for idx, item in enumerate(module.body):
+            if isinstance(item, ast.ImportFrom) and item.module == "base_model":
+                return idx
+        return -1
+
+    @classmethod
+    def replace_base_model_import(cls, module: ast.Module) -> ast.Module:
+        base_model_index = cls.find_base_model_index(module)
+        if base_model_index == -1:
+            raise ValueError("BaseModel not found in module")
+        module.body[base_model_index] = ast.ImportFrom(
+            module="infrahub_sdk.graphql", names=[ast.alias(name=GraphQLReturnTypeModel.__name__)]
+        )
+        return module
+
+    @staticmethod
+    def replace_base_model_class(module: ast.Module) -> ast.Module:
+        """Replace the BaseModel inserted by Ariadne with the GraphQLReturnTypeModel class."""
+        for item in module.body:
+            if not isinstance(item, ast.ClassDef):
+                continue
+
+            for base in item.bases:
+                if isinstance(base, ast.Name) and base.id == "BaseModel":
+                    base.id = GraphQLReturnTypeModel.__name__
+        return module
+
+    def insert_future_annotation(self, module: ast.Module) -> ast.Module:
+        """Insert the future annotation at the beginning of the module."""
+        module.body.insert(0, ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")]))
+        return module
+
+    def generate_result_types_module(
+        self,
+        module: ast.Module,
+        operation_definition: ExecutableDefinitionNode,  # noqa: ARG002
+    ) -> ast.Module:
+        module = self.insert_future_annotation(module)
+        module = self.replace_base_model_import(module)
+        module = self.replace_base_model_class(module)
+
+        return module

infrahub_sdk/graphql/query.py

@@ -0,0 +1,77 @@
+from __future__ import annotations
+
+from typing import Any
+
+from .renderers import render_input_block, render_query_block, render_variables_to_string
+
+
+class BaseGraphQLQuery:
+    query_type: str = "not-defined"
+    indentation: int = 4
+
+    def __init__(self, query: dict, variables: dict | None = None, name: str | None = None):
+        self.query = query
+        self.variables = variables
+        self.name = name or ""
+
+    def render_first_line(self) -> str:
+        first_line = self.query_type
+
+        if self.name:
+            first_line += " " + self.name
+
+        if self.variables:
+            first_line += f" ({render_variables_to_string(self.variables)})"
+
+        first_line += " {"
+
+        return first_line
+
+
+class Query(BaseGraphQLQuery):
+    query_type = "query"
+
+    def render(self, convert_enum: bool = False) -> str:
+        lines = [self.render_first_line()]
+        lines.extend(
+            render_query_block(
+                data=self.query, indentation=self.indentation, offset=self.indentation, convert_enum=convert_enum
+            )
+        )
+        lines.append("}")
+
+        return "\n" + "\n".join(lines) + "\n"
+
+
+class Mutation(BaseGraphQLQuery):
+    query_type = "mutation"
+
+    def __init__(self, *args: Any, mutation: str, input_data: dict, **kwargs: Any):
+        self.input_data = input_data
+        self.mutation = mutation
+        super().__init__(*args, **kwargs)
+
+    def render(self, convert_enum: bool = False) -> str:
+        lines = [self.render_first_line()]
+        lines.append(" " * self.indentation + f"{self.mutation}(")
+        lines.extend(
+            render_input_block(
+                data=self.input_data,
+                indentation=self.indentation,
+                offset=self.indentation * 2,
+                convert_enum=convert_enum,
+            )
+        )
+        lines.append(" " * self.indentation + "){")
+        lines.extend(
+            render_query_block(
+                data=self.query,
+                indentation=self.indentation,
+                offset=self.indentation * 2,
+                convert_enum=convert_enum,
+            )
+        )
+        lines.append(" " * self.indentation + "}")
+        lines.append("}")
+
+        return "\n" + "\n".join(lines) + "\n"