Address reviewer comments

- Add 1 for all fields that call Gitaly (with resolvers or without)
- Clarify comment regarding Gitaly call alert
- Expose predicate `calls_gitaly?` instead of ivar

5 files changed, 50 insertions, 40 deletions

diff --git a/app/graphql/types/base_field.rb b/app/graphql/types/base_field.rb
index 42c7eb6b485..6b377e88e16 100644
--- a/app/graphql/types/base_field.rb
+++ b/app/graphql/types/base_field.rb
@@ -4,8 +4,6 @@ module Types
   class BaseField < GraphQL::Schema::Field
     prepend Gitlab::Graphql::Authorize
 
-    attr_reader :calls_gitaly
-
     DEFAULT_COMPLEXITY = 1
 
     def initialize(*args, **kwargs, &block)
@@ -17,19 +15,12 @@ module Types
 
     def base_complexity
       complexity = DEFAULT_COMPLEXITY
-      complexity += 1 if @calls_gitaly
+      complexity += 1 if calls_gitaly?
       complexity
     end
 
-    def calls_gitaly_check(calls)
-      return if @calls_gitaly
-      return if calls < 1
-
-      # Will inform you if :calls_gitaly should be true or false based on the number of Gitaly calls
-      # involved with the request.
-      raise "Gitaly is called for field '#{name}' #{"on type #{owner.name} " if owner}- please add `calls_gitaly: true` to the field declaration"
-    rescue => e
-      Gitlab::Sentry.track_exception(e)
+    def calls_gitaly?
+      @calls_gitaly
     end
 
     private
@@ -51,6 +42,7 @@ module Types
       proc do |ctx, args, child_complexity|
         # Resolvers may add extra complexity depending on used arguments
         complexity = child_complexity + self.resolver&.try(:resolver_complexity, args, child_complexity: child_complexity).to_i
+        complexity += 1 if calls_gitaly?
 
         field_defn = to_graphql
 
diff --git a/lib/gitlab/graphql/calls_gitaly/instrumentation.rb b/lib/gitlab/graphql/calls_gitaly/instrumentation.rb
index 08e98028755..e2733a1416f 100644
--- a/lib/gitlab/graphql/calls_gitaly/instrumentation.rb
+++ b/lib/gitlab/graphql/calls_gitaly/instrumentation.rb
@@ -7,7 +7,7 @@ module Gitlab
         # Check if any `calls_gitaly: true` declarations need to be added
         def instrument(_type, field)
           type_object = field.metadata[:type_class]
-          return field unless type_object && type_object.respond_to?(:calls_gitaly_check)
+          return field unless type_object && type_object.respond_to?(:calls_gitaly?)
 
           old_resolver_proc = field.resolve_proc
 
@@ -15,16 +15,24 @@ module Gitlab
             previous_gitaly_call_count = Gitlab::GitalyClient.get_request_count
             result = old_resolver_proc.call(typed_object, args, ctx)
             current_gitaly_call_count = Gitlab::GitalyClient.get_request_count
-            type_object.calls_gitaly_check(current_gitaly_call_count - previous_gitaly_call_count)
+            calls_gitaly_check(type_object, current_gitaly_call_count - previous_gitaly_call_count)
             result
-          rescue => e
-            ap "#{e.message}"
           end
 
           field.redefine do
             resolve(gitaly_wrapped_resolve)
           end
         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")
+          Gitlab::Sentry.track_exception(error)
+        end
       end
     end
   end
diff --git a/spec/graphql/types/base_field_spec.rb b/spec/graphql/types/base_field_spec.rb
index 0be83ea60c4..10913e530cf 100644
--- a/spec/graphql/types/base_field_spec.rb
+++ b/spec/graphql/types/base_field_spec.rb
@@ -74,9 +74,11 @@ describe Types::BaseField do
     context 'calls_gitaly' do
       context 'for fields with a resolver' do
         it 'adds 1 if true' do
-          field = described_class.new(name: 'test', type: GraphQL::STRING_TYPE, null: true, calls_gitaly: true)
+          with_gitaly_field = described_class.new(name: 'test', type: GraphQL::STRING_TYPE, resolver_class: resolver, null: true, calls_gitaly: true)
+          without_gitaly_field = described_class.new(name: 'test', type: GraphQL::STRING_TYPE, resolver_class: resolver, null: true)
+          base_result = without_gitaly_field.to_graphql.complexity.call({}, {}, 2)
 
-          expect(field.to_graphql.complexity).to eq 2
+          expect(with_gitaly_field.to_graphql.complexity.call({}, {}, 2)).to eq base_result + 1
         end
       end
 
@@ -100,26 +102,5 @@ describe Types::BaseField do
         expect(field.to_graphql.complexity).to eq 12
       end
     end
-
-    describe '#calls_gitaly_check' do
-      let(:gitaly_field) { described_class.new(name: 'test', type: GraphQL::STRING_TYPE, null: true, calls_gitaly: true) }
-      let(:no_gitaly_field) { described_class.new(name: 'test', type: GraphQL::STRING_TYPE, null: true, calls_gitaly: false) }
-
-      context 'if there are no Gitaly calls' do
-        it 'does not raise an error if calls_gitaly is false' do
-          expect { no_gitaly_field.send(:calls_gitaly_check, 0) }.not_to raise_error
-        end
-      end
-
-      context 'if there is at least 1 Gitaly call' do
-        it 'does not raise an error if calls_gitaly is true' do
-          expect { gitaly_field.send(:calls_gitaly_check, 1) }.not_to raise_error
-        end
-
-        it 'raises an error if calls_gitaly: is false or not defined' do
-          expect { no_gitaly_field.send(:calls_gitaly_check, 1) }.to raise_error(/please add `calls_gitaly: true`/)
-        end
-      end
-    end
   end
 end
diff --git a/spec/lib/gitlab/graphql/calls_gitaly/instrumentation_spec.rb b/spec/lib/gitlab/graphql/calls_gitaly/instrumentation_spec.rb
new file mode 100644
index 00000000000..92d200bfd4e
--- /dev/null
+++ b/spec/lib/gitlab/graphql/calls_gitaly/instrumentation_spec.rb
@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+require 'spec_helper'
+
+describe Gitlab::Graphql::CallsGitaly::Instrumentation do
+  subject { described_class.new }
+
+  context 'when considering complexity' do
+    describe '#calls_gitaly_check' do
+      let(:gitaly_field) { Types::BaseField.new(name: 'test', type: GraphQL::STRING_TYPE, null: true, calls_gitaly: true) }
+      let(:no_gitaly_field) { Types::BaseField.new(name: 'test', type: GraphQL::STRING_TYPE, null: true, calls_gitaly: false) }
+
+      context 'if there are no Gitaly calls' do
+        it 'does not raise an error if calls_gitaly is false' do
+          expect { subject.send(:calls_gitaly_check, no_gitaly_field, 0) }.not_to raise_error
+        end
+      end
+
+      context 'if there is at least 1 Gitaly call' do
+        it 'does not raise an error if calls_gitaly is true' do
+          expect { subject.send(:calls_gitaly_check, gitaly_field, 1) }.not_to raise_error
+        end
+
+        it 'raises an error if calls_gitaly: is false or not defined' do
+          expect { subject.send(:calls_gitaly_check, no_gitaly_field, 1) }.to raise_error(/please add `calls_gitaly: true`/)
+        end
+      end
+    end
+  end
+end
diff --git a/spec/requests/api/graphql_spec.rb b/spec/requests/api/graphql_spec.rb
index d78b17827a6..67371cb35b6 100644
--- a/spec/requests/api/graphql_spec.rb
+++ b/spec/requests/api/graphql_spec.rb
@@ -137,7 +137,7 @@ describe 'GraphQL' do
     let(:user) { create(:user) }
 
     let(:query) do
-      graphql_query_for('project', { 'fullPath' => project.full_path }, %w(forksCount))
+      graphql_query_for('project', { 'fullPath' => project.full_path }, %w(id))
     end
 
     before do