query_stash 0.1.0

data/lib/query_stash/active_record.rb

@@ -0,0 +1,188 @@
+# frozen_string_literal: true
+
+require "digest"
+require_relative "base"
+require_relative "registry/sql"
+
+module QueryStash
+  class ActiveRecord < Base
+    # What the patches ask of Rails' own objects.
+    #
+    # @rbs!
+    #   # An ActiveRecord adapter, up to Rails 8.1.
+    #   interface _InternalExecQueryAdapter
+    #     def internal_exec_query: (String sql, ?String name, ?Array[untyped] binds, **untyped) -> ::ActiveRecord::Result
+    #   end
+    #
+    #   # An ActiveRecord adapter, from Rails 8.2.
+    #   interface _IntentAdapter
+    #     def execute_intent: (_QueryIntent intent) -> untyped
+    #
+    #     def cast_result: (untyped raw_result) -> ::ActiveRecord::Result
+    #   end
+    #
+    #   # A query, from Rails 8.2.
+    #   interface _QueryIntent
+    #     def raw_sql: () -> String
+    #
+    #     def has_binds?: () -> bool
+    #
+    #     def deliver_result: (::ActiveRecord::Result result) -> untyped
+    #
+    #     def cast_result: () -> ::ActiveRecord::Result
+    #   end
+
+    CACHEABLE_SQL_REGEX = /\ASELECT/i #: Regexp
+    private_constant :CACHEABLE_SQL_REGEX
+
+    # Enables the stash for the duration of each request.
+    class Middleware
+      # @rbs @app: _RackApp
+      # @rbs @application: untyped
+
+      #: (_RackApp app, untyped application) -> void
+      def initialize(app, application)
+        @app = app
+        @application = application
+      end
+
+      #: (Hash[String, untyped] env) -> untyped
+      def call(env)
+        QueryStash::ActiveRecord[@application].enable! { @app.call(env) }
+      end
+    end
+
+    # Patch wraps queries run through `internal_exec_query` (rails 7.2 to 8.1).
+    # See `QueryStash::ActiveRecord.implementation`.
+    module InternalExecQuery
+      # @rbs module-self _InternalExecQueryAdapter
+      module AdapterPatch
+        # @rbs override
+        def internal_exec_query(sql, name = "SQL", binds = [], ...) # steep:ignore UnannotatedEmptyCollection
+          stash = QueryStash::ActiveRecord[Rails.application] #: QueryStash::ActiveRecord & Stash
+          stash.wrap(sql, binds) { super(sql, name, binds, ...) }
+        end
+      end
+
+      # @rbs module-self QueryStash::ActiveRecord
+      module Stash
+        # Runs the query through the stash, or just yields when the stash is off or the query isn't
+        # stashable.
+        #: (String sql, Array[untyped]? binds) { () -> ::ActiveRecord::Result } -> ::ActiveRecord::Result
+        def wrap(sql, binds)
+          rule = rule_for(sql, has_binds: binds&.any?)
+          return yield unless rule
+
+          cached_result = read(rule, sql)
+          return cached_result unless cached_result.nil?
+
+          write(rule, yield, sql)
+        end
+      end
+    end
+
+    # Rails 8.2 runs queries as intents: the adapter's `execute_intent` runs one and hands its raw result
+    # back to the intent, which turns it into an ActiveRecord::Result in `QueryIntent#cast_result`.
+    module ExecuteIntent
+      # @rbs module-self _IntentAdapter
+      module AdapterPatch
+        #: (_QueryIntent intent) -> untyped
+        def execute_intent(intent)
+          stash = QueryStash::ActiveRecord[Rails.application] #: QueryStash::ActiveRecord & Stash
+          cached_result = stash.read_intent(intent)
+          return super if cached_result.nil?
+
+          intent.deliver_result(cached_result)
+        end
+
+        private
+
+        # A stash hit is delivered as the ActiveRecord::Result it already is.
+        #: (untyped raw_result) -> ::ActiveRecord::Result
+        def cast_result(raw_result)
+          raw_result.is_a?(::ActiveRecord::Result) ? raw_result : super
+        end
+      end
+
+      # @rbs module-self _QueryIntent
+      module IntentPatch
+        #: () -> ::ActiveRecord::Result
+        def cast_result
+          stash = QueryStash::ActiveRecord[Rails.application] #: QueryStash::ActiveRecord & Stash
+          super.tap { |result| stash.write_intent(self, result) }
+        end
+      end
+
+      # @rbs module-self QueryStash::ActiveRecord
+      module Stash
+        # @rbs @intent_rules: ObjectSpace::WeakMap[_QueryIntent, Registry::Rule]
+
+        # The rule each intent that missed the stash is to be stored under, until its result is cast.
+        #: (QueryStash::ActiveRecord stash) -> void
+        def self.extended(stash)
+          stash.instance_variable_set(:@intent_rules, ObjectSpace::WeakMap.new)
+        end
+
+        # The stashed result of an intent (see AdapterPatch), or nil: when there is none, the intent is
+        # remembered for `write_intent` to store its result.
+        #: (_QueryIntent intent) -> ::ActiveRecord::Result?
+        def read_intent(intent)
+          rule = rule_for(intent.raw_sql, has_binds: intent.has_binds?)
+          return unless rule
+
+          read(rule, intent.raw_sql).tap { |cached_result| @intent_rules[intent] = rule if cached_result.nil? }
+        end
+
+        # Stores the result of an intent `read_intent` found nothing for.
+        #: (_QueryIntent intent, ::ActiveRecord::Result result) -> untyped
+        def write_intent(intent, result)
+          rule = @intent_rules[intent]
+          return unless rule
+
+          @intent_rules.delete(intent)
+          write(rule, result, intent.raw_sql)
+        end
+      end
+    end
+
+    # The implementation for the Rails in use: only to be asked once ActiveRecord is loaded.
+    #: () -> (singleton(InternalExecQuery) | singleton(ExecuteIntent))
+    def self.implementation
+      if ::ActiveRecord::ConnectionAdapters.const_defined?(:QueryIntent)
+        # Rails 8.2 or higher
+        ExecuteIntent
+      else
+        # Rails 8.1 or lower
+        InternalExecQuery
+      end
+    end
+
+    # The rule the query is stashed under, or nil when it isn't to be stashed: the stash is off, it has
+    # binds (not supported), it is not a SELECT, or no rule matches it.
+    #: (String sql, ?has_binds: bool?) -> Registry::Rule?
+    def rule_for(sql, has_binds: false)
+      return unless stashing?
+      return if has_binds
+      return unless CACHEABLE_SQL_REGEX.match?(sql)
+
+      rule = registry.match_rule(sql)
+      rule if rule && (rule.enabled.nil? || rule.enabled.call)
+    end
+
+    #: (String sql) -> String
+    def cache_key_for(sql)
+      # Note that the SQL is already normalized by the adapter here.
+      #
+      # Also note that SQL comments are applied AFTER this since
+      # this module is prepended.
+      "v1:#{::ActiveRecord::Base.current_shard}:#{Digest::SHA256.hexdigest(sql)}"
+    end
+
+    private
+
+    #: () -> Registry::Sql
+    def build_registry
+      Registry::Sql.new
+    end
+  end
+end

data/lib/query_stash/base.rb

@@ -0,0 +1,275 @@
+# frozen_string_literal: true
+
+require "concurrent/map"
+require_relative "registry"
+
+module QueryStash
+  class Base
+    # Where `enable!` and `disable!` keep what they set, when no `isolated_state` is
+    # given: a hash in a thread variable, so every fiber of a thread sees the same values and
+    # other threads see none.
+    #
+    # An isolated state is any object that responds to `[]`, `[]=`, `key?` and `delete`, like
+    # `ActiveSupport::IsolatedExecutionState`.
+    class ThreadVariableState
+      THREAD_VARIABLE = :query_stash_state #: Symbol
+
+      #: (untyped key) -> untyped
+      def [](key)
+        values[key]
+      end
+
+      #: (untyped key, untyped value) -> untyped
+      def []=(key, value)
+        values[key] = value
+      end
+
+      #: (untyped key) -> bool
+      def key?(key)
+        values.key?(key)
+      end
+
+      #: (untyped key) -> untyped
+      def delete(key)
+        values.delete(key)
+      end
+
+      private
+
+      #: () -> Hash[untyped, untyped]
+      def values
+        Thread.current.thread_variable_get(THREAD_VARIABLE) ||
+          Thread.current.thread_variable_set(THREAD_VARIABLE, {})
+      end
+    end
+
+    # Used when no `instrumenter` is given: runs the block and tells no one.
+    class NullInstrumenter
+      #: (String _name, ?Hash[Symbol, untyped] payload) ?{ (Hash[Symbol, untyped]) -> untyped } -> untyped
+      def instrument(_name, payload = {})
+        yield payload if block_given?
+      end
+    end
+
+    # Used when no `error_reporter` is given.
+    class WarnErrorReporter
+      #: (Exception error, ?handled: bool) -> void
+      def report(error, handled: true) # rubocop:disable Lint/UnusedMethodArgument
+        warn("QueryStash: #{error.class}: #{error.message}")
+      end
+    end
+
+    class << self
+      # @rbs self.@instances: Concurrent::Map[untyped, untyped]
+
+      #: (Class subclass) -> void
+      def inherited(subclass)
+        super
+        subclass.instance_variable_set(:@instances, Concurrent::Map.new)
+      end
+
+      #: (untyped app) -> instance
+      def [](app)
+        @instances.fetch(app)
+      end
+
+      #: (untyped app) -> bool
+      def key?(app)
+        !@instances.fetch(app, nil).nil?
+      end
+
+      # Makes `stash` the instance `app` uses.
+      #: (untyped app, instance stash) -> instance
+      def set(app, stash)
+        @instances[app] = stash
+      end
+    end
+
+    # @rbs @enabled: bool
+    # @rbs @isolated_state: _IsolatedState
+    # @rbs @instrumenter: _Instrumenter
+    # @rbs @error_reporter: _ErrorReporter
+    # @rbs @raise_on_cache_error: bool
+    # @rbs @registry: Registry
+    # @rbs @cache: _CacheStore
+    # @rbs @adapter: Symbol
+
+    attr_reader :registry #: Registry
+
+    attr_reader :cache #: _CacheStore
+
+    attr_reader :enabled #: bool
+
+    # What each stash is given:
+    #
+    # - `enabled`: whether queries are stashed outside of `enable!` and `disable!`.
+    # - `isolated_state`: where those keep what they set. By default, a thread variable.
+    # - `instrumenter`: told about what the stash does. By default, no one.
+    # - `error_reporter`: told about the errors it handles. By default, a warning on stderr.
+    # - `raise_on_cache_error`: raise the errors of the cache store, instead of reporting them.
+    # - `cache_store`: where results are kept, any object that responds to `read(key)`,
+    #   `write(key, value, expires_in:)` and `clear`, like an ActiveSupport cache store.
+    #
+    # @rbs cache_store: _CacheStore
+    # @rbs enabled: bool
+    # @rbs isolated_state: _IsolatedState?
+    # @rbs instrumenter: _Instrumenter?
+    # @rbs error_reporter: _ErrorReporter?
+    # @rbs raise_on_cache_error: bool
+    # @rbs return: void
+    def initialize(
+      cache_store:,
+      enabled: true,
+      isolated_state: nil,
+      instrumenter: nil,
+      error_reporter: nil,
+      raise_on_cache_error: false
+    )
+      @enabled = enabled
+      @isolated_state = isolated_state || ThreadVariableState.new
+      @instrumenter = instrumenter || NullInstrumenter.new
+      @error_reporter = error_reporter || WarnErrorReporter.new
+      @raise_on_cache_error = raise_on_cache_error
+      @registry = build_registry
+      @cache = cache_store
+      @adapter = self.class.name.to_s.split("::").last.to_s.gsub(/(?<=[a-z])(?=[A-Z])/, "_").downcase.to_sym
+    end
+
+    #: () { (Registry registry) -> void } -> void
+    def configure
+      yield registry
+    end
+
+    # Turn the stash on or off for the block. Outside of either, it follows `enabled`.
+    #: [T] () { () -> T } -> T
+    def enable!(&block)
+      isolate(true, &block)
+    end
+
+    #: [T] () { () -> T } -> T
+    def disable!(&block)
+      isolate(false, &block)
+    end
+
+    #: (Registry::Rule rule) -> _CacheStore
+    def cache_for(rule)
+      rule.store || cache
+    end
+
+    # ---------------- What an adapter's stash provides ----------------
+    #
+    # An adapter patches whatever runs its queries so that it asks its stash, in this order:
+    #
+    # 1. `rule_for(*args)`: is this query stashed, and under which rule? When it isn't, the query just runs.
+    # 2. `read(rule, *args)`: is its result in the stash? When it is, that is the result.
+    # 3. `write(rule, result, *args)`: otherwise it runs, and its result goes in the stash.
+    #
+    # `*args` is whatever the adapter has that identifies a query (the SQL, a driver's view, ...). It is the
+    # same for `read`, `write` and `cache_key_for`. `rule_for` may be given more, for what only decides
+    # whether to stash.
+    #
+    # A stash gets `read` and `write` from here, which keep the reading and writing of the cache store, the
+    # instrumentation and the handling of errors in one place. It has to define `rule_for` and
+    # `cache_key_for`, and can override the others.
+
+    # The rule of the registry the query is to be stashed under, or nil when it isn't to be: the stash is off
+    # (see `stashing?`), the query is not one that can be, or no rule matches it.
+    #: (*untyped) -> Registry::Rule?
+    def rule_for(*)
+      raise NotImplementedError, "#{self.class} must define #rule_for"
+    end
+
+    # The key the query's result is stored under.
+    #: (*untyped) -> String
+    def cache_key_for(*)
+      raise NotImplementedError, "#{self.class} must define #cache_key_for"
+    end
+
+    # The result stored for the query, or nil when there is none (or the cache store failed: that is
+    # reported, or raised when `raise_on_cache_error` is set). Emits `cache_read` and, on a hit,
+    # `cache_hit`.
+    #: (Registry::Rule rule, *untyped args) -> untyped
+    def read(rule, *args)
+      cache_key = cache_key_for(*args)
+      cache_store = cache_for(rule)
+
+      rescue_cache_error do
+        cached_result = @instrumenter.instrument(
+          "cache_read.query_stash",
+          adapter: @adapter,
+          rule: rule.name
+        ) do |payload|
+          cache_store.read(cache_key).tap do |value|
+            payload[:hit] = !value.nil?
+          end
+        end
+
+        # Results are never nil, so nil means a miss.
+        unless cached_result.nil?
+          @instrumenter.instrument("cache_hit.query_stash", adapter: @adapter, rule: rule.name)
+          return cached_result
+        end
+      end
+
+      nil
+    end
+
+    # Stores `result` for the query when the rule keeps results for any time (else it is skipped), and
+    # returns it. Emits `cache_miss` or, when skipped, `cache_skip`.
+    # A failure of the cache store is reported (or raised), and does not lose the result.
+    #: (Registry::Rule rule, untyped result, *untyped args) -> untyped
+    def write(rule, result, *args)
+      ttl_seconds = rule.ttl_seconds_for(result)
+      ttl_jitter_seconds = rule.ttl_jitter_seconds_for(result)
+
+      rescue_cache_error do
+        if ttl_seconds.positive?
+          @instrumenter.instrument("cache_miss.query_stash", adapter: @adapter, rule: rule.name)
+          cache_for(rule).write(cache_key_for(*args), result, expires_in: ttl_seconds + rand(ttl_jitter_seconds))
+        else
+          @instrumenter.instrument("cache_skip.query_stash", adapter: @adapter, rule: rule.name)
+        end
+      end
+
+      result
+    end
+
+    private
+
+    # Whether the stash is on for the current execution context: what `enable!` or `disable!` set
+    # there, or else `enabled`.
+    #: () -> bool
+    def stashing?
+      @isolated_state.key?(self) ? @isolated_state[self] : enabled
+    end
+
+    # Sets the stash on or off in the isolated state for the block, then puts back what was there.
+    #: [T] (bool value) { () -> T } -> T
+    def isolate(value)
+      had_value = @isolated_state.key?(self)
+      previous = @isolated_state[self]
+      @isolated_state[self] = value
+      yield
+    ensure
+      if had_value
+        @isolated_state[self] = previous
+      else
+        @isolated_state.delete(self)
+      end
+    end
+
+    #: () -> Registry
+    def build_registry
+      Registry.new
+    end
+
+    #: [T] () { () -> T } -> untyped
+    def rescue_cache_error(&block)
+      block.call
+    rescue StandardError => e
+      raise e if @raise_on_cache_error
+
+      @error_reporter.report(e, handled: true)
+    end
+  end
+end

data/lib/query_stash/mongo.rb

@@ -0,0 +1,152 @@
+# frozen_string_literal: true
+
+require "digest"
+require "json"
+require_relative "base"
+
+module QueryStash
+  class Mongo < Base
+    # What the patch and `wrap` ask of the driver's own objects.
+    #
+    # @rbs!
+    #   # Has a name: a collection, or a database.
+    #   interface _Named
+    #     def name: () -> String
+    #   end
+    #
+    #   # A `Mongo::Collection::View`, or an aggregation view.
+    #   interface _View
+    #     def collection: () -> _Named
+    #
+    #     def database: () -> _Named
+    #
+    #     def filter: () -> untyped
+    #
+    #     def options: () -> Hash[Symbol, untyped]
+    #
+    #     def each: () ?{ (untyped document) -> untyped } -> untyped
+    #   end
+    #
+    #   # A selector to match rules against.
+    #   type selector = Hash[Symbol, untyped]
+
+    # ---------------- Constants ----------------
+    # Wildcard sentinel for selector patterns
+    ANY = Object.new.freeze #: Object
+
+    class MongoFilterRegistry < Registry
+      private
+
+      # Pattern semantics:
+      #
+      # - ANY sentinel:
+      #     pattern.equal?(ANY) => match anything as long as the key exists
+      #: (untyped actual, untyped pattern) -> bool
+      def value_matches_pattern?(actual, pattern)
+        # Wildcard – anything goes (key presence already checked by caller)
+        return true if pattern.equal?(ANY)
+
+        super
+      end
+    end
+
+    # Enables the stash for the duration of each request. `application` is the Rails application the
+    # middleware stack belongs to; `app` is the next middleware.
+    class Middleware
+      # @rbs @app: _RackApp
+      # @rbs @application: untyped
+
+      #: (_RackApp app, untyped application) -> void
+      def initialize(app, application)
+        @app = app
+        @application = application
+      end
+
+      #: (Hash[String, untyped] env) -> untyped
+      def call(env)
+        QueryStash::Mongo[@application].enable! { @app.call(env) }
+      end
+    end
+
+    # @rbs module-self _View
+    module AdapterPatch
+      #: () ?{ (untyped document) -> untyped } -> untyped
+      def each(&block)
+        QueryStash::Mongo[Rails.application].wrap(self, block) { |&iterate| super(&iterate) }
+      end
+    end
+
+    # Runs the view's cursor through the stash: `cursor` iterates the view's documents into the block
+    # it is given, and `consumer` is the block `each` was called with. (RBS can't say that a block takes
+    # a block, so the parameters of `cursor` are not typed.)
+    #: (_View view, ?(^(untyped) -> untyped)? consumer) { (?) -> untyped } -> untyped
+    def wrap(view, consumer = nil, &cursor)
+      rule = rule_for(view)
+      # Avoid eager loading the entire result set unless this query actually matches a cache rule.
+      # This can cause very quick OOMs if it has to load a whole collection at once.
+      return cursor.call(&consumer) unless rule
+
+      documents = read(rule, view) || write(rule, load_documents(&cursor), view)
+
+      consumer ? documents.each(&consumer) : documents.to_enum
+    end
+
+    # The rule the view is stashed under, or nil when it isn't to be stashed: the stash is off, or no rule
+    # matches its selector.
+    #: (_View view) -> Registry::Rule?
+    def rule_for(view)
+      return unless stashing?
+
+      registry.match_rule(selector_for(view))
+    end
+
+    # Based off of #cache_key from Mongo::QueryCache
+    # https://github.com/mongodb/mongo-ruby-driver/blob/master/lib/mongo/query_cache.rb#L19
+    #: (_View view) -> String
+    def cache_key_for(view)
+      selector = selector_for(view)
+      raw = [
+        selector[:collection_name],
+        selector[:filter],
+        selector[:limit],
+        selector[:skip],
+        selector[:sort],
+        selector[:projection],
+        selector[:pipeline]
+      ].to_json
+
+      "v1:#{view.database.name}:#{Digest::SHA256.hexdigest(raw)}"
+    end
+
+    private
+
+    #: () -> MongoFilterRegistry
+    def build_registry
+      MongoFilterRegistry.new
+    end
+
+    # Based off of #cache_options in Mongo::Collection::View::Iterable
+    # https://github.com/mongodb/mongo-ruby-driver/blob/v2.19.3/lib/mongo/collection/view/iterable.rb
+    #: (_View view) -> selector
+    def selector_for(view)
+      # An aggregation view has a pipeline where the others have a filter.
+      pipelined = view #: untyped
+      {
+        collection_name: view.collection.name,
+        filter: view.filter,
+        limit: view.options[:limit],
+        skip: view.options[:skip],
+        sort: view.options[:sort],
+        projection: view.options[:projection],
+        pipeline: pipelined.respond_to?(:pipeline) ? pipelined.pipeline : nil
+      }
+    end
+
+    #: () { (?) -> untyped } -> Array[untyped]
+    def load_documents(&cursor)
+      docs = [] #: Array[untyped]
+      cursor.call { |doc| docs << doc }
+      docs
+    end
+  end
+end