Merge pull request #129 from evo-company/fix-pass-context

fix: pass context to dispatch
evo-company · Sep 7, 2023 · aefbf68 · aefbf68
2 parents 198e0aa + 203cf76
commit aefbf68
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 5 deletions.
diff --git a/hiku/endpoint/graphql.py b/hiku/endpoint/graphql.py
@@ -368,13 +368,17 @@ async def dispatch(
             return list(
                 await gather(
                     *(
-                        super(AsyncGraphQLEndpoint, self).dispatch(item)
+                        super(AsyncGraphQLEndpoint, self).dispatch(
+                            item, context
+                        )
                         for item in data
                     )
                 )
             )
         else:
-            return await super(AsyncGraphQLEndpoint, self).dispatch(data)
+            return await super(AsyncGraphQLEndpoint, self).dispatch(
+                data, context
+            )
 
 
 class AsyncBatchGraphQLEndpoint(AsyncGraphQLEndpoint):

diff --git a/hiku/federation/endpoint.py b/hiku/federation/endpoint.py
@@ -373,10 +373,12 @@ async def dispatch(
     ) -> SingleOrBatchedResponse:
         if isinstance(data, list):
             return list(
-                await gather(*(super().dispatch(item) for item in data))
+                await gather(
+                    *(super().dispatch(item, context) for item in data)
+                )
             )
 
-        return await super().dispatch(data)
+        return await super().dispatch(data, context)
 
 
 class AsyncBatchFederatedGraphQLEndpoint(AsyncFederatedGraphQLEndpoint):

diff --git a/tests/test_source_sqlalchemy.py b/tests/test_source_sqlalchemy.py
@@ -356,5 +356,5 @@ def check(self, src, value):
             self.graph,
         )
 
-        result = endpoint.dispatch({"query": src}, context={SA_ENGINE_KEY: sa_engine})
+        result = endpoint.dispatch({"query": src}, {SA_ENGINE_KEY: sa_engine})
         check_result(result['data'], value)