From ceb0d9d78ff51ff05629707c99fd7d49cd82ee6c Mon Sep 17 00:00:00 2001
From: Jens Vogler <vogler@init7.net>
Date: Wed, 2 Feb 2022 14:21:45 +0100
Subject: [PATCH] feat: plugins can extend the GraphQL schema

---
 netbox/extras/plugins/__init__.py | 34 +++++++++++++++++++++++++
 netbox/netbox/graphql/schema.py   | 42 ++++++++++++++++++++++++++++---
 2 files changed, 73 insertions(+), 3 deletions(-)

diff --git a/netbox/extras/plugins/__init__.py b/netbox/extras/plugins/__init__.py
index f9a7856ea..6afa3abb0 100644
--- a/netbox/extras/plugins/__init__.py
+++ b/netbox/extras/plugins/__init__.py
@@ -1,7 +1,10 @@
 import collections
 import inspect
+import graphene
 from packaging import version
 
+from graphql.type.definition import GraphQLType
+
 from django.apps import AppConfig
 from django.core.exceptions import ImproperlyConfigured
 from django.template.loader import get_template
@@ -15,6 +18,10 @@ from extras.plugins.utils import import_object
 # Initialize plugin registry stores
 registry['plugin_template_extensions'] = collections.defaultdict(list)
 registry['plugin_menu_items'] = {}
+registry['graphql_queries'] = []
+registry['graphql_mutations'] = []
+registry['graphql_subscriptions'] = []
+registry['graphql_types'] = []
 
 
 #
@@ -50,6 +57,12 @@ class PluginConfig(AppConfig):
     # Django-rq queues dedicated to the plugin
     queues = []
 
+    # Graphene GraphQL schema extension points
+    graphql_queries = []
+    graphql_mutations = []
+    graphql_subscriptions = []
+    graphql_types = []
+
     # Default integration paths. Plugin authors can override these to customize the paths to
     # integrated components.
     template_extensions = 'template_content.template_extensions'
@@ -67,6 +80,10 @@ class PluginConfig(AppConfig):
         if menu_items is not None:
             register_menu_items(self.verbose_name, menu_items)
 
+        # Register GraphQL schema extensions
+        register_graphql_schema(self.graphql_queries, self.graphql_mutations, self.graphql_subscriptions,
+                                self.graphql_types)
+
     @classmethod
     def validate(cls, user_config, netbox_version):
 
@@ -242,3 +259,20 @@ def register_menu_items(section_name, class_list):
                 raise TypeError(f"{button} must be an instance of extras.plugins.PluginMenuButton")
 
     registry['plugin_menu_items'][section_name] = class_list
+
+
+def register_graphql_schema(queries, mutations, subscriptions, types):
+    """
+    Register extensions to the GraphQL schema
+    """
+    for object_type in [*queries, *mutations, *subscriptions]:
+        if not issubclass(object_type, graphene.ObjectType):
+            raise TypeError(f"{object_type} must be a subclass of graphene.types.objecttype.ObjectType")
+
+    for graphql_type in types:
+        if not issubclass(graphql_type, GraphQLType):
+            raise TypeError(f"f{graphql_type} must be a subclass of graphql.type.definition.GraphQLType")
+
+    registry['graphql_queries'].extend(queries)
+    registry['graphql_subscriptions'].extend(subscriptions)
+    registry['graphql_types'].extend(types)
diff --git a/netbox/netbox/graphql/schema.py b/netbox/netbox/graphql/schema.py
index 812c1656d..c8d73f6b2 100644
--- a/netbox/netbox/graphql/schema.py
+++ b/netbox/netbox/graphql/schema.py
@@ -3,14 +3,14 @@ import graphene
 from circuits.graphql.schema import CircuitsQuery
 from dcim.graphql.schema import DCIMQuery
 from extras.graphql.schema import ExtrasQuery
+from extras.plugins import registry
 from ipam.graphql.schema import IPAMQuery
 from tenancy.graphql.schema import TenancyQuery
 from users.graphql.schema import UsersQuery
 from virtualization.graphql.schema import VirtualizationQuery
 from wireless.graphql.schema import WirelessQuery
 
-
-class Query(
+query_types = [
     CircuitsQuery,
     DCIMQuery,
     ExtrasQuery,
@@ -19,9 +19,45 @@ class Query(
     UsersQuery,
     VirtualizationQuery,
     WirelessQuery,
+    *registry['graphql_queries'],
+]
+
+mutation_types = [
+    *registry['graphql_mutations'],
+]
+
+subscription_types = [
+    *registry['graphql_subscriptions'],
+]
+
+extra_types = [
+    *registry['graphql_types'],
+]
+
+
+class Query(
+    *query_types,
     graphene.ObjectType
 ):
     pass
 
 
-schema = graphene.Schema(query=Query, auto_camelcase=False)
+class Mutation(
+    *mutation_types,
+    graphene.ObjectType
+):
+    pass
+
+
+class Subscription(
+    *subscription_types,
+    graphene.ObjectType
+):
+    pass
+
+
+schema = graphene.Schema(query=Query,
+                         mutation=Mutation if len(mutation_types) > 0 else None,
+                         subscription=Subscription if len(subscription_types) > 0 else None,
+                         types=extra_types if len(extra_types) > 0 else None,
+                         auto_camelcase=False)