persisted_operation_document_cache.rb

require "bundler/inline"
require "logger"

gemfile do
  gem "graphql", "1.12.19"
  gem "graphql-pro", "1.20.1"
  gem "redis"
end

# Here's an example of caching parsed AST documents for persisted queries.
#
# The cache integrates at two points:
# - First, before calling `.execute(...)`, check the cache to see if there's already a document for the current `operation_id`.
#   If there is a document there, pass it along to `.execute`. (It's fine to pass `document: nil`, too, so you don't actually have to check it.)
# - Second, hooking into the OperationStore runtime instrumentation. That's where it fetches the operation from the database.
#   We need to add two behaviors there:
#     - Check if the query already has a `.document` -- if it does, skip the lookup in the operation store altogether.
#       (The library only checks for `query.query_string` -- this change should really be upstreamed, because `query.document` is good, too.)
#     - Then, if the query _did_ fetch a query string from the OperationStore, update the cache.
#       That way, subsequent runs with this operation ID will use the already-cached document.
#
# This cache is not thread-safe -- it's possible that during startup, some keys will override each other.
# I don't *think* this is a big deal, since we expect the same keys and values to end up in the cache anyways.
# The `GraphQL::Language::Nodes::Document` instances in the cache are immutable; no problem sharing those between threads.
# But if we need thread-safety, we could use a `Concurrent::Map` instead.
#
# This implementation will work fine if you don't need to limit the number of documents in the cache.
# If you *do* need to limit the number of documents, you could swap out the plain Hash
# stored in `MySchema.persisted_operation_document_cache` for a more sophisticated data structure.
# For example, a simple LRU cache: https://stackoverflow.com/a/16161783 Or pick a third-party library of your choice...
# (If you change the data structure, you might have to update the `..._cache[operation_id]` and `..._cache[operation_id] = ...` calls, too)
#
module InMemoryPersistedOperationCache
  def before_query(query)
    if query.document
      # We're good -- it was cached in memory
    else
      super
      if (operation_id = query.context[:operation_id]) && query.valid?
        query.schema.persisted_operation_document_cache[operation_id] = query.document
      end
    end
  end
end

class GraphQL::Pro::OperationStore::QueryInstrumentation
  prepend InMemoryPersistedOperationCache
end

class MySchema < GraphQL::Schema
  class Query < GraphQL::Schema::Object
    field :int, Integer, null: false

    def int
      100
    end
  end

  query(Query)
  use(GraphQL::Pro::OperationStore, redis: Redis.new(logger: Logger.new(STDOUT)))

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

MySchema.operation_store.upsert_client("client-1", "1234")
MySchema.operation_store.add(body: "query GetInts { i1: int i2: int }", client_name: "client-1", operation_alias: "query-1")

2.times do |i|
  puts "\n\nRun ##{i + 1}"
  operation_id = "client-1/query-1"
  # For example, this call would go in GraphqlController, before calling `.execute(...)`:
  document = MySchema.persisted_operation_document_cache[operation_id]
  p "Found cached #{document.class}"
  pp MySchema.execute(document: document, context: { operation_id: operation_id })
end

# Notice that the first run here logs the Redis traffic for reading the operation from the operation store. 
# But the second run has no Redis traffic, since it's read from memory.
#
# Run #1
# "Found cached NilClass"
# D, [2021-11-09T15:17:15.924745 #11989] DEBUG -- : [Redis] command=HMGET args="gql:opstore:clop:client-1:query-1" "digest" "is_archived"
# ... 
# D, [2021-11-09T15:17:15.926060 #11989] DEBUG -- : [Redis] call_time=0.09 ms
# #<GraphQL::Query::Result @query=... @to_h={"data"=>{"i1"=>100, "i2"=>100}}>
#
#
# Run #2
# "Found cached GraphQL::Language::Nodes::Document"
# #<GraphQL::Query::Result @query=... @to_h={"data"=>{"i1"=>100, "i2"=>100}}>