Allow CTEs to be used as the from of queries

Author: Carlos Silva

lib/torque/postgresql/auxiliary_statement.rb

@@ -125,18 +125,18 @@ def initialize(*, **options)
       end
 
       # Build the statement on the given arel and return the WITH statement
-      def build(base)
+      def build(base, joining = true)
         @bound_attributes.clear
         @join_sources.clear
 
         # Prepare all the data for the statement
         prepare(base, configure(base, self))
 
         # Add the join condition to the list
-        @join_sources << build_join(base)
+        @join_sources << build_join(base) if joining
 
         # Return the statement with its dependencies
-        [@dependencies, ::Arel::Nodes::As.new(table, build_query(base))]
+        [@dependencies, ::Arel::Nodes::As.new(table, build_query(base, joining))]
       end
 
       private
@@ -151,7 +151,7 @@ def prepare(base, settings)
 
           # Call a proc to get the real query
           if @query.respond_to?(:call)
-            call_args = @query.try(:arity) === 0 ? [] : [OpenStruct.new(@args)]
+            call_args = @query.try(:arity) === 0 ? nil : [OpenStruct.new(@args)]
             @query = @query.call(*call_args)
           end
 
@@ -171,18 +171,19 @@ def prepare(base, settings)
         end
 
         # Build the string or arel query
-        def build_query(base)
+        def build_query(base, joining = true)
           # Expose columns and get the list of the ones for select
-          columns = expose_columns(base, @query.try(:arel_table))
+          columns = expose_columns(base, @query.try(:arel_table), joining)
 
           # Prepare the query depending on its type
           if @query.is_a?(String)
             args = @args.map{ |k, v| [k, base.connection.quote(v)] }.to_h
             ::Arel.sql("(#{@query})" % args)
           elsif relation_query?(@query)
             @query = @query.where(@where) if @where.present?
+            @query = add_selected_columns(@query, columns, joining)
             @bound_attributes.concat(@query.send(:bound_attributes))
-            @query.select(*columns).arel
+            @query.arel
           else
             raise ArgumentError, <<-MSG.squish
               Only String and ActiveRecord::Base objects are accepted as query objects,
@@ -235,26 +236,14 @@ def build_join(base)
           arel_join.new(table, table.create_on(conditions))
         end
 
-        # Get the class of the join on arel
-        def arel_join
-          case @join_type
-          when :inner then ::Arel::Nodes::InnerJoin
-          when :left  then ::Arel::Nodes::OuterJoin
-          when :right then ::Arel::Nodes::RightOuterJoin
-          when :full  then ::Arel::Nodes::FullOuterJoin
-          else
-            raise ArgumentError, <<-MSG.squish
-              The '#{@join_type}' is not implemented as a join type.
-            MSG
-          end
-        end
-
         # Mount the list of selected attributes
-        def expose_columns(base, query_table = nil)
+        def expose_columns(base, query_table = nil, joining = true)
           # 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)
+          list = !joining ? @select : begin
+            @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
           end
 
@@ -268,6 +257,31 @@ def expose_columns(base, query_table = nil)
           end
         end
 
+        # Project a column on a given table, or use the column table
+        def project(column, arel_table = nil)
+          if column.respond_to?(:as)
+            return column
+          elsif (as_string = TABLE_COLUMN_AS_STRING.match(column.to_s))
+            column = as_string[2]
+            arel_table = ::Arel::Table.new(as_string[1]) unless as_string[1].nil?
+          end
+
+          arel_table ||= table
+          arel_table[column.to_s]
+        end
+
+        # Add the provided +columns+ to the provided +query+ in a way that
+        # complies with the given +joining+ condition
+        def add_selected_columns(query, columns, joining)
+          query.tap do
+            if joining
+              query._select!(*columns)
+            else
+              query.select_extra_values += columns
+            end
+          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
@@ -288,17 +302,18 @@ def ensure_dependencies(list, base)
           end
         end
 
-        # Project a column on a given table, or use the column table
-        def project(column, arel_table = nil)
-          if column.respond_to?(:as)
-            return column
-          elsif (as_string = TABLE_COLUMN_AS_STRING.match(column.to_s))
-            column = as_string[2]
-            arel_table = ::Arel::Table.new(as_string[1]) unless as_string[1].nil?
+        # Get the class of the join on arel
+        def arel_join
+          case @join_type
+          when :inner then ::Arel::Nodes::InnerJoin
+          when :left  then ::Arel::Nodes::OuterJoin
+          when :right then ::Arel::Nodes::RightOuterJoin
+          when :full  then ::Arel::Nodes::FullOuterJoin
+          else
+            raise ArgumentError, <<-MSG.squish
+              The '#{@join_type}' is not implemented as a join type.
+            MSG
           end
-
-          arel_table ||= table
-          arel_table[column.to_s]
         end
     end
   end

lib/torque/postgresql/auxiliary_statement/recursive.rb

@@ -26,33 +26,33 @@ def initialize(*, **options)
         private
 
           # Build the string or arel query
-          def build_query(base)
+          def build_query(base, joining = true)
             # Expose columns and get the list of the ones for select
-            columns = expose_columns(base, @query.try(:arel_table))
+            columns = expose_columns(base, @query.try(:arel_table), joining)
             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)
+            extra_columns(base, columns, sub_columns, joining)
 
             # 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?
+              @query = add_selected_columns(@query, columns, joining)
               @bound_attributes.concat(@query.send(:bound_attributes))
 
               if relation_query?(@sub_query)
+                sub_query = add_selected_columns(@sub_query, sub_columns, joining).arel
+                sub_query = sub_query.from([@sub_query.arel_table, table])
                 @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)
+              @query.arel.union(type, sub_query)
             else
               raise ArgumentError, <<-MSG.squish
                 Only String and ActiveRecord::Base objects are accepted as query and sub query
@@ -104,17 +104,17 @@ def prepare_sub_query(base, settings)
               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)]
+              call_args = @sub_query.try(:arity) === 0 ? nil : [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)
+          def extra_columns(base, columns, sub_columns, joining)
             return if @query.is_a?(String) || @sub_query.is_a?(String)
 
             # Add the connect attribute to the query
-            if defined?(@connect)
+            if defined?(@connect) && joining
               columns.unshift(@query.arel_table[@connect[0]])
               sub_columns.unshift(@sub_query.arel_table[@connect[0]])
             end

lib/torque/postgresql/base.rb

@@ -279,7 +279,7 @@ def dynamic_attribute(name, &block)
           # Changes the type of the join
           #
           # query key:
-          # Save the query command to be performand
+          # Save the query command to be performed
           #
           # requires key:
           # Indicates dependencies with another statements

lib/torque/postgresql/relation/auxiliary_statement.rb

@@ -10,6 +10,16 @@ def auxiliary_statements_values; get_value(:auxiliary_statements); end
         # :nodoc:
         def auxiliary_statements_values=(value); set_value(:auxiliary_statements, value); end
 
+        # Hook into the +from+ method to allow querying from a CTE
+        def from(value, subquery_name = nil, **options)
+          if value.is_a?(Symbol) && auxiliary_statements_list.key?(value)
+            value = auxiliary_statements_list[value]
+            value = value.new(**options)
+          end
+
+          super(value, subquery_name)
+        end
+
         # Set use of an auxiliary statement
         def with(*args, **settings)
           spawn.with!(*args, **settings)
@@ -38,10 +48,29 @@ def bound_attributes
 
           # Hook arel build to add the distinct on clause
           def build_arel(*)
+            # Check if CTE was included as part of the from of the query
+            from = from_clause.value
+            from_cte = from.is_a?(PostgreSQL::AuxiliaryStatement)
+
+            # Build the arel normally and then get the type of the statements
             arel = super
             type = auxiliary_statement_type
-            sub_queries = build_auxiliary_statements(arel)
-            sub_queries.nil? ? arel : arel.with(*type, *sub_queries)
+
+            # Build all the statements and add them to the arel
+            sub_queries = build_auxiliary_statements(arel).flatten
+            sub_queries << from.build(self, false) if from_cte
+            arel.with(*type, *sub_queries) unless sub_queries.empty?
+            arel
+          end
+
+          # Intercept when the FROM clause is being generated to properly build
+          # from a setup CTE
+          def build_from
+            opts = from_clause.value
+            return super unless opts.is_a?(PostgreSQL::AuxiliaryStatement)
+
+            name = from_clause.name
+            name ? opts.table.as(name.to_s) : opts.table_name
           end
 
           # Instantiate one or more auxiliary statements for the given +klass+
@@ -60,7 +89,7 @@ def instantiate_auxiliary_statements(*args, **options)
 
           # Build all necessary data for auxiliary statements
           def build_auxiliary_statements(arel)
-            return unless auxiliary_statements_values.present?
+            return [] unless auxiliary_statements_values.present?
             auxiliary_statements_values.map do |klass|
               klass.build(self).tap { arel.join_sources.concat(klass.join_sources) }
             end
@@ -69,7 +98,8 @@ def build_auxiliary_statements(arel)
           # Return recursive if any auxiliary statement is recursive
           def auxiliary_statement_type
             klass = PostgreSQL::AuxiliaryStatement::Recursive
-            :recursive if auxiliary_statements_values.any?(klass)
+            :recursive if auxiliary_statements_values.any?(klass) ||
+              from_clause.value.is_a?(klass)
           end
 
           # Throw an error showing that an auxiliary statement of the given

lib/torque/postgresql/version.rb

@@ -2,6 +2,6 @@
 
 module Torque
   module PostgreSQL
-    VERSION = '3.2.2'
+    VERSION = '3.2.3'
   end
 end

spec/tests/auxiliary_statement_spec.rb

@@ -27,6 +27,17 @@
       expect(subject.with(:comments).arel.to_sql).to eql(result)
     end
 
+    it 'can be used as the from clause' do
+      klass.send(:auxiliary_statement, :admin_users) do |cte|
+        cte.query User.where(role: :admin)
+      end
+
+      result = 'WITH "admin_users" AS'
+      result << ' (SELECT "users".* FROM "users" WHERE "users"."role" = $1)'
+      result << ' SELECT "users".* FROM "admin_users" AS users'
+      expect(subject.from(:admin_users, :users).arel.to_sql).to eql(result)
+    end
+
     it 'can perform more complex queries' do
       klass.send(:auxiliary_statement, :comments) do |cte|
         cte.query Comment.distinct_on(:user_id).order(:user_id, id: :desc)
@@ -456,6 +467,21 @@
         expect(subject.with(:all_categories).arel.to_sql).to eql(result)
       end
 
+      it 'can be used as the from clause' do
+        Category.send(:recursive_auxiliary_statement, :all_categories) do |cte|
+          cte.query Category.all
+        end
+
+        result = 'WITH RECURSIVE "all_categories" AS ('
+        result << ' SELECT "categories".* FROM "categories"'
+        result << ' WHERE "categories"."parent_id" IS NULL'
+        result << ' UNION'
+        result << ' SELECT "categories".* FROM "categories", "all_categories"'
+        result << ' WHERE "categories"."parent_id" = "all_categories"."id"'
+        result << ' ) SELECT "categories".* FROM "all_categories" AS categories'
+        expect(Category.from(:all_categories, :categories).arel.to_sql).to eql(result)
+      end
+
       it 'allows connect to be set to something different using a single value' do
         klass.send(:recursive_auxiliary_statement, :all_categories) do |cte|
           cte.query Category.all
@@ -615,6 +641,25 @@
         expect(subject.with(:all_categories).arel.to_sql).to eql(result)
       end
 
+      it 'can be used as the from clause with depth and path' do
+        Category.send(:recursive_auxiliary_statement, :all_categories) do |cte|
+          cte.query Category.all
+          cte.with_depth
+          cte.with_path
+        end
+
+        result = 'WITH RECURSIVE "all_categories" AS ('
+        result << ' SELECT "categories".*, 0 AS depth, ARRAY["categories"."id"]::varchar[] AS path'
+        result << ' FROM "categories"'
+        result << ' WHERE "categories"."parent_id" IS NULL'
+        result << ' UNION'
+        result << ' SELECT "categories".*, ("all_categories"."depth" + 1) AS depth, array_append("all_categories"."path", "categories"."id"::varchar) AS path'
+        result << ' FROM "categories", "all_categories"'
+        result << ' WHERE "categories"."parent_id" = "all_categories"."id"'
+        result << ' ) SELECT "categories".* FROM "all_categories" AS categories'
+        expect(Category.from(:all_categories, :categories).arel.to_sql).to eql(result)
+      end
+
       it 'works with string queries' do
         klass.send(:recursive_auxiliary_statement, :all_categories) do |cte|
           cte.query 'SELECT * FROM categories WHERE a IS NULL'