From 8b809837f44bbebdac65eebadb09a45fb60c6a65 Mon Sep 17 00:00:00 2001
From: charlieablett <cablett@gitlab.com>
Date: Sun, 16 Jun 2019 22:12:56 +1200
Subject: 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
---
 app/graphql/types/base_field.rb       | 23 +++++++++-
 app/graphql/types/project_type.rb     |  2 +-
 app/graphql/types/tree/tree_type.rb   |  4 +-
 spec/graphql/types/base_field_spec.rb | 81 +++++++++++++++++++++++++++++++++++
 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
-- 
cgit v1.2.3


From c99c30fdd6f3adf4fb29aad4b10e265be69d2d67 Mon Sep 17 00:00:00 2001
From: charlieablett <cablett@gitlab.com>
Date: Wed, 19 Jun 2019 14:37:54 +0200
Subject: Remove potentially noisy warning

- If Gitaly calls are missing, it could be due to a conditional and
may just become noise
---
 app/graphql/types/base_field.rb       | 2 --
 spec/graphql/types/base_field_spec.rb | 4 ----
 2 files changed, 6 deletions(-)

diff --git a/app/graphql/types/base_field.rb b/app/graphql/types/base_field.rb
index ee23146f711..95db116d6f9 100644
--- a/app/graphql/types/base_field.rb
+++ b/app/graphql/types/base_field.rb
@@ -60,8 +60,6 @@ module Types
       # 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)
diff --git a/spec/graphql/types/base_field_spec.rb b/spec/graphql/types/base_field_spec.rb
index ebdfa3eaf4d..d7360872508 100644
--- a/spec/graphql/types/base_field_spec.rb
+++ b/spec/graphql/types/base_field_spec.rb
@@ -127,10 +127,6 @@ describe Types::BaseField do
         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
-- 
cgit v1.2.3


From f4890d90782ad42a802b89c2a17c83bf9fb9d123 Mon Sep 17 00:00:00 2001
From: charlieablett <cablett@gitlab.com>
Date: Fri, 21 Jun 2019 16:20:00 +0200
Subject: Alert if `calls_gitaly` declaration missing

- Move `calls_gitaly_check` to public
- Add instrumentation for flagging missing CallsGitaly declarations
- Wrap resolver proc in before-and-after Gitaly counts to get the net
Gitaly call count for the resolver.
---
 app/graphql/gitlab_schema.rb                       |  1 +
 app/graphql/types/base_field.rb                    | 22 ++++++++-------
 lib/gitlab/graphql/authorize.rb                    |  2 +-
 lib/gitlab/graphql/calls_gitaly.rb                 | 15 +++++++++++
 lib/gitlab/graphql/calls_gitaly/instrumentation.rb | 30 +++++++++++++++++++++
 spec/graphql/gitlab_schema_spec.rb                 |  4 +++
 spec/graphql/types/base_field_spec.rb              | 18 ++++---------
 spec/requests/api/graphql_spec.rb                  | 31 ++++++++++++++++++++++
 8 files changed, 99 insertions(+), 24 deletions(-)
 create mode 100644 lib/gitlab/graphql/calls_gitaly.rb
 create mode 100644 lib/gitlab/graphql/calls_gitaly/instrumentation.rb

diff --git a/app/graphql/gitlab_schema.rb b/app/graphql/gitlab_schema.rb
index 5615909c4ec..152ebb930e2 100644
--- a/app/graphql/gitlab_schema.rb
+++ b/app/graphql/gitlab_schema.rb
@@ -13,6 +13,7 @@ class GitlabSchema < GraphQL::Schema
   use BatchLoader::GraphQL
   use Gitlab::Graphql::Authorize
   use Gitlab::Graphql::Present
+  use Gitlab::Graphql::CallsGitaly
   use Gitlab::Graphql::Connections
   use Gitlab::Graphql::GenericTracing
 
diff --git a/app/graphql/types/base_field.rb b/app/graphql/types/base_field.rb
index 95db116d6f9..64bc7e6474f 100644
--- a/app/graphql/types/base_field.rb
+++ b/app/graphql/types/base_field.rb
@@ -21,6 +21,18 @@ module Types
       complexity
     end
 
+    def calls_gitaly_check(calls)
+      return if @calls_gitaly
+
+      # Will inform you if :calls_gitaly should be true or false based on the number of Gitaly calls
+      # involved with the request.
+      if calls > 0
+        raise "Gitaly is called for field '#{name}' - please add `calls_gitaly: true` to the field declaration"
+      end
+    rescue => e
+      Gitlab::Sentry.track_exception(e)
+    end
+
     private
 
     def field_complexity(resolver_class)
@@ -54,15 +66,5 @@ 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"
-      end
-    rescue => e
-      Gitlab::Sentry.track_exception(e)
-    end
   end
 end
diff --git a/lib/gitlab/graphql/authorize.rb b/lib/gitlab/graphql/authorize.rb
index f8d0208e275..e83b567308b 100644
--- a/lib/gitlab/graphql/authorize.rb
+++ b/lib/gitlab/graphql/authorize.rb
@@ -8,7 +8,7 @@ module Gitlab
       extend ActiveSupport::Concern
 
       def self.use(schema_definition)
-        schema_definition.instrument(:field, Instrumentation.new, after_built_ins: true)
+        schema_definition.instrument(:field, Gitlab::Graphql::Authorize::Instrumentation.new, after_built_ins: true)
       end
     end
   end
diff --git a/lib/gitlab/graphql/calls_gitaly.rb b/lib/gitlab/graphql/calls_gitaly.rb
new file mode 100644
index 00000000000..f75941e269f
--- /dev/null
+++ b/lib/gitlab/graphql/calls_gitaly.rb
@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+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.
+    module CallsGitaly
+      extend ActiveSupport::Concern
+
+      def self.use(schema_definition)
+        schema_definition.instrument(:field, Gitlab::Graphql::CallsGitaly::Instrumentation.new, after_built_ins: true)
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/graphql/calls_gitaly/instrumentation.rb b/lib/gitlab/graphql/calls_gitaly/instrumentation.rb
new file mode 100644
index 00000000000..ca54e12c049
--- /dev/null
+++ b/lib/gitlab/graphql/calls_gitaly/instrumentation.rb
@@ -0,0 +1,30 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Graphql
+    module CallsGitaly
+      class Instrumentation
+        # 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)
+
+          old_resolver_proc = field.resolve_proc
+          wrapped_proc = gitaly_wrapped_resolve(old_resolver_proc, type_object)
+          field.redefine { resolve(wrapped_proc) }
+        end
+
+        def gitaly_wrapped_resolve(old_resolver_proc, type_object)
+          proc do |parent_typed_object, args, ctx|
+            previous_gitaly_call_count = Gitlab::GitalyClient.get_request_count
+
+            old_resolver_proc.call(parent_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)
+          end
+        end
+      end
+    end
+  end
+end
diff --git a/spec/graphql/gitlab_schema_spec.rb b/spec/graphql/gitlab_schema_spec.rb
index d36e428a8ee..93b86b9b812 100644
--- a/spec/graphql/gitlab_schema_spec.rb
+++ b/spec/graphql/gitlab_schema_spec.rb
@@ -21,6 +21,10 @@ describe GitlabSchema do
     expect(field_instrumenters).to include(instance_of(::Gitlab::Graphql::Present::Instrumentation))
   end
 
+  it 'enables using gitaly call checker' do
+    expect(field_instrumenters).to include(instance_of(::Gitlab::Graphql::CallsGitaly::Instrumentation))
+  end
+
   it 'has the base mutation' do
     expect(described_class.mutation).to eq(::Types::MutationType.to_graphql)
   end
diff --git a/spec/graphql/types/base_field_spec.rb b/spec/graphql/types/base_field_spec.rb
index d7360872508..0be83ea60c4 100644
--- a/spec/graphql/types/base_field_spec.rb
+++ b/spec/graphql/types/base_field_spec.rb
@@ -106,26 +106,18 @@ describe Types::BaseField do
       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`/)
+          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
-        before do
-          allow(Gitlab::GitalyClient).to receive(:get_request_count).and_return(1)
+        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 'does not raise an error if calls_gitaly is true' do
-          expect { gitaly_field.send(:calls_gitaly_check) }.not_to raise_error
+        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
diff --git a/spec/requests/api/graphql_spec.rb b/spec/requests/api/graphql_spec.rb
index 656d6f8b50b..d78b17827a6 100644
--- a/spec/requests/api/graphql_spec.rb
+++ b/spec/requests/api/graphql_spec.rb
@@ -131,4 +131,35 @@ describe 'GraphQL' do
       end
     end
   end
+
+  describe 'testing for Gitaly calls' do
+    let(:project) { create(:project, :repository) }
+    let(:user) { create(:user) }
+
+    let(:query) do
+      graphql_query_for('project', { 'fullPath' => project.full_path }, %w(forksCount))
+    end
+
+    before do
+      project.add_developer(user)
+    end
+
+    it_behaves_like 'a working graphql query' do
+      before do
+        post_graphql(query, current_user: user)
+      end
+    end
+
+    context 'when Gitaly is called' do
+      before do
+        allow(Gitlab::GitalyClient).to receive(:get_request_count).and_return(1, 2)
+      end
+
+      it "logs a warning that the 'calls_gitaly' field declaration is missing" do
+        expect(Gitlab::Sentry).to receive(:track_exception).once
+
+        post_graphql(query, current_user: user)
+      end
+    end
+  end
 end
-- 
cgit v1.2.3


From a11fe5de4408595cc8b2b091cbbb76e423c98f34 Mon Sep 17 00:00:00 2001
From: charlieablett <cablett@gitlab.com>
Date: Wed, 26 Jun 2019 22:42:25 +1200
Subject: Wrap proc properly in gitaly call counts

- Add `calls_gitaly: true` to some fields missing (hey, it works!)
- Clarify proc wrapping
- Add kwargs argument to `mount_mutation`
---
 app/graphql/types/base_field.rb                     |  5 ++---
 app/graphql/types/merge_request_type.rb             |  2 +-
 app/graphql/types/mutation_type.rb                  |  2 +-
 app/graphql/types/permission_types/merge_request.rb |  4 ++--
 app/graphql/types/project_type.rb                   |  2 +-
 app/graphql/types/repository_type.rb                |  6 +++---
 lib/gitlab/graphql/calls_gitaly/instrumentation.rb  | 17 +++++++++--------
 lib/gitlab/graphql/mount_mutation.rb                |  5 +++--
 8 files changed, 22 insertions(+), 21 deletions(-)

diff --git a/app/graphql/types/base_field.rb b/app/graphql/types/base_field.rb
index 64bc7e6474f..42c7eb6b485 100644
--- a/app/graphql/types/base_field.rb
+++ b/app/graphql/types/base_field.rb
@@ -23,12 +23,11 @@ module Types
 
     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.
-      if calls > 0
-        raise "Gitaly is called for field '#{name}' - please add `calls_gitaly: true` to the field declaration"
-      end
+      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)
     end
diff --git a/app/graphql/types/merge_request_type.rb b/app/graphql/types/merge_request_type.rb
index 577ccd48ef8..6734d4761c2 100644
--- a/app/graphql/types/merge_request_type.rb
+++ b/app/graphql/types/merge_request_type.rb
@@ -43,7 +43,7 @@ module Types
     field :allow_collaboration, GraphQL::BOOLEAN_TYPE, null: true
     field :should_be_rebased, GraphQL::BOOLEAN_TYPE, method: :should_be_rebased?, null: false
     field :rebase_commit_sha, GraphQL::STRING_TYPE, null: true
-    field :rebase_in_progress, GraphQL::BOOLEAN_TYPE, method: :rebase_in_progress?, null: false
+    field :rebase_in_progress, GraphQL::BOOLEAN_TYPE, method: :rebase_in_progress?, null: false, calls_gitaly: true
     field :merge_commit_message, GraphQL::STRING_TYPE, method: :default_merge_commit_message, null: true, deprecation_reason: "Renamed to defaultMergeCommitMessage"
     field :default_merge_commit_message, GraphQL::STRING_TYPE, null: true
     field :merge_ongoing, GraphQL::BOOLEAN_TYPE, method: :merge_ongoing?, null: false
diff --git a/app/graphql/types/mutation_type.rb b/app/graphql/types/mutation_type.rb
index 6ef1d816b7c..bc5fb709522 100644
--- a/app/graphql/types/mutation_type.rb
+++ b/app/graphql/types/mutation_type.rb
@@ -9,6 +9,6 @@ module Types
     mount_mutation Mutations::AwardEmojis::Add
     mount_mutation Mutations::AwardEmojis::Remove
     mount_mutation Mutations::AwardEmojis::Toggle
-    mount_mutation Mutations::MergeRequests::SetWip
+    mount_mutation Mutations::MergeRequests::SetWip, calls_gitaly: true
   end
 end
diff --git a/app/graphql/types/permission_types/merge_request.rb b/app/graphql/types/permission_types/merge_request.rb
index 13995d3ea8f..d877fc177d2 100644
--- a/app/graphql/types/permission_types/merge_request.rb
+++ b/app/graphql/types/permission_types/merge_request.rb
@@ -10,8 +10,8 @@ module Types
       abilities :read_merge_request, :admin_merge_request,
                 :update_merge_request, :create_note
 
-      permission_field :push_to_source_branch, method: :can_push_to_source_branch?
-      permission_field :remove_source_branch, method: :can_remove_source_branch?
+      permission_field :push_to_source_branch, method: :can_push_to_source_branch?, calls_gitaly: true
+      permission_field :remove_source_branch, method: :can_remove_source_branch?, calls_gitaly: true
       permission_field :cherry_pick_on_current_merge_request, method: :can_cherry_pick_on_current_merge_request?
       permission_field :revert_on_current_merge_request, method: :can_revert_on_current_merge_request?
     end
diff --git a/app/graphql/types/project_type.rb b/app/graphql/types/project_type.rb
index 87d5351f80f..13be71c26ee 100644
--- a/app/graphql/types/project_type.rb
+++ b/app/graphql/types/project_type.rb
@@ -40,7 +40,7 @@ module Types
     field :lfs_enabled, GraphQL::BOOLEAN_TYPE, null: true
     field :merge_requests_ff_only_enabled, GraphQL::BOOLEAN_TYPE, null: true
 
-    field :avatar_url, GraphQL::STRING_TYPE, null: true, resolve: -> (project, args, ctx) do
+    field :avatar_url, GraphQL::STRING_TYPE, null: true, calls_gitaly: true, resolve: -> (project, args, ctx) do
       project.avatar_url(only_path: false)
     end
 
diff --git a/app/graphql/types/repository_type.rb b/app/graphql/types/repository_type.rb
index 5987467e1ea..b024eca61fc 100644
--- a/app/graphql/types/repository_type.rb
+++ b/app/graphql/types/repository_type.rb
@@ -6,9 +6,9 @@ module Types
 
     authorize :download_code
 
-    field :root_ref, GraphQL::STRING_TYPE, null: true
-    field :empty, GraphQL::BOOLEAN_TYPE, null: false, method: :empty?
+    field :root_ref, GraphQL::STRING_TYPE, null: true, calls_gitaly: true
+    field :empty, GraphQL::BOOLEAN_TYPE, null: false, method: :empty?, calls_gitaly: true
     field :exists, GraphQL::BOOLEAN_TYPE, null: false, method: :exists?
-    field :tree, Types::Tree::TreeType, null: true, resolver: Resolvers::TreeResolver
+    field :tree, Types::Tree::TreeType, null: true, resolver: Resolvers::TreeResolver, calls_gitaly: true
   end
 end
diff --git a/lib/gitlab/graphql/calls_gitaly/instrumentation.rb b/lib/gitlab/graphql/calls_gitaly/instrumentation.rb
index ca54e12c049..08e98028755 100644
--- a/lib/gitlab/graphql/calls_gitaly/instrumentation.rb
+++ b/lib/gitlab/graphql/calls_gitaly/instrumentation.rb
@@ -10,18 +10,19 @@ module Gitlab
           return field unless type_object && type_object.respond_to?(:calls_gitaly_check)
 
           old_resolver_proc = field.resolve_proc
-          wrapped_proc = gitaly_wrapped_resolve(old_resolver_proc, type_object)
-          field.redefine { resolve(wrapped_proc) }
-        end
 
-        def gitaly_wrapped_resolve(old_resolver_proc, type_object)
-          proc do |parent_typed_object, args, ctx|
+          gitaly_wrapped_resolve = -> (typed_object, args, ctx) do
             previous_gitaly_call_count = Gitlab::GitalyClient.get_request_count
-
-            old_resolver_proc.call(parent_typed_object, args, ctx)
-
+            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)
+            result
+          rescue => e
+            ap "#{e.message}"
+          end
+
+          field.redefine do
+            resolve(gitaly_wrapped_resolve)
           end
         end
       end
diff --git a/lib/gitlab/graphql/mount_mutation.rb b/lib/gitlab/graphql/mount_mutation.rb
index 9048967d4e1..b10e963170a 100644
--- a/lib/gitlab/graphql/mount_mutation.rb
+++ b/lib/gitlab/graphql/mount_mutation.rb
@@ -6,11 +6,12 @@ module Gitlab
       extend ActiveSupport::Concern
 
       class_methods do
-        def mount_mutation(mutation_class)
+        def mount_mutation(mutation_class, **custom_kwargs)
           # Using an underscored field name symbol will make `graphql-ruby`
           # standardize the field name
           field mutation_class.graphql_name.underscore.to_sym,
-                mutation: mutation_class
+                mutation: mutation_class,
+                **custom_kwargs
         end
       end
     end
-- 
cgit v1.2.3


From cf1b0d10bcdde69f05695a2e9a0d380c6badb6d1 Mon Sep 17 00:00:00 2001
From: charlieablett <cablett@gitlab.com>
Date: Fri, 28 Jun 2019 13:24:47 +1200
Subject: 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
---
 app/graphql/types/base_field.rb                    | 16 +++---------
 lib/gitlab/graphql/calls_gitaly/instrumentation.rb | 16 +++++++++---
 spec/graphql/types/base_field_spec.rb              | 27 +++-----------------
 .../graphql/calls_gitaly/instrumentation_spec.rb   | 29 ++++++++++++++++++++++
 spec/requests/api/graphql_spec.rb                  |  2 +-
 5 files changed, 50 insertions(+), 40 deletions(-)
 create mode 100644 spec/lib/gitlab/graphql/calls_gitaly/instrumentation_spec.rb

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
-- 
cgit v1.2.3


From 675c9b9f6bec35f1e6988a42c4fa6a6f8331d14f Mon Sep 17 00:00:00 2001
From: charlieablett <cablett@gitlab.com>
Date: Tue, 2 Jul 2019 17:32:44 +1200
Subject: Address reviewer comments

- Remove Gitaly call check for fields that have a constant complexity
declared
- Add associated test
---
 app/graphql/types/base_field.rb                    |  5 +++++
 app/graphql/types/tree/tree_type.rb                |  2 +-
 lib/gitlab/graphql/calls_gitaly.rb                 |  4 ++--
 lib/gitlab/graphql/calls_gitaly/instrumentation.rb |  7 +++---
 spec/graphql/types/base_field_spec.rb              | 15 ++++++++++---
 .../graphql/calls_gitaly/instrumentation_spec.rb   | 26 +++++++++-------------
 6 files changed, 34 insertions(+), 25 deletions(-)

diff --git a/app/graphql/types/base_field.rb b/app/graphql/types/base_field.rb
index 6b377e88e16..efeee4a7a4d 100644
--- a/app/graphql/types/base_field.rb
+++ b/app/graphql/types/base_field.rb
@@ -8,6 +8,7 @@ module Types
 
     def initialize(*args, **kwargs, &block)
       @calls_gitaly = !!kwargs.delete(:calls_gitaly)
+      @constant_complexity = !!kwargs[:complexity]
       kwargs[:complexity] ||= field_complexity(kwargs[:resolver_class])
 
       super(*args, **kwargs, &block)
@@ -23,6 +24,10 @@ module Types
       @calls_gitaly
     end
 
+    def constant_complexity?
+      @constant_complexity
+    end
+
     private
 
     def field_complexity(resolver_class)
diff --git a/app/graphql/types/tree/tree_type.rb b/app/graphql/types/tree/tree_type.rb
index 2023abc13f9..fbdc1597461 100644
--- a/app/graphql/types/tree/tree_type.rb
+++ b/app/graphql/types/tree/tree_type.rb
@@ -7,7 +7,7 @@ module Types
       graphql_name 'Tree'
 
       # Complexity 10 as it triggers a Gitaly call on each render
-      field :last_commit, Types::CommitType, null: true, complexity: 10, resolve: -> (tree, args, ctx) do
+      field :last_commit, Types::CommitType, null: true, complexity: 10, calls_gitaly: true, resolve: -> (tree, args, ctx) do
         tree.repository.last_commit_for_path(tree.sha, tree.path)
       end
 
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
diff --git a/spec/graphql/types/base_field_spec.rb b/spec/graphql/types/base_field_spec.rb
index 10913e530cf..77ef8933717 100644
--- a/spec/graphql/types/base_field_spec.rb
+++ b/spec/graphql/types/base_field_spec.rb
@@ -96,10 +96,19 @@ describe Types::BaseField do
         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)
+      context 'with declared constant complexity value' do
+        it 'has complexity set to that constant' do
+          field = described_class.new(name: 'test', type: GraphQL::STRING_TYPE, null: true, complexity: 12)
 
-        expect(field.to_graphql.complexity).to eq 12
+          expect(field.to_graphql.complexity).to eq 12
+        end
+
+        it 'does not raise an error even with Gitaly calls' do
+          allow(Gitlab::GitalyClient).to receive(:get_request_count).and_return([0, 1])
+          field = described_class.new(name: 'test', type: GraphQL::STRING_TYPE, null: true, complexity: 12)
+
+          expect(field.to_graphql.complexity).to eq 12
+        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
index 92d200bfd4e..d93ce464a92 100644
--- a/spec/lib/gitlab/graphql/calls_gitaly/instrumentation_spec.rb
+++ b/spec/lib/gitlab/graphql/calls_gitaly/instrumentation_spec.rb
@@ -4,25 +4,19 @@ 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) }
+  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
+    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
+    context 'if there is at least 1 Gitaly call' do
+      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(/specify a constant complexity or add `calls_gitaly: true`/)
       end
     end
   end
-- 
cgit v1.2.3