redis-memo 0.0.0.alpha → 0.0.0.beta.4

data/lib/redis_memo/memoize_query.rb

@@ -0,0 +1,151 @@
+# frozen_string_literal: true
+require_relative 'memoize_method'
+
+if defined?(ActiveRecord)
+  # Hook into ActiveRecord to cache SQL queries and perform auto cache
+  # invalidation
+  module RedisMemo::MemoizeQuery
+    require_relative 'memoize_query/cached_select'
+    require_relative 'memoize_query/invalidation'
+    require_relative 'memoize_query/model_callback'
+
+    # Only editable columns will be used to create memos that are invalidatable
+    # after each record save
+    def memoize_table_column(*raw_columns, editable: true)
+      RedisMemo::MemoizeQuery.using_active_record!(self)
+      return if ENV["REDIS_MEMO_DISABLE_#{self.table_name.upcase}"] == 'true'
+
+      columns = raw_columns.map(&:to_sym).sort
+
+      RedisMemo::MemoizeQuery.memoized_columns(self, editable_only: true) << columns if editable
+      RedisMemo::MemoizeQuery.memoized_columns(self, editable_only: false) << columns
+
+      RedisMemo::MemoizeQuery::ModelCallback.install(self)
+      RedisMemo::MemoizeQuery::Invalidation.install(self)
+
+      if ENV['REDIS_MEMO_DISABLE_CACHED_SELECT'] != 'true'
+        RedisMemo::MemoizeQuery::CachedSelect.install(ActiveRecord::Base.connection)
+      end
+
+      # The code below might fail due to missing DB/table errors
+      columns.each do |column|
+        unless self.columns_hash.include?(column.to_s)
+          raise(
+            RedisMemo::ArgumentError,
+            "'#{self.name}' does not contain column '#{column}'",
+          )
+        end
+      end
+
+      unless ENV["REDIS_MEMO_DISABLE_QUERY_#{self.table_name.upcase}"] == 'true'
+        RedisMemo::MemoizeQuery::CachedSelect.enabled_models[self.table_name] = self
+      end
+    rescue ActiveRecord::NoDatabaseError, ActiveRecord::StatementInvalid
+      # no-opts: models with memoize_table_column decleared might be loaded in
+      # rake tasks that are used to create databases
+    end
+
+    def self.using_active_record!(model_class)
+      unless using_active_record?(model_class)
+        raise RedisMemo::ArgumentError, "'#{model_class.name}' does not use ActiveRecord"
+      end
+    end
+
+    def self.using_active_record?(model_class)
+      model_class.respond_to?(:<) && model_class < ActiveRecord::Base
+    end
+
+    @@memoized_columns = Hash.new { |h, k| h[k] = [Set.new, Set.new] }
+
+    def self.memoized_columns(model_or_table, editable_only: false)
+      table = model_or_table.is_a?(Class) ? model_or_table.table_name : model_or_table
+      @@memoized_columns[table.to_sym][editable_only ? 1 : 0]
+    end
+
+    # extra_props are considered as AND conditions on the model class
+    def self.create_memo(model_class, **extra_props)
+      using_active_record!(model_class)
+
+      keys = extra_props.keys.sort
+      if !keys.empty? && !memoized_columns(model_class).include?(keys)
+        raise(
+          RedisMemo::ArgumentError,
+          "'#{model_class.name}' has not memoized columns: #{keys}",
+        )
+      end
+
+      extra_props.each do |key, value|
+        # The data type is ensured by the database, thus we don't need to cast
+        # types here for better performance
+        column_name = key.to_s
+        extra_props[key] =
+          if model_class.defined_enums.include?(column_name)
+            enum_mapping = model_class.defined_enums[column_name]
+            # Assume a value is a converted enum if it does not exist in the
+            # enum mapping
+            (enum_mapping[value.to_s] || value).to_s
+          else
+            value.to_s
+          end
+      end
+
+      RedisMemo::Memoizable.new(
+        __redis_memo_memoize_query_table_name__: model_class.table_name,
+        **extra_props,
+      )
+    end
+
+    def self.invalidate_all(model_class)
+      RedisMemo::Tracer.trace(
+        'redis_memo.memoizable.invalidate_all',
+        model_class.name,
+      ) do
+        RedisMemo::Memoizable.invalidate([model_class.redis_memo_class_memoizable])
+      end
+    end
+
+    def self.invalidate(*records)
+      RedisMemo::Memoizable.invalidate(
+        records.map { |record| to_memos(record) }.flatten,
+      )
+    end
+
+    def self.to_memos(record)
+      # Invalidate memos with current values
+      memos_to_invalidate = memoized_columns(record.class).map do |columns|
+        props = {}
+        columns.each do |column|
+          props[column] = record.send(column)
+        end
+
+        create_memo(record.class, **props)
+      end
+
+      # Create memos with previous values if
+      #  - there are saved changes
+      #  - this is not creating a new record
+      if !record.saved_changes.empty? && !record.saved_changes.include?(record.class.primary_key)
+        previous_values = {}
+        record.saved_changes.each do |column, (previous_value, _)|
+          previous_values[column.to_sym] = previous_value
+        end
+
+        memoized_columns(record.class, editable_only: true).each do |columns|
+          props = previous_values.slice(*columns)
+          next if props.empty?
+
+          # Fill the column values that have not changed
+          columns.each do |column|
+            next if props.include?(column)
+
+            props[column] = record.send(column)
+          end
+
+          memos_to_invalidate << create_memo(record.class, **props)
+        end
+      end
+
+      memos_to_invalidate
+    end
+  end
+end

data/lib/redis_memo/memoize_query/cached_select.rb

@@ -0,0 +1,380 @@
+# frozen_string_literal: true
+
+#
+# Inspect a SQL's AST to memoize SELECT statements
+#
+# As Rails applies additional logic on top of the rows returned from the
+# database:
+#
+#  - `belongs_to ..., inverse_of: ...`: By using `inverse_of`, Rails could
+#  prevent instantiating the different objects from the DB when the objects are
+#  really the same.
+#
+#  - Associations may have scopes that add more filtering to the existing query
+#
+#  - +ActiveRecord::Relation+ defers the data fetching until the end
+#
+#  - +ActiveRecord::Relation+ could preload associations to avoid N+1 queries
+#
+# Memoizing each SQL query by inspecting its AST is the best approach we have
+# to reliably perform query caching with ActiveRecord.
+#
+# Here's how this works at a high level:
+#
+# First, we extract dependencies from SQL queries. Consider the following query
+#
+#     SELECT * FROM my_records WHERE value = 'a'
+#
+# The rows returned from the database would not change unless records with the
+# value 'a' have been updated. Therefore, if we are to cache this query, we
+# need to set dependencies on this query and discard the cache if the
+# dependencies have changed.
+#
+# Here's the dependency (aka a +Memoizable+) for the above query:
+#
+#     Memoizable.new(model: MyRecord, value: 'a')
+#
+# We bump the column dependencies automatically when updating a record that has
+# the `memoize_table_column` declaration on the model class.
+#
+#     class MyRecord < ApplicationRecord
+#       extend RedisMemo::MemoizeQuery
+#       memoize_table_column :value
+#     end
+#
+# After saving any MyRecord, we will bump the dependencies versions filled with
+# the record's current and past values:
+#
+#     my_record.update(value: 'new_value') # from 'old_value'
+#
+# Then we will bump the versions for at least two memoizables:
+#
+#     Memoizable.new(model: MyRecord, value: 'new_value')
+#     Memoizable.new(model: MyRecord, value: 'old_value')
+#
+#       When the another_value column is also memoized, we have another
+#       memoizable to bump version for, regardless whether the another_value
+#       filed of my_record has been changed:
+#
+#         Memoizable.new(model: MyRecord, another_value:  'current_value')
+#
+#       We need to do this because other columns could be cached in
+#
+#         SELECT * FROM ... WHERE another_value = ?
+#
+#       queries. Those query result sets become stale after the update.
+#
+# By setting dependencies on the query, we will use the dependencies versions
+# as a part of the query cache key. After we bump the dependencies versions,
+# the following request will produce a different new query cache key, so the
+# request will end up with a cache_miss:
+#   - Compute the fresh query result and it will actually send the query to the database
+#   - Fill the new query cache key with the fresh query result
+#
+# After saving my_record and bumping the dependencies versions, all currently
+# cached SQL queries that have `value = 'new_value'` or `value = 'old_value'`
+# in their WHERE clause (or any WHERE conditions that's using the current
+# memoized column values of my_record) can no longer be accessed by any new
+# requests; Those entries will be automatically deleted through cache expiry or
+# cache eviction.
+#
+# We can only memoize SQL queries that can be automatically invalidated through
+# this mechanism:
+#
+#   - The query contains only =, IN conditions
+#   - And those conditions are on table columns that have been memoized via
+#   +memoized_table_column+
+#
+# See +extract_bind_params+ for the precise detection logic.
+#
+class RedisMemo::MemoizeQuery::CachedSelect
+  require_relative 'cached_select/bind_params'
+  require_relative 'cached_select/connection_adapter'
+  require_relative 'cached_select/statement_cache'
+
+  @@enabled_models = {}
+
+  def self.enabled_models
+    @@enabled_models
+  end
+
+  def self.install(connection)
+    klass = connection.class
+    return if klass.singleton_class < RedisMemo::MemoizeMethod
+
+    klass.class_eval do
+      extend RedisMemo::MemoizeMethod
+
+      memoize_method(
+        :exec_query,
+        method_id: proc do |_, sql, *args|
+          sql.gsub(/(\$\d+)/, '?')      # $1 -> ?
+             .gsub(/((, *)*\?)+/, '?')  # (?, ?, ? ...) -> (?)
+        end,
+      ) do |_, sql, name, binds, **kwargs|
+        RedisMemo::MemoizeQuery::CachedSelect
+          .current_query_bind_params
+          .params
+          .each do |model, attrs_set|
+            attrs_set.each do |attrs|
+              depends_on model, **attrs
+            end
+          end
+
+        depends_on RedisMemo::Memoizable.new(
+          __redis_memo_memoize_query_memoize_query_sql__: sql,
+          __redis_memo_memoize_query_memoize_query_binds__: binds.map do |bind|
+            if bind.respond_to?(:value_for_database)
+              bind.value_for_database
+            else
+              # In activerecord >= 6, a bind could be an actual database value
+              bind
+            end
+          end
+        )
+      end
+    end
+
+    klass.prepend(ConnectionAdapter)
+    ActiveRecord::StatementCache.prepend(StatementCache)
+
+    # Cached result objects could be sampled to compare against fresh result
+    # objects. Overwrite the == operator to make the comparison meaningful.
+    ActiveRecord::Result.class_eval do
+      def ==(other)
+        columns == other.columns && rows == other.rows
+      end
+    end
+
+    ActiveRecord::StatementCache::BindMap.class_eval do
+      def map_substitutes(values)
+        ret = {}
+        @indexes.each_with_index do |offset, i|
+          bound_attr = @bound_attributes[offset]
+          substitute = bound_attr.value
+          ret[substitute] = values[i]
+        end
+        ret
+      end
+    end
+  end
+
+  def self.extract_bind_params(sql)
+    ast = Thread.current[THREAD_KEY_AREL]&.ast
+    return false unless ast.is_a?(Arel::Nodes::SelectStatement)
+    return false unless ast.to_sql == sql
+
+    Thread.current[THREAD_KEY_SUBSTITUTES] ||= {}
+    # Iterate through the Arel AST in a Depth First Search
+    bind_params = extract_bind_params_recurse(ast)
+    return false unless bind_params
+
+    bind_params.uniq!
+    return false unless bind_params.memoizable?
+
+    Thread.current[THREAD_KEY_AREL_BIND_PARAMS] = bind_params
+    true
+  end
+
+  def self.current_query_bind_params
+    Thread.current[THREAD_KEY_AREL_BIND_PARAMS]
+  end
+
+  def self.current_query=(arel)
+    Thread.current[THREAD_KEY_AREL] = arel
+  end
+
+  def self.current_substitutes=(substitutes)
+    Thread.current[THREAD_KEY_SUBSTITUTES] = substitutes
+  end
+
+  def self.reset_current_query
+    Thread.current[THREAD_KEY_AREL] = nil
+    Thread.current[THREAD_KEY_SUBSTITUTES] = nil
+    Thread.current[THREAD_KEY_AREL_BIND_PARAMS] = nil
+  end
+
+  def self.with_new_query_context
+    prev_arel = Thread.current[THREAD_KEY_AREL]
+    prev_substitutes = Thread.current[THREAD_KEY_SUBSTITUTES]
+    prev_bind_params = Thread.current[THREAD_KEY_AREL_BIND_PARAMS]
+    RedisMemo::MemoizeQuery::CachedSelect.reset_current_query
+
+    yield
+  ensure
+    Thread.current[THREAD_KEY_AREL] = prev_arel
+    Thread.current[THREAD_KEY_SUBSTITUTES] = prev_substitutes
+    Thread.current[THREAD_KEY_AREL_BIND_PARAMS] = prev_bind_params
+  end
+
+  private
+
+  # A pre-order Depth First Search
+  #
+  # Note: Arel::Nodes#each returns a list in post-order, and it does not step
+  # into Union nodes. So we're implementing our own DFS
+  def self.extract_bind_params_recurse(node)
+    bind_params = BindParams.new
+
+    case node
+    when NodeHasFilterCondition
+      attr_node = node.left
+      return unless attr_node.is_a?(Arel::Attributes::Attribute)
+
+      table_node =
+        case attr_node.relation
+        when Arel::Table
+          attr_node.relation
+        when Arel::Nodes::TableAlias
+          attr_node.relation.left
+        else
+          # Not yet supported
+          return
+        end
+
+      binding_relation = extract_binding_relation(table_node)
+      return unless binding_relation
+
+      rights = node.right.is_a?(Array) ? node.right : [node.right]
+      substitutes = Thread.current[THREAD_KEY_SUBSTITUTES]
+
+      rights.each do |right|
+        case right
+        when Arel::Nodes::BindParam
+          # No need to type cast as they're only used to create +memoizables+
+          # (used as strings)
+          value = right.value.value_before_type_cast
+
+          if value.is_a?(ActiveRecord::StatementCache::Substitute)
+            value = substitutes[value]
+          end
+
+          bind_params.params[binding_relation] << {
+            right.value.name.to_sym => value,
+          }
+        when Arel::Nodes::Casted
+          bind_params.params[binding_relation] << {
+            right.attribute.name.to_sym =>
+              if right.respond_to?(:val)
+                right.val
+              else
+                # activerecord >= 6
+                right.value
+              end,
+          }
+        else
+          bind_params = bind_params.union(extract_bind_params_recurse(right))
+          if bind_params
+            next
+          else
+            return
+          end
+        end
+      end
+
+      bind_params
+    when Arel::Nodes::SelectStatement
+      # No OREDER BY
+      return unless node.orders.empty?
+
+      node.cores.each do |core|
+        # We don't support JOINs
+        return unless core.source.right.empty?
+
+        # Should have a WHERE if directly selecting from a table
+        source_node = core.source.left
+        binding_relation = nil
+        case source_node
+        when Arel::Table
+          binding_relation = extract_binding_relation(source_node)
+
+          return if core.wheres.empty? || binding_relation.nil?
+        when Arel::Nodes::TableAlias
+          bind_params = bind_params.union(
+            extract_bind_params_recurse(source_node.left)
+          )
+
+          return unless bind_params
+        else
+          return
+        end
+
+        # Binds wheres before havings
+        core.wheres.each do |where|
+          bind_params = bind_params.union(
+            extract_bind_params_recurse(where)
+          )
+
+          return unless bind_params
+        end
+
+        core.havings.each do |having|
+          bind_params = bind_params.union(
+            extract_bind_params_recurse(having)
+          )
+
+          return unless bind_params
+        end
+
+        # Reject any unbound select queries
+        return if binding_relation && bind_params.params[binding_relation].empty?
+      end
+
+      bind_params
+    when Arel::Nodes::Grouping
+      # Inline SQL
+      return if node.expr.is_a?(Arel::Nodes::SqlLiteral)
+
+      extract_bind_params_recurse(node.expr)
+    when Arel::Nodes::And
+      node.children.each do |child|
+        bind_params = bind_params.product(
+          extract_bind_params_recurse(child)
+        )
+
+        return unless bind_params
+      end
+
+      bind_params
+    when Arel::Nodes::Union, Arel::Nodes::Or
+      [node.left, node.right].each do |child|
+        bind_params = bind_params.union(
+          extract_bind_params_recurse(child)
+        )
+
+        return unless bind_params
+      end
+
+      bind_params
+    else
+      # Not yet supported
+      return
+    end
+  end
+
+  def self.extract_binding_relation(table_node)
+    enabled_models[table_node.try(:name)]
+  end
+
+  class NodeHasFilterCondition
+    def self.===(node)
+      case node
+      when Arel::Nodes::Equality, Arel::Nodes::In
+        true
+      else
+        # In activerecord >= 6, a new arel node HomogeneousIn is introduced
+        if defined?(Arel::Nodes::HomogeneousIn) &&
+           node.is_a?(Arel::Nodes::HomogeneousIn)
+          true
+        else
+          false
+        end
+      end
+    end
+  end
+
+  # Thread locals to exchange information between RedisMemo and ActiveRecord
+  THREAD_KEY_AREL = :__redis_memo_memoize_query_cached_select_arel__
+  THREAD_KEY_SUBSTITUTES = :__redis_memo_memoize_query_cached_select_substitues__
+  THREAD_KEY_AREL_BIND_PARAMS = :__redis_memo_memoize_query_cached_select_arel_bind_params__
+end