Address reviewer comments

- Remove Gitaly call check for fields that have a constant complexity
declared
- Add associated test

2 files changed, 6 insertions, 5 deletions

diff --git a/lib/gitlab/graphql/calls_gitaly.rb b/lib/gitlab/graphql/calls_gitaly.rb
index f75941e269f..40cd74a34f2 100644
--- a/lib/gitlab/graphql/calls_gitaly.rb
+++ b/lib/gitlab/graphql/calls_gitaly.rb
@@ -2,8 +2,8 @@
 
 module Gitlab
   module Graphql
-    # Allow fields to declare permissions their objects must have. The field
-    # will be set to nil unless all required permissions are present.
+    # Wraps the field resolution to count Gitaly calls before and after.
+    # Raises an error if the field calls Gitaly but hadn't declared such.
     module CallsGitaly
       extend ActiveSupport::Concern
 
diff --git a/lib/gitlab/graphql/calls_gitaly/instrumentation.rb b/lib/gitlab/graphql/calls_gitaly/instrumentation.rb
index e2733a1416f..fbd5e348c7d 100644
--- a/lib/gitlab/graphql/calls_gitaly/instrumentation.rb
+++ b/lib/gitlab/graphql/calls_gitaly/instrumentation.rb
@@ -5,9 +5,11 @@ module Gitlab
     module CallsGitaly
       class Instrumentation
         # Check if any `calls_gitaly: true` declarations need to be added
+        # Do nothing if a constant complexity was provided
         def instrument(_type, field)
           type_object = field.metadata[:type_class]
-          return field unless type_object && type_object.respond_to?(:calls_gitaly?)
+          return field unless type_object.respond_to?(:calls_gitaly?)
+          return field if type_object.constant_complexity? || type_object.calls_gitaly?
 
           old_resolver_proc = field.resolve_proc
 
@@ -25,12 +27,11 @@ module Gitlab
         end
 
         def calls_gitaly_check(type_object, calls)
-          return if type_object.calls_gitaly?
           return if calls < 1
 
           # Will inform you if there needs to be `calls_gitaly: true` as a kwarg in the field declaration
           # if there is at least 1 Gitaly call involved with the field resolution.
-          error = RuntimeError.new("Gitaly is called for field '#{type_object.name}' on #{type_object.owner.try(:name)} - please add `calls_gitaly: true` to the field declaration")
+          error = RuntimeError.new("Gitaly is called for field '#{type_object.name}' on #{type_object.owner.try(:name)} - please either specify a constant complexity or add `calls_gitaly: true` to the field declaration")
           Gitlab::Sentry.track_exception(error)
         end
       end