lib/graphql/batch/loader.rb in graphql-batch-0.2.5 vs lib/graphql/batch/loader.rb in graphql-batch-0.3.0

@@ -1,11 +1,13 @@
 module GraphQL::Batch
   class Loader
     def self.for(*group_args)
       loader_key = [self].concat(group_args)
-      Executor.current.loaders[loader_key] ||= new(*group_args).tap do |loader|
+      executor = Executor.current
+      executor.loaders[loader_key] ||= new(*group_args).tap do |loader|
         loader.loader_key = loader_key
+        loader.executor = executor
       end
     end
 
     def self.load(key)
       self.for.load(key)
@@ -13,25 +15,21 @@
 
     def self.load_many(keys)
       self.for.load_many(keys)
     end
 
-    attr_accessor :loader_key
+    attr_accessor :loader_key, :executor
 
     def load(key)
-      loader = Executor.current.loaders[loader_key] ||= self
-      if loader != self
-        raise "load called on loader that wasn't registered with executor"
-      end
       cache[cache_key(key)] ||= begin
         queue << key
-        Promise.new
+        Promise.new.tap { |promise| promise.source = self }
       end
     end
 
     def load_many(keys)
-      Promise.all(keys.map { |key| load(key) })
+      ::Promise.all(keys.map { |key| load(key) })
     end
 
     def resolve #:nodoc:
       load_keys = queue
       return if load_keys.empty?
@@ -42,10 +40,20 @@
       each_pending_promise(load_keys) do |key, promise|
         promise.reject(err)
       end
     end
 
+    # For Promise#sync
+    def wait #:nodoc:
+      if executor
+        executor.loaders.delete(loader_key)
+        executor.resolve(self)
+      else
+        resolve
+      end
+    end
+
     protected
 
     # Fulfill the key with provided value, for use in #perform
     def fulfill(key, value)
       promise_for(key).fulfill(value)
@@ -89,10 +97,10 @@
       end
     end
 
     def check_for_broken_promises(load_keys)
       each_pending_promise(load_keys) do |key, promise|
-        promise.reject(BrokenPromiseError.new("#{self.class} didn't fulfill promise for key #{key.inspect}"))
+        promise.reject(::Promise::BrokenError.new("#{self.class} didn't fulfill promise for key #{key.inspect}"))
       end
     end
   end
 end