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`

6 files changed, 10 insertions, 11 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