Repository: mirumee/ariadne-codegen (python)

Issue:
Feature request: support Apollo Federation
```
File "/.venv/lib/python3.11/site-packages/ariadne_codegen/schema.py", line 93, in get_graphql_schema_from_path
    schema: GraphQLSchema = build_ast_schema(graphql_ast)
                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/.venv/lib/python3.11/site-packages/graphql/utilities/build_ast_schema.py", line 45, in build_ast_schema
    assert_valid_sdl(document_ast)
  File "/.venv/lib/python3.11/site-packages/graphql/validation/validate.py", line 119, in assert_valid_sdl
    raise TypeError("\n\n".join(error.message for error in errors))
TypeError: Unknown directive '@key'.

Unknown directive '@extends'.

Unknown directive '@external'.
```

Maybe support `assume_valid = true` or/and `assume_valid_sdl = true` in `pyproject.toml` ?

Pull request:
Process schema hook
This pr:
- Adds `process_schema` plugin hook
- Changes flow of parsing schema: we parse schema allowing invalid ones(`assume_valid=True`), then `process_schema` hook is called, only then we validate processed schema.
- Adds `plugins` setting to `graphqlschema` strategy.

Resolves #149 Utilizing `process_schema` plugin will be able to inject apollo federation directives definitions into schema.

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 336fff8..e1556cf 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -14,6 +14,7 @@
 - Fixed `ShorterResultsPlugin` that generated faulty code for discriminated unions.
 - Changed generator to ignore unused fragments which should be unpacked in queries.
 - Changed type hints for parse and serialize methods of scalars to `typing.Any`.
+- Added `process_schema` plugin hook.
 
 
 ## 0.6.0 (2023-04-18)
diff --git a/PLUGINS.md b/PLUGINS.md
index 567f6db..0ecbb60 100644
--- a/PLUGINS.md
+++ b/PLUGINS.md
@@ -305,6 +305,16 @@ def generate_fragments_module(
 Hook executed on generation of fragments module. Module has classes representing all fragments from provided queries. Later this module will be saved as `{fragments_module_name}.py`, `fragments_module_name` is taken from config.
 
 
+### process_schema
+
+```py
+def process_schema(self, schema: GraphQLSchema) -> GraphQLSchema:
+```
+
+Hook executed on creating `GraphQLSchema` object from path or url provided in settings. During parsing `assume_valid` is set to `True`. Then this hook is called, and only after that `graphql.assert_valid_schema` is used to validate schema.
+To ensure all plugins have current version of schema, result of this hook is propagated to all plugins, updating their `schema` field. 
+
+
 ## Example
 
 This example plugin adds `__version__ = "..."` to generated `__init__.py` file.
diff --git a/README.md b/README.md
index 056bd84..b9c0350 100644
--- a/README.md
+++ b/README.md
@@ -255,7 +255,7 @@ Instead of generating client, you can generate file with a copy of GraphQL schem
 ariadne-codegen graphqlschema
 ```
 
-`graphqlschema` mode reads configuration from the same place as [`client`](#configuration) but uses only `schema_path`, `remote_schema_url`, `remote_schema_headers` and `remote_schema_verify_ssl` options with addition to some extra options specific to it:
+`graphqlschema` mode reads configuration from the same place as [`client`](#configuration) but uses only `schema_path`, `remote_schema_url`, `remote_schema_headers`, `remote_schema_verify_ssl` and `plugins` options with addition to some extra options specific to it:
 
 - `target_file_path` (defaults to `"schema.py"`) - destination path for generated file
 - `schema_variable_name` (defaults to `"schema"`) - name for schema variable, must be valid python identifier
diff --git a/ariadne_codegen/main.py b/ariadne_codegen/main.py
index d9c123c..296ae89 100644
--- a/ariadne_codegen/main.py
+++ b/ariadne_codegen/main.py
@@ -1,6 +1,7 @@
 import sys
 
 import click
+from graphql import assert_valid_schema
 
 from .client_generators.package import PackageGenerator
 from .config import get_client_settings, get_config_dict, get_graphql_schema_settings
@@ -37,6 +38,7 @@ def main(strategy=Strategy.CLIENT, config=None):
 
 def client(config_dict):
     settings = get_client_settings(config_dict)
+
     if settings.schema_path:
         schema = get_graphql_schema_from_path(settings.schema_path)
         schema_source = settings.schema_path
@@ -48,6 +50,14 @@ def client(config_dict):
         )
         schema_source = settings.remote_schema_url
 
+    plugin_manager = PluginManager(
+        schema=schema,
+        config_dict=config_dict,
+        plugins_types=get_plugins_types(settings.plugins),
+    )
+    schema = plugin_manager.process_schema(schema)
+    assert_valid_schema(schema)
+
     definitions = get_graphql_queries(settings.queries_path)
     queries = filter_operations_definitions(definitions)
     fragments = filter_fragments_definitions(definitions)
@@ -72,11 +82,7 @@ def client(config_dict):
         async_client=settings.async_client,
         files_to_include=settings.files_to_include,
         custom_scalars=settings.scalars,
-        plugin_manager=PluginManager(
-            schema=schema,
-            config_dict=config_dict,
-            plugins_types=get_plugins_types(settings.plugins),
-        ),
+        plugin_manager=plugin_manager,
     )
     for query in queries:
         package_generator.add_operation(query)
@@ -87,7 +93,6 @@ def client(config_dict):
 
 def graphql_schema(config_dict):
     settings = get_graphql_schema_settings(config_dict)
-    sys.stdout.write(settings.used_settings_message)
 
     schema = (
         get_graphql_schema_from_path(settings.schema_path)
@@ -98,6 +103,15 @@ def graphql_schema(config_dict):
             verify_ssl=settings.remote_schema_verify_ssl,
         )
     )
+    plugin_manager = PluginManager(
+        schema=schema,
+        config_dict=config_dict,
+        plugins_types=get_plugins_types(settings.plugins),
+    )
+    schema = plugin_manager.process_schema(schema)
+    assert_valid_schema(schema)
+
+    sys.stdout.write(settings.used_settings_message)
 
     generate_graphql_schema_file(
         schema=schema,
diff --git a/ariadne_codegen/plugins/base.py b/ariadne_codegen/plugins/base.py
index d515888..2bd77ec 100644
--- a/ariadne_codegen/plugins/base.py
+++ b/ariadne_codegen/plugins/base.py
@@ -158,3 +158,6 @@ class Plugin:
         fragments_definitions: Dict[str, FragmentDefinitionNode],
     ) -> ast.Module:
         return module
+
+    def process_schema(self, schema: GraphQLSchema) -> GraphQLSchema:
+        return schema
diff --git a/ariadne_codegen/plugins/manager.py b/ariadne_codegen/plugins/manager.py
index a736532..9c547d7 100644
--- a/ariadne_codegen/plugins/manager.py
+++ b/ariadne_codegen/plugins/manager.py
@@ -202,3 +202,13 @@ class PluginManager:
             module,
             fragments_definitions=fragments_definitions,
         )
+
+    def process_schema(self, schema: GraphQLSchema) -> GraphQLSchema:
+        processed_schema = schema
+        for plugin in self.plugins:
+            processed_schema = plugin.process_schema(processed_schema)
+
+            for plugin in self.plugins:
+                plugin.schema = processed_schema
+
+        return processed_schema
diff --git a/ariadne_codegen/schema.py b/ariadne_codegen/schema.py
index f6e56aa..10f7c36 100644
--- a/ariadne_codegen/schema.py
+++ b/ariadne_codegen/schema.py
@@ -9,7 +9,6 @@ from graphql import (
     GraphQLSyntaxError,
     IntrospectionQuery,
     OperationDefinitionNode,
-    assert_valid_schema,
     build_ast_schema,
     build_client_schema,
     get_introspection_query,
@@ -44,7 +43,8 @@ def get_graphql_schema_from_url(
     url: str, headers: Optional[Dict[str, str]] = None, verify_ssl: bool = True
 ) -> GraphQLSchema:
     return build_client_schema(
-        introspect_remote_schema(url=url, headers=headers, verify_ssl=verify_ssl)
+        introspect_remote_schema(url=url, headers=headers, verify_ssl=verify_ssl),
+        assume_valid=True,
     )
 
 
@@ -90,8 +90,7 @@ def get_graphql_schema_from_path(schema_path: str) -> GraphQLSchema:
     """Get graphql schema build from provided path."""
     schema_str = load_graphql_files_from_path(Path(schema_path))
     graphql_ast = parse(schema_str)
-    schema: GraphQLSchema = build_ast_schema(graphql_ast)
-    assert_valid_schema(schema)
+    schema: GraphQLSchema = build_ast_schema(graphql_ast, assume_valid=True)
     return schema
 
 
diff --git a/ariadne_codegen/settings.py b/ariadne_codegen/settings.py
index 0dc7c7e..ebc69e3 100644
--- a/ariadne_codegen/settings.py
+++ b/ariadne_codegen/settings.py
@@ -25,6 +25,7 @@ class BaseSettings:
     remote_schema_url: Optional[str] = None
     remote_schema_headers: dict = field(default_factory=dict)
     remote_schema_verify_ssl: bool = True
+    plugins: List[str] = field(default_factory=list)
 
     def __post_init__(self):
         if not self.schema_path and not self.remote_schema_url:
@@ -54,7 +55,6 @@ class ClientSettings(BaseSettings):
     convert_to_snake_case: bool = True
     async_client: bool = True
     files_to_include: List[str] = field(default_factory=list)
-    plugins: List[str] = field(default_factory=list)
     scalars: Dict[str, ScalarData] = field(default_factory=dict)
 
     def __post_init__(self):
@@ -156,6 +156,12 @@ class GraphQLSchemaSettings(BaseSettings):
 
     @property
     def used_settings_message(self):
+        plugins_list = ",".join(self.plugins)
+        plugins_msg = (
+            f"Plugins to use: {plugins_list}"
+            if self.plugins
+            else "No plugin is being used."
+        )
         return dedent(
             f"""\
             Selected strategy: {Strategy.GRAPHQL_SCHEMA}
@@ -163,6 +169,7 @@ class GraphQLSchemaSettings(BaseSettings):
             Saving graphql schema to: {self.target_file_path}.
             Using {self.schema_variable_name} as variable name for schema.
             Using {self.type_map_variable_name} as variable name for type map.
+            {plugins_msg}
             """
         )
 


Test patch:
diff --git a/tests/plugins/test_manager.py b/tests/plugins/test_manager.py
index 9c9e94b..f0ff8c0 100644
--- a/tests/plugins/test_manager.py
+++ b/tests/plugins/test_manager.py
@@ -342,3 +342,39 @@ def test_generate_fragments_module_calls_plugins_generate_fragments_module(
     plugin1, plugin2 = plugin_manager_with_mocked_plugins.plugins
     assert plugin1.generate_fragments_module.called
     assert plugin2.generate_fragments_module.called
+
+
+def test_process_schema_calls_plugins_process_schema(
+    plugin_manager_with_mocked_plugins,
+):
+    plugin_manager_with_mocked_plugins.process_schema(GraphQLSchema())
+
+    plugin1, plugin2 = plugin_manager_with_mocked_plugins.plugins
+    assert plugin1.process_schema.called
+    assert plugin2.process_schema.called
+
+
+def test_process_schema_updates_plugins_schema_field():
+    class SchemaPugin(Plugin):
+        def process_schema(self, schema: GraphQLSchema) -> GraphQLSchema:
+            return GraphQLSchema()
+
+    class DumbPlugin(Plugin):
+        pass
+
+    org_schema = GraphQLSchema()
+    manager = PluginManager(
+        schema=org_schema,
+        plugins_types=[DumbPlugin, SchemaPugin, DumbPlugin],
+    )
+
+    dumb_plugin1, schema_plugin, dump_plugin2 = manager.plugins
+    assert dumb_plugin1.schema is org_schema
+    assert dump_plugin2.schema is org_schema
+    assert schema_plugin.schema is org_schema
+
+    manager.process_schema(org_schema)
+
+    assert dumb_plugin1.schema is not org_schema
+    assert dump_plugin2.schema is not org_schema
+    assert schema_plugin.schema is not org_schema
diff --git a/tests/test_schema.py b/tests/test_schema.py
index 02502bc..9b5b072 100644
--- a/tests/test_schema.py
+++ b/tests/test_schema.py
@@ -29,7 +29,13 @@ SECOND_SCHEMA = """
     }
 """
 
-INCORRECT_SCHEMA = """
+INVALID_SCHEMA = """
+    type Query {
+        test: String! @unknownDirective
+    }
+"""
+
+INVALID_SYNTAX_SCHEMA = """
     type Query {
         test: Custom
 
@@ -93,9 +99,16 @@ def single_file_schema(tmp_path_factory):
 
 
 @pytest.fixture
-def incorrect_schema_file(tmp_path_factory):
+def invalid_schema_file(tmp_path_factory):
+    file_ = tmp_path_factory.mktemp("schema").joinpath("schema.graphql")
+    file_.write_text(INVALID_SCHEMA, encoding="utf-8")
+    return file_
+
+
+@pytest.fixture
+def invalid_syntax_schema_file(tmp_path_factory):
     file_ = tmp_path_factory.mktemp("schema").joinpath("schema.graphql")
-    file_.write_text(INCORRECT_SCHEMA, encoding="utf-8")
+    file_.write_text(INVALID_SYNTAX_SCHEMA, encoding="utf-8")
     return file_
 
 
@@ -155,11 +168,11 @@ def test_read_graphql_file_returns_content_of_file(single_file_schema):
 
 
 def test_read_graphql_file_with_invalid_file_raises_invalid_graphql_syntax_exception(
-    incorrect_schema_file,
+    invalid_syntax_schema_file,
 ):
     with pytest.raises(InvalidGraphqlSyntax) as exc:
-        read_graphql_file(incorrect_schema_file)
-    assert str(incorrect_schema_file) in str(exc)
+        read_graphql_file(invalid_syntax_schema_file)
+    assert str(invalid_syntax_schema_file) in str(exc)
 
 
 def test_walk_graphql_files_returns_graphql_files_from_directory(schemas_directory):
@@ -200,20 +213,25 @@ def test_load_graphql_files_from_path_returns_schema_from_nested_directory(
 
 @pytest.mark.parametrize(
     "path_fixture",
-    ["single_file_schema", "schemas_directory", "schemas_nested_directories"],
+    [
+        "single_file_schema",
+        "invalid_schema_file",
+        "schemas_directory",
+        "schemas_nested_directories",
+    ],
     indirect=True,
 )
-def test_get_graphql_schema_returns_graphql_schema(path_fixture):
+def test_get_graphql_schema_from_path_returns_graphql_schema(path_fixture):
     assert isinstance(
         get_graphql_schema_from_path(path_fixture.as_posix()), GraphQLSchema
     )
 
 
-def test_get_graphql_schema_with_invalid_schema_raises_invalid_graphql_syntax_exception(
-    incorrect_schema_file,
+def test_get_graphql_schema_from_path_with_invalid_syntax_raises_invalid_graphql_syntax(
+    invalid_syntax_schema_file,
 ):
     with pytest.raises(InvalidGraphqlSyntax):
-        get_graphql_schema_from_path(incorrect_schema_file.as_posix())
+        get_graphql_schema_from_path(invalid_syntax_schema_file.as_posix())
 
 
 def test_introspect_remote_schema_called_with_invalid_url_raises_introspection_error(