[GraphQL] verify batch load

tracer.rb

# frozen_string_literal: true

module Graphql
  # This module implements the functionality to track whether the batch loading is
  # required for a field to solve N+1 problem.
  #
  # If you want to check whether a particular field doesn't use batch loading in vain,
  # mark it with `verify_batch_load: true` option.
  #
  # By default, it logs the following information when detects a possible N+1 (thus, confirms that batch loading is required).
  #
  # To enable this tracer do the following:
  #   - extend the base field class, e.g. in BaseObject: `field_class.prepend(Graphql::BatchLoaderVerify::FieldExt)`
  #   - enabled tracer globally, i.e. in your schema: `tracer(Graphql::BatchLoaderVerify)`
  #   - OR per-query (e.g., if you want to add sampling): `execute(query, tracers: [Graphql::BatchLoaderVerify])`
  module BatchLoaderVerify
    module FieldExt
      def initialize(*args, verify_batch_load: false, **kwargs, &block)
        super(*args, **kwargs, &block)
        @verify_batch_load = verify_batch_load
      end

      def to_graphql
        super.tap do |new_field|
          new_field.instance_variable_set(:@verify_batch_load, true) if @verify_batch_load
        end
      end
    end

    class << self
      attr_accessor :reporter

      STORE_KEY = :__blv_store__

      def trace(event, metadata)
        if event == "execute_query"
          reset_store
          store[:query] = metadata[:query].operation_name
          yield.tap { report }
        elsif event == "execute_field"
          track(metadata[:path]) if metadata[:field].instance_variable_get(:@verify_batch_load)
          yield
        else
          yield
        end
      end

      def store
        Thread.current[STORE_KEY] ||= {
          paths: Hash.new { |h, k| h[k] = 0 },
        }
      end

      def reset_store
        Thread.current[STORE_KEY] = nil
      end

      def track(path)
        store[:paths][normalize_path(path)] += 1
      end

      def normalize_path(path)
        path.map { |segment| segment.is_a?(Numeric) ? "[]" : segment }.join(".").tap do |p|
          p.gsub!(/\.\[\]/, "[]")
        end
      end

      def report
        batch_loadable_paths = store[:paths].select { |_, v| v > 1 }
        return if batch_loadable_paths.empty?

        batch_loadable_paths.each do |path, _|
          reporter.call(store[:query], path)
        end
      end
    end

    self.reporter = lambda { |query, path|
      Rails.logger.info "[query: #{query}] batch loading is required for #{path}"
    }
  end
end