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

@@ -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

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

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

lib/gitlab/graphql/calls_gitaly.rb

+# 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

lib/gitlab/graphql/calls_gitaly/instrumentation.rb

+# 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

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

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

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