From 1c6983ff4ab9674c83aee0b67468bc968d0ff18b Mon Sep 17 00:00:00 2001
From: Noah Metz <noah@metznet.ca>
Date: Sun, 2 Jul 2023 09:34:36 -0600
Subject: [PATCH] Exposed load function maps in graph context constructor

---
 graph.go      | 10 +++++++++-
 graph_test.go |  4 ++--
 2 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/graph.go b/graph.go
index 0ad3fd7..020f000 100644
--- a/graph.go
+++ b/graph.go
@@ -195,7 +195,7 @@ func LoadNodeRecurse(ctx * GraphContext, id NodeID, loaded_nodes map[NodeID]Grap
   return node, nil
 }
 
-func NewGraphContext(db * badger.DB, log Logger, types TypeList, type_map ObjTypeMap, queries FieldMap, subscriptions FieldMap, mutations FieldMap) * GraphContext {
+func NewGraphContext(db * badger.DB, log Logger, state_loads StateLoadMap, node_loads NodeLoadMap, types TypeList, type_map ObjTypeMap, queries FieldMap, subscriptions FieldMap, mutations FieldMap) * GraphContext {
   gql, err := NewGQLContext(types, type_map, queries, subscriptions, mutations)
   if err != nil {
     panic(err)
@@ -217,6 +217,14 @@ func NewGraphContext(db * badger.DB, log Logger, types TypeList, type_map ObjTyp
     },
   }
 
+  for name, fn := range(state_loads) {
+    ctx.StateLoadFuncs[name] = fn
+  }
+
+  for name, fn := range(node_loads) {
+    ctx.NodeLoadFuncs[name] = fn
+  }
+
   return &ctx
 }
 
diff --git a/graph_test.go b/graph_test.go
index 9673995..a723a12 100644
--- a/graph_test.go
+++ b/graph_test.go
@@ -58,7 +58,7 @@ func logTestContext(t * testing.T, components []string) * GraphContext {
     t.Fatal(err)
   }
 
-  return NewGraphContext(db, NewConsoleLogger(components), TypeList{}, ObjTypeMap{}, FieldMap{}, FieldMap{}, FieldMap{})
+  return NewGraphContext(db, NewConsoleLogger(components), StateLoadMap{}, NodeLoadMap{}, TypeList{}, ObjTypeMap{}, FieldMap{}, FieldMap{}, FieldMap{})
 }
 
 func testContext(t * testing.T) * GraphContext {
@@ -67,7 +67,7 @@ func testContext(t * testing.T) * GraphContext {
     t.Fatal(err)
   }
 
-  return NewGraphContext(db, NewConsoleLogger([]string{}), TypeList{}, ObjTypeMap{}, FieldMap{}, FieldMap{}, FieldMap{})
+  return NewGraphContext(db, NewConsoleLogger([]string{}), StateLoadMap{}, NodeLoadMap{}, TypeList{}, ObjTypeMap{}, FieldMap{}, FieldMap{}, FieldMap{})
 }
 
 func fatalErr(t * testing.T, err error) {