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.

3 files changed, 46 insertions, 1 deletions

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