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`

2 files changed, 12 insertions, 10 deletions

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