timeout.rb in graphql-1.11.5

- old
+ new

@@ -5,11 +5,11 @@
     # This plugin will stop resolving new fields after `max_seconds` have elapsed.
     # After the time has passed, any remaining fields will be `nil`, with errors added
     # to the `errors` key. Any already-resolved fields will be in the `data` key, so
     # you'll get a partial response.
     #
-    # You can subclass `GraphQL::Schema::Timeout` and override the `handle_timeout` method
+    # You can subclass `GraphQL::Schema::Timeout` and override `max_seconds` and/or `handle_timeout`
     # to provide custom logic when a timeout error occurs.
     #
     # Note that this will stop a query _in between_ field resolutions, but
     # it doesn't interrupt long-running `resolve` functions. Be sure to use
     # timeout options for external connections. For more info, see
@@ -31,12 +31,10 @@
     #   class MySchema < GraphQL::Schema
     #     use MyTimeout, max_seconds: 2
     #   end
     #
     class Timeout
-      attr_reader :max_seconds
-
       def self.use(schema, **options)
         tracer = new(**options)
         schema.tracer(tracer)
       end
 
@@ -46,44 +44,60 @@
       end
 
       def trace(key, data)
         case key
         when 'execute_multiplex'
-          timeout_state = {
-            timeout_at: Process.clock_gettime(Process::CLOCK_MONOTONIC, :millisecond) + max_seconds * 1000,
-            timed_out: false
-          }
-
           data.fetch(:multiplex).queries.each do |query|
+            timeout_duration_s = max_seconds(query)
+            timeout_state = if timeout_duration_s == false
+              # if the method returns `false`, don't apply a timeout
+              false
+            else
+              now = Process.clock_gettime(Process::CLOCK_MONOTONIC, :millisecond)
+              timeout_at = now + (max_seconds(query) * 1000)
+              {
+                timeout_at: timeout_at,
+                timed_out: false
+              }
+            end
             query.context.namespace(self.class)[:state] = timeout_state
           end
 
           yield
         when 'execute_field', 'execute_field_lazy'
-          query = data[:context] ? data.fetch(:context).query : data.fetch(:query)
-          timeout_state = query.context.namespace(self.class).fetch(:state)
-          if Process.clock_gettime(Process::CLOCK_MONOTONIC, :millisecond) > timeout_state.fetch(:timeout_at)
+          query_context = data[:context] || data[:query].context
+          timeout_state = query_context.namespace(self.class).fetch(:state)
+          # If the `:state` is `false`, then `max_seconds(query)` opted out of timeout for this query.
+          if timeout_state != false && Process.clock_gettime(Process::CLOCK_MONOTONIC, :millisecond) > timeout_state.fetch(:timeout_at)
             error = if data[:context]
-              context = data.fetch(:context)
-              GraphQL::Schema::Timeout::TimeoutError.new(context.parent_type, context.field)
+              GraphQL::Schema::Timeout::TimeoutError.new(query_context.parent_type, query_context.field)
             else
               field = data.fetch(:field)
               GraphQL::Schema::Timeout::TimeoutError.new(field.owner, field)
             end
 
             # Only invoke the timeout callback for the first timeout
-            unless timeout_state[:timed_out]
+            if !timeout_state[:timed_out]
               timeout_state[:timed_out] = true
-              handle_timeout(error, query)
+              handle_timeout(error, query_context.query)
             end
 
             error
           else
             yield
           end
         else
           yield
         end
+      end
+
+      # Called at the start of each query.
+      # The default implementation returns the `max_seconds:` value from installing this plugin.
+      #
+      # @param query [GraphQL::Query] The query that's about to run
+      # @return [Integer, false] The number of seconds after which to interrupt query execution and call {#handle_error}, or `false` to bypass the timeout.
+      def max_seconds(query)
+        @max_seconds
       end
 
       # Invoked when a query times out.
       # @param error [GraphQL::Schema::Timeout::TimeoutError]
       # @param query [GraphQL::Error]