From 681455a0db8185013d9d3dcc6f376ac22f8a72cb Mon Sep 17 00:00:00 2001
From: Arun Suresh Kumar <asuresh960@gmail.com>
Date: Sun, 10 Dec 2023 18:34:30 +0530
Subject: [PATCH] Bug Fix: build_schema

---
 graphene_federation/main.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/graphene_federation/main.py b/graphene_federation/main.py
index d8f177a..e1a1fa0 100644
--- a/graphene_federation/main.py
+++ b/graphene_federation/main.py
@@ -30,5 +30,9 @@ def build_schema(
     schema = schema or Schema(query=query, mutation=mutation, **kwargs)
     schema.auto_camelcase = kwargs.get("auto_camelcase", True)
     schema.federation_version = 2 if enable_federation_2 else 1
-    federation_query = _get_query(schema, query)
-    return Schema(query=federation_query, mutation=mutation, **kwargs)
+    federation_query = _get_query(schema, schema.query if schema else query)
+    return Schema(
+        query=federation_query,
+        mutation=schema.mutation if schema else mutation,
+        **kwargs
+    )