GraphQL-Ruby in-memory query cache with OperationStore and a custom analyzer

example.rb

require "bundler/inline"
require "logger"

gemfile do
  gem "graphql", "2.2.10"
  gem "graphql-pro", "1.26.3"
  gem "redis"
end

# In-memory cache: populate the cache after queries are run.
# If the query was run from the operation store and was `.valid?`, add it to the cache.
# (At this point, `.valid?` reflects static validation _and_ analysis).
module InMemoryPersistedOperationCache
  def execute_multiplex(multiplex:)
    result = super
    multiplex.queries.each do |query|
      if (operation_id = query.context[:operation_id]) && query.valid?
        # Assign this document to the cache if it isn't already there:
        query.schema.persisted_operation_document_cache[operation_id] ||= query.document
      end
    end
    result
  end
end

# Here's that custom analyzer:
class PreventNestedCards < GraphQL::Analysis::AST::Analyzer
  def initialize(query)
    super
    @inside_players_selection = false
  end

  def analyze?
    query.validate
  end

  def on_enter_field(node, parent, visitor)
    if (field_defn = visitor.field_definition)
      if field_defn.graphql_name == "players"
        @inside_players_selection = true
      elsif @inside_players_selection && field_defn.graphql_name == "cards"
        raise GraphQL::AnalysisError.new(
          "Selecting `cards` within a list of `players` is not supported; Select cards for a specific player instead.",
          ast_node: node,
        )
      end
    end
  end

  def on_leave_field(node, parent, visitor)
    if @inside_players_selection && (field_defn = visitor.field_definition) && field_defn.graphql_name == "players"
      @inside_players_selection = false
    end
  end

  def result
    nil # nothing required here, since the error is raised above if anything goes wrong
  end
end

# Here's a minimal schema to demonstrate the behavior when valid/invalid:
class Schema < GraphQL::Schema
  class Card < GraphQL::Schema::Object
    field :id, String, null: false
  end

  class Player < GraphQL::Schema::Object
    field :id, String, null: false
    field :name, String, null: false
    field :cards, Card.connection_type, null: false
  end

  class Query < GraphQL::Schema::Object
    field :players, Player.connection_type
    field :player, Player do
      argument :id, String
    end
    # I added this to demonstrate that other selections are allowed:
    field :cards, Card.connection_type
  end

  query(Query)
  query_analyzer(PreventNestedCards)
  use GraphQL::Pro::OperationStore, redis: Redis.new
  trace_with InMemoryPersistedOperationCache

  class << self
    def persisted_operation_document_cache
      @persisted_operation_document_cache ||= {}
    end
  end
end

# Prepare the operation store:
Schema.operation_store.upsert_client("test", "test")

# Prepare a set of operations to test:
operations = [
  {
    body: "query op1 { players { nodes { cards { nodes { id } } } } }",
    alias: "op1",
    errors: 1,
  },
  {
    body: <<~GRAPHQL,
    query op2 {
      players { ...Players }
      __typename
    }
    fragment Players on PlayerConnection {
      nodes {
        cards {
          nodes { id }
        }
      }
    }
    GRAPHQL
    alias: "op2",
    errors: 1,
  },
  {
    body: <<~GRAPHQL,
    query op3 {
      players { nodes { ...Player } }
      __typename
    }
    fragment Player on Player {
      cards {
        nodes { id }
      }
    }
    GRAPHQL
    alias: "op3",
    errors: 1,
  },
  {
    body: <<~GRAPHQL,
    query op4 {
      players {
        nodes {
          ...on Player {
            cards { nodes { id } }
          }
        }
      }
      __typename
    }
    GRAPHQL
    alias: "op4",
    errors: 1,
  },
  {
    body: "query op5 { player(id: \"abc\") { cards { nodes { id } } } }",
    alias: "op5",
    errors: 0,
  },
  {
    body: "query op6 { players { nodes { name } } cards { nodes { id } } }",
    alias: "op6",
    errors: 0,
  },
]

# Load them into the operation store
operations.each do |op|
  Schema.operation_store.add(body: op[:body], client_name: "test", operation_alias: op[:alias])
end

# Then run each one, asserting that it's first uncached, then cached:
operations.each do |op|
  puts op[:body]
  operation_id = "test/#{op[:alias]}"
  results = 2.times.map do |i|
    # First, this will be null, but it should be present on the second run:
    document = Schema.persisted_operation_document_cache[operation_id]
    res = Schema.execute(
      query: nil,
      context: { operation_id: operation_id },
      document: document,
      root_value: {},
      validate: !document
    ).to_h

    # Make sure the result has the expected errors:
    errors_size = (res["errors"] || []).size
    if errors_size != op[:errors]
      raise "Unexpected Errors for #{op[:alias]} on #{i + 1} (expected #{op[:errors]}): #{res.inspect}"
    end

    # Return the cache status and result for comparing between runs:
    [!!document, res]
  end

  if results.first.last != results.last.last
    pp results
    raise "Results didn't match before and after"
  end

  if results.first.first != false || results.last.first != (op[:errors] == 0 ? true : false)
    pp results
    raise "Unexpected document cache behavior"
  end

  pp results.first.last
  puts "\n\n"
end

# Only two documents (the valid ones) were cached:
if Schema.persisted_operation_document_cache.size != 2
  raise "Wrong cache size (found: #{Schema.persisted_operation_document_cache.size})"
else
  puts "Found two cached documents"
end