redis-memo 0.1.1 → 1.1.0

data/lib/redis_memo/memoize_query/cached_select.rb

@@ -1,6 +1,6 @@
 # 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
@@ -94,6 +94,12 @@ class RedisMemo::MemoizeQuery::CachedSelect
 
   @@enabled_models = {}
 
+  # Thread locals to exchange information between RedisMemo and ActiveRecord
+  RedisMemo::ThreadLocalVar.define :arel
+  RedisMemo::ThreadLocalVar.define :substitues
+  RedisMemo::ThreadLocalVar.define :arel_bind_params
+
+  # @return [Hash] models enabled for caching
   def self.enabled_models
     @@enabled_models
   end
@@ -107,10 +113,7 @@ class RedisMemo::MemoizeQuery::CachedSelect
 
       memoize_method(
         :exec_query,
-        method_id: proc do |_, sql, *args|
-          sql.gsub(/(\$\d+)/, '?')      # $1 -> ?
-             .gsub(/((, *)*\?)+/, '?')  # (?, ?, ? ...) -> (?)
-        end,
+        method_id: proc { |_, sql, *| RedisMemo::Util.tagify_parameterized_sql(sql) },
       ) do |_, sql, _, binds, **|
         depends_on RedisMemo::MemoizeQuery::CachedSelect.current_query_bind_params
 
@@ -123,7 +126,7 @@ class RedisMemo::MemoizeQuery::CachedSelect
               # In activerecord >= 6, a bind could be an actual database value
               bind
             end
-          end
+          end,
         )
       end
     end
@@ -152,61 +155,74 @@ class RedisMemo::MemoizeQuery::CachedSelect
     end
   end
 
+  # Extract bind params from the query by inspecting the SQL's AST recursively
+  # The bind params will be passed into the local thread variables. See
+  # +construct_bind_params_recurse+ for how to construct binding params
+  # recursively.
+  #
+  # @param sql [String] SQL query
+  # @return [Boolean] indicating whether a query should be cached
   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
+    RedisMemo::Tracer.trace(
+      'redis_memo.memoize_query.extract_bind_params',
+      RedisMemo::Util.tagify_parameterized_sql(sql),
+    ) do
+      ast = RedisMemo::ThreadLocalVar.arel&.ast
+      return false unless ast.is_a?(Arel::Nodes::SelectStatement)
+      return false unless ast.to_sql == sql
+
+      RedisMemo::ThreadLocalVar.substitues ||= {}
+      # Iterate through the Arel AST in a Depth First Search
+      bind_params = construct_bind_params_recurse(ast)
+      return false unless bind_params&.should_cache?
+
+      bind_params.extract!
+      RedisMemo::ThreadLocalVar.arel_bind_params = bind_params
+      true
+    end
   end
 
   def self.current_query_bind_params
-    Thread.current[THREAD_KEY_AREL_BIND_PARAMS]
+    RedisMemo::ThreadLocalVar.arel_bind_params
   end
 
   def self.current_query=(arel)
-    Thread.current[THREAD_KEY_AREL] = arel
+    RedisMemo::ThreadLocalVar.arel = arel
   end
 
   def self.current_substitutes=(substitutes)
-    Thread.current[THREAD_KEY_SUBSTITUTES] = substitutes
+    RedisMemo::ThreadLocalVar.substitues = 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
+    RedisMemo::ThreadLocalVar.arel = nil
+    RedisMemo::ThreadLocalVar.substitues = nil
+    RedisMemo::ThreadLocalVar.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]
+    prev_arel = RedisMemo::ThreadLocalVar.arel
+    prev_substitutes = RedisMemo::ThreadLocalVar.substitues
+    prev_bind_params = RedisMemo::ThreadLocalVar.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
+    RedisMemo::ThreadLocalVar.arel = prev_arel
+    RedisMemo::ThreadLocalVar.substitues = prev_substitutes
+    RedisMemo::ThreadLocalVar.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)
+  #
+  # @param node [Arel::Nodes::Node]
+  #
+  # @return [RedisMemo::MemoizeQuery::CachedSelect::BindParams]
+  def self.construct_bind_params_recurse(node)
+    # rubocop: disable Lint/NonLocalExitFromIterator
     bind_params = BindParams.new
 
     case node
@@ -229,7 +245,7 @@ class RedisMemo::MemoizeQuery::CachedSelect
       return unless binding_relation
 
       rights = node.right.is_a?(Array) ? node.right : [node.right]
-      substitutes = Thread.current[THREAD_KEY_SUBSTITUTES]
+      substitutes = RedisMemo::ThreadLocalVar.substitues
 
       rights.each do |right|
         case right
@@ -256,20 +272,13 @@ class RedisMemo::MemoizeQuery::CachedSelect
               end,
           }
         else
-          bind_params = bind_params.union(extract_bind_params_recurse(right))
-          if bind_params
-            next
-          else
-            return
-          end
+          bind_params = bind_params.union(construct_bind_params_recurse(right))
+          return if !bind_params
         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?
@@ -284,7 +293,7 @@ class RedisMemo::MemoizeQuery::CachedSelect
           return if core.wheres.empty? || binding_relation.nil?
         when Arel::Nodes::TableAlias
           bind_params = bind_params.union(
-            extract_bind_params_recurse(source_node.left)
+            construct_bind_params_recurse(source_node.left),
           )
 
           return unless bind_params
@@ -295,7 +304,7 @@ class RedisMemo::MemoizeQuery::CachedSelect
         # Binds wheres before havings
         core.wheres.each do |where|
           bind_params = bind_params.union(
-            extract_bind_params_recurse(where)
+            construct_bind_params_recurse(where),
           )
 
           return unless bind_params
@@ -303,26 +312,23 @@ class RedisMemo::MemoizeQuery::CachedSelect
 
         core.havings.each do |having|
           bind_params = bind_params.union(
-            extract_bind_params_recurse(having)
+            construct_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)
+      construct_bind_params_recurse(node.expr)
+    when Arel::Nodes::LessThan, Arel::Nodes::LessThanOrEqual, Arel::Nodes::GreaterThan, Arel::Nodes::GreaterThanOrEqual, Arel::Nodes::NotEqual
+      bind_params
     when Arel::Nodes::And
       node.children.each do |child|
         bind_params = bind_params.product(
-          extract_bind_params_recurse(child)
+          construct_bind_params_recurse(child),
         )
 
         return unless bind_params
@@ -332,7 +338,7 @@ class RedisMemo::MemoizeQuery::CachedSelect
     when Arel::Nodes::Union, Arel::Nodes::Or
       [node.left, node.right].each do |child|
         bind_params = bind_params.union(
-          extract_bind_params_recurse(child)
+          construct_bind_params_recurse(child),
         )
 
         return unless bind_params
@@ -341,14 +347,24 @@ class RedisMemo::MemoizeQuery::CachedSelect
       bind_params
     else
       # Not yet supported
-      return
+      nil
     end
+    # rubocop: enable Lint/NonLocalExitFromIterator
   end
 
+  # Retrieve the model info from the table node
+  # table node is an Arel::Table object, e.g. <Arel::Table @name="sites" ...>
+  # and we can retrieve the model info by inspecting thhe table name
+  # See +RedisMemo::MemoizeQuery::memoize_table_column+ for how to construct enabled_models
+  #
+  # @params table_node [Arel::Table]
   def self.extract_binding_relation(table_node)
     enabled_models[table_node.try(:name)]
   end
 
+  #
+  # Identify whether the node has filter condition
+  #
   class NodeHasFilterCondition
     def self.===(node)
       case node
@@ -365,9 +381,4 @@ class RedisMemo::MemoizeQuery::CachedSelect
       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

data/lib/redis_memo/memoize_query/cached_select/bind_params.rb

@@ -2,47 +2,122 @@
 
 class RedisMemo::MemoizeQuery::CachedSelect
   class BindParams
-    def params
-      #
-      # Bind params is hash of sets: each key is a model class, each value is a
-      # set of hashes for memoized column conditions. Example:
-      #
-      #   {
-      #      Site => [
-      #        {name: 'a', city: 'b'},
-      #        {name: 'a', city: 'c'},
-      #        {name: 'b', city: 'b'},
-      #        {name: 'b', city: 'c'},
-      #      ],
-      #   }
-      #
-      @params ||= Hash.new do |models, model|
-        models[model] = []
-      end
+    def initialize(left = nil, right = nil, operator = nil)
+      @left = left
+      @right = right
+      @operator = operator
     end
 
     def union(other)
       return unless other
 
-      # The tree is almost always right-heavy. Merge into the right node for better
-      # performance.
-      other.params.merge!(params) do |_, other_attrs_set, attrs_set|
-        if other_attrs_set.empty?
-          attrs_set
-        elsif attrs_set.empty?
-          other_attrs_set
-        else
-          attrs_set + other_attrs_set
+      self.class.new(self, other, __method__)
+    end
+
+    def product(other)
+      return unless other
+
+      self.class.new(self, other, __method__)
+    end
+
+    def should_cache?
+      plan!
+
+      if plan.model_attrs.empty? || plan.dependency_size_estimation.to_i > RedisMemo::DefaultOptions.max_query_dependency_size
+        return false
+      end
+
+      plan.model_attrs.each do |model, attrs_set|
+        return false if attrs_set.empty?
+
+        attrs_set.each do |attrs|
+          return false unless RedisMemo::MemoizeQuery
+            .memoized_columns(model)
+            .include?(attrs.keys.sort)
         end
       end
 
-      other
+      true
     end
 
-    def product(other)
+    #
+    # Extracted bind params is hash of sets: each key is a model class, each
+    # value is a set of hashes for memoized column conditions. Example:
+    #
+    #   {
+    #      Site => [
+    #        {name: 'a', city: 'b'},
+    #        {name: 'a', city: 'c'},
+    #        {name: 'b', city: 'b'},
+    #        {name: 'b', city: 'c'},
+    #      ],
+    #   }
+    #
+    def extract!
+      return if operator.nil?
+
+      left.extract!
+      right.extract!
+      __send__(:"#{operator}!")
+    end
+
+    def params
+      @params ||= Hash.new do |models, model|
+        models[model] = Set.new
+      end
+    end
+
+    protected
+
+    # BindParams is built recursively when iterating through the Arel AST
+    # nodes. BindParams represents a binary tree. Query parameters are added to
+    # the leaf nodes of the tree, and the leaf nodes are connected by
+    # operators, such as `union` (or conditions) or `product` (and conditions).
+    attr_accessor :left
+    attr_accessor :right
+    attr_accessor :operator
+    attr_accessor :plan
+
+    def plan!
+      self.plan = Plan.new(self)
+      return if operator.nil?
+
+      left.plan!
+      right.plan!
+      __send__(:"plan_#{operator}")
+    end
+
+    def plan_union
+      plan.dependency_size_estimation = left.plan.dependency_size_estimation + right.plan.dependency_size_estimation
+      plan.model_attrs = union_attrs_set(left.plan.model_attrs, right.plan.model_attrs)
+    end
+
+    def plan_product
+      plan.dependency_size_estimation = left.plan.dependency_size_estimation * right.plan.dependency_size_estimation
+      plan.model_attrs = product_attrs_set(left.plan.model_attrs, right.plan.model_attrs)
+    end
+
+    def union!
+      @params = union_attrs_set(left.params, right.params)
+    end
+
+    def product!
+      @params = product_attrs_set(left.params, right.params)
+    end
+
+    def union_attrs_set(left, right)
+      left.merge(right) do |_, attrs_set, other_attrs_set|
+        next attrs_set if other_attrs_set.empty?
+        next other_attrs_set if attrs_set.empty?
+
+        attrs_set + other_attrs_set
+      end
+    end
+
+    def product_attrs_set(left, right)
       #  Example:
       #
-      #  and(
+      #  product(
       #    [{a: 1}, {a: 2}],
       #    [{b: 1}, {b: 2}],
       #  )
@@ -55,29 +130,16 @@ class RedisMemo::MemoizeQuery::CachedSelect
       #    {a: 2, b: 1},
       #    {a: 2, b: 2},
       #  ]
-      return unless other
-
-      # The tree is almost always right-heavy. Merge into the right node for better
-      # performance.
-      params.each do |model, attrs_set|
-        next if attrs_set.empty?
-
-        # The other model does not have any conditions so far: carry the
-        # attributes over to the other node
-        if other.params[model].empty?
-          other.params[model] = attrs_set
-          next
-        end
-
-        # Distribute the current attrs into the other
-        other_attrs_set_size = other.params[model].size
-        other_attrs_set = other.params[model]
-        merged_attrs_set = Array.new(other_attrs_set_size * attrs_set.size)
+      left.merge(right) do |_, attrs_set, other_attrs_set|
+        next attrs_set if other_attrs_set.empty?
+        next other_attrs_set if attrs_set.empty?
 
-        attrs_set.each_with_index do |attrs, i|
-          other_attrs_set.each_with_index do |other_attrs, j|
-            k = i * other_attrs_set_size + j
-            merged_attrs = merged_attrs_set[k] = other_attrs.dup
+        # distribute the current attrs into the other
+        merged_attrs_set = Set.new
+        attrs_set.each do |attrs|
+          other_attrs_set.each do |other_attrs|
+            merged_attrs = other_attrs.dup
+            should_add_attrs = true
             attrs.each do |name, val|
               # Conflict detected. For example:
               #
@@ -86,42 +148,112 @@ class RedisMemo::MemoizeQuery::CachedSelect
               #  Keep:     a = 1 and b = 2, a = 2 and b = 1
               #  Discard:  a = 1 and a = 2, b = 1 and b = 2
               if merged_attrs.include?(name) && merged_attrs[name] != val
-                merged_attrs_set[k] = nil
+                should_add_attrs = false
                 break
               end
 
               merged_attrs[name] = val
             end
+            merged_attrs_set << merged_attrs if should_add_attrs
           end
         end
 
-        merged_attrs_set.compact!
-        other.params[model] = merged_attrs_set
+        merged_attrs_set
       end
-
-      other
     end
 
-    def uniq!
-      params.each do |_, attrs_set|
-        attrs_set.uniq!
-      end
-    end
+    # Prior to actually extracting the bind parameters, we first quickly
+    # estimate if it makes sense to do so. If a query contains too many
+    # dependencies, or contains dependencies that have not been memoized, then
+    # the query itself cannot be cached correctly/efficiently, so there’s no
+    # point to actually extract.
+    #
+    # The planning phase is similar to the extraction phase. Though in the
+    # planning phase, we can ignore all the actual attribute values and only
+    # look at the attribute names. This way, we can precompute the dependency
+    # size without populating their actual values.
+    #
+    # For example, in the planning phase,
+    #
+    #   {a:nil} x {b: nil} => {a: nil, b: nil}
+    #   {a:nil, b:nil} x {a: nil: b: nil} => {a: nil, b: nil}
+    #
+    # and in the extraction phase, that's where the # of dependency can
+    # actually grow significantly:
+    #
+    #   {a: [1,2,3]} x {b: [1,2,3]} => [{a: 1, b: 1}, ....]
+    #   {a:[1,2], b:[1,2]} x {a: [1,2,3]: b: [1,2,3]} => [{a: 1, b: 1}, ...]
+    #
+    class Plan
+      class DependencySizeEstimation
+        def initialize(hash = nil)
+          @hash = hash
+        end
 
-    def memoizable?
-      return false if params.empty?
+        def +(other)
+          merged_hash = hash.dup
+          other.hash.each do |k, v|
+            merged_hash[k] += v
+          end
+          self.class.new(merged_hash)
+        end
 
-      params.each do |model, attrs_set|
-        return false if attrs_set.empty?
+        def *(other)
+          merged_hash = hash.dup
+          other.hash.each do |k, v|
+            if merged_hash.include?(k)
+              merged_hash[k] *= v
+            else
+              merged_hash[k] = v
+            end
+          end
+          self.class.new(merged_hash)
+        end
 
-        attrs_set.each do |attrs|
-          return false unless RedisMemo::MemoizeQuery
-            .memoized_columns(model)
-            .include?(attrs.keys.sort)
+        def [](key)
+          hash[key]
+        end
+
+        def []=(key, val)
+          hash[key] = val
+        end
+
+        def to_i
+          ret = 0
+          hash.each do |_, v|
+            ret += v
+          end
+          ret
+        end
+
+        protected
+
+        def hash
+          @hash ||= Hash.new(0)
         end
       end
 
-      true
+      attr_accessor :dependency_size_estimation
+      attr_accessor :model_attrs
+
+      def initialize(bind_params)
+        @dependency_size_estimation = DependencySizeEstimation.new
+        @model_attrs = Hash.new do |models, model|
+          models[model] = Set.new
+        end
+
+        # An aggregated bind_params node can only obtain params by combining
+        # its children nodes
+        return if !bind_params.__send__(:operator).nil?
+
+        bind_params.params.each do |model, attrs_set|
+          @dependency_size_estimation[model] += attrs_set.size
+          attrs_set.each do |attrs|
+            # [k, nil]: Ignore the attr value and keep the name only
+            @model_attrs[model] << attrs.keys.map { |k| [k, nil] }.to_h
+          end
+        end
+      end
     end
   end
 end