Merge pull request #75 from crashtech/recursive_cte

Implement Recursive CTE

lib/torque/postgresql/auxiliary_statement.rb

@@ -1,40 +1,44 @@
 # frozen_string_literal: true
 
 require_relative 'auxiliary_statement/settings'
+require_relative 'auxiliary_statement/recursive'
 
 module Torque
   module PostgreSQL
     class AuxiliaryStatement
       TABLE_COLUMN_AS_STRING = /\A(?:"?(\w+)"?\.)?"?(\w+)"?\z/.freeze
 
       class << self
-        attr_reader :config
+        attr_reader :config, :table_name
 
         # Find or create the class that will handle statement
         def lookup(name, base)
           const = name.to_s.camelize << '_' << self.name.demodulize
           return base.const_get(const, false) if base.const_defined?(const, false)
-          base.const_set(const, Class.new(AuxiliaryStatement))
+
+          base.const_set(const, Class.new(self)).tap do |klass|
+            klass.instance_variable_set(:@table_name, name.to_s)
+          end
         end
 
         # Create a new instance of an auxiliary statement
-        def instantiate(statement, base, options = nil)
+        def instantiate(statement, base, **options)
           klass = while base < ActiveRecord::Base
             list = base.auxiliary_statements_list
             break list[statement] if list.present? && list.key?(statement)
 
             base = base.superclass
           end
 
-          return klass.new(options) unless klass.nil?
+          return klass.new(**options) unless klass.nil?
           raise ArgumentError, <<-MSG.squish
             There's no '#{statement}' auxiliary statement defined for #{base.class.name}.
           MSG
         end
 
         # Fast access to statement build
-        def build(statement, base, options = nil, bound_attributes = [], join_sources = [])
-          klass = instantiate(statement, base, options)
+        def build(statement, base, bound_attributes = [], join_sources = [], **options)
+          klass = instantiate(statement, base, **options)
           result = klass.build(base)
 
           bound_attributes.concat(klass.bound_attributes)
@@ -56,7 +60,7 @@ def arel_query?(obj)
         # A way to create auxiliary statements outside of models configurations,
         # being able to use on extensions
         def create(table_or_settings, &block)
-          klass = Class.new(AuxiliaryStatement)
+          klass = Class.new(self)
 
           if block_given?
             klass.instance_variable_set(:@table_name, table_or_settings)
@@ -89,7 +93,8 @@ def configurator(config)
         def configure(base, instance)
           return @config unless @config.respond_to?(:call)
 
-          settings = Settings.new(base, instance)
+          recursive = self < AuxiliaryStatement::Recursive
+          settings = Settings.new(base, instance, recursive)
           settings.instance_exec(settings, &@config)
           settings
         end
@@ -98,11 +103,6 @@ def configure(base, instance)
         def table
           @table ||= ::Arel::Table.new(table_name)
         end
-
-        # Get the name of the table of the configurated statement
-        def table_name
-          @table_name ||= self.name.demodulize.split('_').first.underscore
-        end
       end
 
       delegate :config, :table, :table_name, :relation, :configure, :relation_query?,
@@ -111,15 +111,14 @@ def table_name
       attr_reader :bound_attributes, :join_sources
 
       # Start a new auxiliary statement giving extra options
-      def initialize(*args)
-        options = args.extract_options!
+      def initialize(*, **options)
         args_key = Torque::PostgreSQL.config.auxiliary_statement.send_arguments_key
 
         @join = options.fetch(:join, {})
         @args = options.fetch(args_key, {})
         @where = options.fetch(:where, {})
         @select = options.fetch(:select, {})
-        @join_type = options.fetch(:join_type, nil)
+        @join_type = options[:join_type]
 
         @bound_attributes = []
         @join_sources = []
@@ -131,7 +130,7 @@ def build(base)
         @join_sources.clear
 
         # Prepare all the data for the statement
-        prepare(base)
+        prepare(base, configure(base, self))
 
         # Add the join condition to the list
         @join_sources << build_join(base)
@@ -141,24 +140,22 @@ def build(base)
       end
 
       private
+
         # Setup the statement using the class configuration
-        def prepare(base)
-          settings = configure(base, self)
+        def prepare(base, settings)
           requires = Array.wrap(settings.requires).flatten.compact
           @dependencies = ensure_dependencies(requires, base).flatten.compact
 
           @join_type ||= settings.join_type || :inner
           @query = settings.query
 
           # Call a proc to get the real query
-          if @query.methods.include?(:call)
+          if @query.respond_to?(:call)
             call_args = @query.try(:arity) === 0 ? [] : [OpenStruct.new(@args)]
             @query = @query.call(*call_args)
-            @args = []
           end
 
-          # Manually set the query table when it's not an relation query
-          @query_table = settings.query_table unless relation_query?(@query)
+          # Merge select attributes provided on the instance creation
           @select = settings.attributes.merge(@select) if settings.attributes.present?
 
           # Merge join settings
@@ -168,7 +165,7 @@ def prepare(base)
             @association = settings.through.to_s
           elsif relation_query?(@query)
             @association = base.reflections.find do |name, reflection|
-              break name if @query.klass.eql? reflection.klass
+              break name if @query.klass.eql?(reflection.klass)
             end
           end
         end
@@ -234,15 +231,6 @@ def build_join(base)
             as a query object on #{self.class.name}.
           MSG
 
-          # Expose join columns
-          if relation_query?(@query)
-            query_table = @query.arel_table
-            conditions.children.each do |item|
-              @query.select_values += [query_table[item.left.name]] \
-                if item.left.relation.eql?(table)
-            end
-          end
-
           # Build the join based on the join type
           arel_join.new(table, table.create_on(conditions))
         end
@@ -263,29 +251,40 @@ def arel_join
 
         # Mount the list of selected attributes
         def expose_columns(base, query_table = nil)
+          # Add the columns necessary for the join
+          list = @join_sources.each_with_object(@select) do |join, hash|
+            join.right.expr.children.each do |item|
+              hash[item.left.name] = nil if item.left.relation.eql?(table)
+            end
+          end
+
           # Add select columns to the query and get exposed columns
-          @select.map do |left, right|
-            base.select_extra_values += [table[right.to_s]]
-            project(left, query_table).as(right.to_s) if query_table
+          list.filter_map do |left, right|
+            base.select_extra_values += [table[right.to_s]] unless right.nil?
+            next unless query_table
+
+            col = project(left, query_table)
+            right.nil? ? col : col.as(right.to_s)
           end
         end
 
         # Ensure that all the dependencies are loaded in the base relation
         def ensure_dependencies(list, base)
           with_options = list.extract_options!.to_a
-          (list + with_options).map do |dependent, options|
-            dependent_klass = base.model.auxiliary_statements_list[dependent]
+          (list + with_options).map do |name, options|
+            dependent_klass = base.model.auxiliary_statements_list[name]
 
             raise ArgumentError, <<-MSG.squish if dependent_klass.nil?
-              The '#{dependent}' auxiliary statement dependency can't found on
+              The '#{name}' auxiliary statement dependency can't found on
               #{self.class.name}.
             MSG
 
             next if base.auxiliary_statements_values.any? do |cte|
               cte.is_a?(dependent_klass)
             end
 
-            AuxiliaryStatement.build(dependent, base, options, bound_attributes, join_sources)
+            options ||= {}
+            AuxiliaryStatement.build(name, base, bound_attributes, join_sources, **options)
           end
         end
 

lib/torque/postgresql/auxiliary_statement/recursive.rb

@@ -0,0 +1,149 @@
+# frozen_string_literal: true
+
+module Torque
+  module PostgreSQL
+    class AuxiliaryStatement
+      class Recursive < AuxiliaryStatement
+        # Setup any additional option in the recursive mode
+        def initialize(*, **options)
+          super
+
+          @connect = options[:connect]&.to_a&.first
+          @union_all = options[:union_all]
+          @sub_query = options[:sub_query]
+
+          if options.key?(:with_depth)
+            @depth = options[:with_depth].values_at(:name, :start, :as)
+            @depth[0] ||= 'depth'
+          end
+
+          if options.key?(:with_path)
+            @path = options[:with_path].values_at(:name, :source, :as)
+            @path[0] ||= 'path'
+          end
+        end
+
+        private
+
+          # Build the string or arel query
+          def build_query(base)
+            # Expose columns and get the list of the ones for select
+            columns = expose_columns(base, @query.try(:arel_table))
+            sub_columns = columns.dup
+            type = @union_all.present? ? 'all' : ''
+
+            # Build any extra columns that are dynamic and from the recursion
+            extra_columns(base, columns, sub_columns)
+
+            # Prepare the query depending on its type
+            if @query.is_a?(String) && @sub_query.is_a?(String)
+              args = @args.each_with_object({}) { |h, (k, v)| h[k] = base.connection.quote(v) }
+              ::Arel.sql("(#{@query} UNION #{type.upcase} #{@sub_query})" % args)
+            elsif relation_query?(@query)
+              @query = @query.where(@where) if @where.present?
+              @bound_attributes.concat(@query.send(:bound_attributes))
+
+              if relation_query?(@sub_query)
+                @bound_attributes.concat(@sub_query.send(:bound_attributes))
+
+                sub_query = @sub_query.select(*sub_columns).arel
+                sub_query.from([@sub_query.arel_table, table])
+              else
+                sub_query = ::Arel.sql(@sub_query)
+              end
+
+              @query.select(*columns).arel.union(type, sub_query)
+            else
+              raise ArgumentError, <<-MSG.squish
+                Only String and ActiveRecord::Base objects are accepted as query and sub query
+                objects, #{@query.class.name} given for #{self.class.name}.
+              MSG
+            end
+          end
+
+          # Setup the statement using the class configuration
+          def prepare(base, settings)
+            super
+
+            prepare_sub_query(base, settings)
+          end
+
+          # Make sure that both parts of the union are ready
+          def prepare_sub_query(base, settings)
+            @union_all = settings.union_all if @union_all.nil?
+            @sub_query ||= settings.sub_query
+            @depth ||= settings.depth
+            @path ||= settings.path
+
+            # Collect the connection
+            @connect ||= settings.connect || begin
+              key = base.primary_key
+              [key.to_sym, :"parent_#{key}"] unless key.nil?
+            end
+
+            raise ArgumentError, <<-MSG.squish if @sub_query.nil? && @query.is_a?(String)
+              Unable to generate sub query from a string query. Please provide a `sub_query`
+              property on the "#{table_name}" settings.
+            MSG
+
+            if @sub_query.nil?
+              raise ArgumentError, <<-MSG.squish if @connect.blank?
+                Unable to generate sub query without setting up a proper way to connect it
+                with the main query. Please provide a `connect` property on the "#{table_name}"
+                settings.
+              MSG
+
+              left, right = @connect.map(&:to_s)
+              condition = @query.arel_table[right].eq(table[left])
+
+              if @query.where_values_hash.key?(right)
+                @sub_query = @query.unscope(where: right.to_sym).where(condition)
+              else
+                @sub_query = @query.where(condition)
+                @query = @query.where(right => nil)
+              end
+            elsif @sub_query.respond_to?(:call)
+              # Call a proc to get the real sub query
+              call_args = @sub_query.try(:arity) === 0 ? [] : [OpenStruct.new(@args)]
+              @sub_query = @sub_query.call(*call_args)
+            end
+          end
+
+          # Add depth and path if they were defined in settings
+          def extra_columns(base, columns, sub_columns)
+            return if @query.is_a?(String) || @sub_query.is_a?(String)
+
+            # Add the connect attribute to the query
+            if defined?(@connect)
+              columns.unshift(@query.arel_table[@connect[0]])
+              sub_columns.unshift(@sub_query.arel_table[@connect[0]])
+            end
+
+            # Build a column to represent the depth of the recursion
+            if @depth.present?
+              name, start, as = @depth
+              col = table[name]
+              base.select_extra_values += [col.as(as)] unless as.nil?
+
+              columns << ::Arel.sql(start.to_s).as(name)
+              sub_columns << (col + ::Arel.sql('1')).as(name)
+            end
+
+            # Build a column to represent the path of the record access
+            if @path.present?
+              name, source, as = @path
+              source = @query.arel_table[source || @connect[0]]
+
+              col = table[name]
+              base.select_extra_values += [col.as(as)] unless as.nil?
+              parts = [col, source.cast(:varchar)]
+
+              columns << ::Arel.array([source]).cast(:varchar, true).as(name)
+              sub_columns << ::Arel::Nodes::NamedFunction.new('array_append', parts).as(name)
+            end
+          end
+
+      end
+    end
+  end
+end