Enumerate fields with Gitaly calls

 - Add a complexity of 1 if Gitaly is called at least once
 - Add an error notification if `calls_gitaly` isn't right for a
 particular field

4 files changed, 106 insertions, 4 deletions

diff --git a/app/graphql/types/base_field.rb b/app/graphql/types/base_field.rb
index dd0d9105df6..ee23146f711 100644
--- a/app/graphql/types/base_field.rb
+++ b/app/graphql/types/base_field.rb
@@ -4,21 +4,30 @@ module Types
   class BaseField < GraphQL::Schema::Field
     prepend Gitlab::Graphql::Authorize
 
+    attr_reader :calls_gitaly
+
     DEFAULT_COMPLEXITY = 1
 
     def initialize(*args, **kwargs, &block)
+      @calls_gitaly = !!kwargs.delete(:calls_gitaly)
       kwargs[:complexity] ||= field_complexity(kwargs[:resolver_class])
 
       super(*args, **kwargs, &block)
     end
 
+    def base_complexity
+      complexity = DEFAULT_COMPLEXITY
+      complexity += 1 if @calls_gitaly
+      complexity
+    end
+
     private
 
     def field_complexity(resolver_class)
       if resolver_class
         field_resolver_complexity
       else
-        DEFAULT_COMPLEXITY
+        base_complexity
       end
     end
 
@@ -45,5 +54,17 @@ module Types
         complexity.to_i
       end
     end
+
+    def calls_gitaly_check
+      # Will inform you if :calls_gitaly should be true or false based on the number of Gitaly calls
+      # involved with the request.
+      if @calls_gitaly && Gitlab::GitalyClient.get_request_count == 0
+        raise "Gitaly is called for field '#{name}' - please add `calls_gitaly: true` to the field declaration"
+      elsif !@calls_gitaly && Gitlab::GitalyClient.get_request_count > 0
+        raise "Gitaly not called for field '#{name}' - please remove `calls_gitaly: true` from the field declaration"
+      end
+    rescue => e
+      Gitlab::Sentry.track_exception(e)
+    end
   end
 end
diff --git a/app/graphql/types/project_type.rb b/app/graphql/types/project_type.rb
index c25688ab043..87d5351f80f 100644
--- a/app/graphql/types/project_type.rb
+++ b/app/graphql/types/project_type.rb
@@ -26,7 +26,7 @@ module Types
     field :web_url, GraphQL::STRING_TYPE, null: true
 
     field :star_count, GraphQL::INT_TYPE, null: false
-    field :forks_count, GraphQL::INT_TYPE, null: false
+    field :forks_count, GraphQL::INT_TYPE, null: false, calls_gitaly: true # 4 times
 
     field :created_at, Types::TimeType, null: true
     field :last_activity_at, Types::TimeType, null: true
diff --git a/app/graphql/types/tree/tree_type.rb b/app/graphql/types/tree/tree_type.rb
index b947713074e..2023abc13f9 100644
--- a/app/graphql/types/tree/tree_type.rb
+++ b/app/graphql/types/tree/tree_type.rb
@@ -15,9 +15,9 @@ module Types
         Gitlab::Graphql::Representation::TreeEntry.decorate(obj.trees, obj.repository)
       end
 
-      field :submodules, Types::Tree::SubmoduleType.connection_type, null: false
+      field :submodules, Types::Tree::SubmoduleType.connection_type, null: false, calls_gitaly: true
 
-      field :blobs, Types::Tree::BlobType.connection_type, null: false, resolve: -> (obj, args, ctx) do
+      field :blobs, Types::Tree::BlobType.connection_type, null: false, calls_gitaly: true, resolve: -> (obj, args, ctx) do
         Gitlab::Graphql::Representation::TreeEntry.decorate(obj.blobs, obj.repository)
       end
       # rubocop: enable Graphql/AuthorizeTypes
diff --git a/spec/graphql/types/base_field_spec.rb b/spec/graphql/types/base_field_spec.rb
index 0d3c3e37daf..ebdfa3eaf4d 100644
--- a/spec/graphql/types/base_field_spec.rb
+++ b/spec/graphql/types/base_field_spec.rb
@@ -22,6 +22,24 @@ describe Types::BaseField do
       expect(field.to_graphql.complexity).to eq 1
     end
 
+    describe '#base_complexity' do
+      context 'with no gitaly calls' do
+        it 'defaults to 1' do
+          field = described_class.new(name: 'test', type: GraphQL::STRING_TYPE, null: true)
+
+          expect(field.base_complexity).to eq 1
+        end
+      end
+
+      context 'with a gitaly call' do
+        it 'adds 1 to the default value' do
+          field = described_class.new(name: 'test', type: GraphQL::STRING_TYPE, null: true, calls_gitaly: true)
+
+          expect(field.base_complexity).to eq 2
+        end
+      end
+    end
+
     it 'has specified value' do
       field = described_class.new(name: 'test', type: GraphQL::STRING_TYPE, null: true, complexity: 12)
 
@@ -52,5 +70,68 @@ describe Types::BaseField do
         end
       end
     end
+
+    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)
+
+          expect(field.to_graphql.complexity).to eq 2
+        end
+      end
+
+      context 'for fields without a resolver' do
+        it 'adds 1 if true' do
+          field = described_class.new(name: 'test', type: GraphQL::STRING_TYPE, null: true, calls_gitaly: true)
+
+          expect(field.to_graphql.complexity).to eq 2
+        end
+      end
+
+      it 'defaults to false' do
+        field = described_class.new(name: 'test', type: GraphQL::STRING_TYPE, null: true)
+
+        expect(field.base_complexity).to eq Types::BaseField::DEFAULT_COMPLEXITY
+      end
+
+      it 'is overridden by declared complexity value' do
+        field = described_class.new(name: 'test', type: GraphQL::STRING_TYPE, null: true, calls_gitaly: true, complexity: 12)
+
+        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
+        before do
+          allow(Gitlab::GitalyClient).to receive(:get_request_count).and_return(0)
+        end
+
+        it 'does not raise an error if calls_gitaly is false' do
+          expect { no_gitaly_field.send(:calls_gitaly_check) }.not_to raise_error
+        end
+
+        it 'raises an error if calls_gitaly: true appears' do
+          expect { gitaly_field.send(:calls_gitaly_check) }.to raise_error(/please add `calls_gitaly: true`/)
+        end
+      end
+
+      context 'if there is at least 1 Gitaly call' do
+        before do
+          allow(Gitlab::GitalyClient).to receive(:get_request_count).and_return(1)
+        end
+
+        it 'does not raise an error if calls_gitaly is true' do
+          expect { gitaly_field.send(:calls_gitaly_check) }.not_to raise_error
+        end
+
+        it 'raises an error if calls_gitaly is not decalared' do
+          expect { no_gitaly_field.send(:calls_gitaly_check) }.to raise_error(/please remove `calls_gitaly: true`/)
+        end
+      end
+    end
   end
 end