Merge pull request rails#37944 from vlado/with_cte

Common Table Expression support added "out-of-the-box"

Author: byroot

activerecord/CHANGELOG.md

@@ -1,3 +1,13 @@
+*   `.with` query method added. Construct common table expressions with ease and get `ActiveRecord::Relation` back.
+
+    ```ruby
+    Post.with(posts_with_comments: Post.where("comments_count > ?", 0))
+    # => ActiveRecord::Relation
+    # WITH posts_with_comments AS (SELECT * FROM posts WHERE (comments_count > 0)) SELECT * FROM posts
+    ```
+
+    *Vlado Cingel*
+
 *   Don't establish a new connection if an identical pool exists already.
 
     Previously, if `establish_connection` was called on a class that already had an established connection, the existing connection would be removed regardless of whether it was the same config. Now if a pool is found with the same values as the new connection, the existing connection will be returned instead of creating a new one.

activerecord/lib/active_record/querying.rb

@@ -17,7 +17,7 @@ module Querying
       :and, :or, :annotate, :optimizer_hints, :extending,
       :having, :create_with, :distinct, :references, :none, :unscope, :merge, :except, :only,
       :count, :average, :minimum, :maximum, :sum, :calculate,
-      :pluck, :pick, :ids, :strict_loading, :excluding, :without,
+      :pluck, :pick, :ids, :strict_loading, :excluding, :without, :with,
       :async_count, :async_average, :async_minimum, :async_maximum, :async_sum, :async_pluck, :async_pick,
     ].freeze # :nodoc:
     delegate(*QUERYING_METHODS, to: :all)

activerecord/lib/active_record/relation.rb

@@ -5,13 +5,14 @@ module ActiveRecord
   class Relation
     MULTI_VALUE_METHODS  = [:includes, :eager_load, :preload, :select, :group,
                             :order, :joins, :left_outer_joins, :references,
-                            :extending, :unscope, :optimizer_hints, :annotate]
+                            :extending, :unscope, :optimizer_hints, :annotate,
+                            :with]
 
     SINGLE_VALUE_METHODS = [:limit, :offset, :lock, :readonly, :reordering, :strict_loading,
                             :reverse_order, :distinct, :create_with, :skip_query_cache]
 
     CLAUSE_METHODS = [:where, :having, :from]
-    INVALID_METHODS_FOR_DELETE_ALL = [:distinct]
+    INVALID_METHODS_FOR_DELETE_ALL = [:distinct, :with]
 
     VALUE_METHODS = MULTI_VALUE_METHODS + SINGLE_VALUE_METHODS + CLAUSE_METHODS
 

activerecord/lib/active_record/relation/query_methods.rb

@@ -324,6 +324,63 @@ def _select!(*fields) # :nodoc:
       self
     end
 
+    # Add a Common Table Expression (CTE) that you can then reference within another SELECT statement.
+    #
+    #   Post.with(posts_with_tags: Post.where("tags_count > ?", 0))
+    #   # => ActiveRecord::Relation
+    #   # WITH posts_with_tags AS (
+    #   #   SELECT * FROM posts WHERE (tags_count > 0)
+    #   # )
+    #   # SELECT * FROM posts
+    #
+    # Once you define Common Table Expression you can use custom `FROM` value or `JOIN` to reference it.
+    #
+    #   Post.with(posts_with_tags: Post.where("tags_count > ?", 0)).from("posts_with_tags AS posts")
+    #   # => ActiveRecord::Relation
+    #   # WITH posts_with_tags AS (
+    #   #  SELECT * FROM posts WHERE (tags_count > 0)
+    #   # )
+    #   # SELECT * FROM posts_with_tags AS posts
+    #
+    #   Post.with(posts_with_tags: Post.where("tags_count > ?", 0)).joins("JOIN posts_with_tags ON posts_with_tags.id = posts.id")
+    #   # => ActiveRecord::Relation
+    #   # WITH posts_with_tags AS (
+    #   #   SELECT * FROM posts WHERE (tags_count > 0)
+    #   # )
+    #   # SELECT * FROM posts JOIN posts_with_tags ON posts_with_tags.id = posts.id
+    #
+    # It is recommended to pass a query as `ActiveRecord::Relation`. If that is not possible
+    # and you have verified it is safe for the database, you can pass it as SQL literal
+    # using `Arel`.
+    #
+    #   Post.with(popular_posts: Arel.sql("... complex sql to calculate posts popularity ..."))
+    #
+    # Great caution should be taken to avoid SQL injection vulnerabilities. This method should not
+    # be used with unsafe values that include unsanitized input.
+    #
+    # To add multiple CTEs just pass multiple key-value pairs
+    #
+    #   Post.with(
+    #     posts_with_comments: Post.where("comments_count > ?", 0),
+    #     posts_with_tags: Post.where("tags_count > ?", 0)
+    #   )
+    #
+    # or chain multiple `.with` calls
+    #
+    #   Post
+    #     .with(posts_with_comments: Post.where("comments_count > ?", 0))
+    #     .with(posts_with_tags: Post.where("tags_count > ?", 0))
+    def with(*args)
+      check_if_method_has_arguments!(__callee__, args)
+      spawn.with!(*args)
+    end
+
+    # Like #with, but modifies relation in place.
+    def with!(*args) # :nodoc:
+      self.with_values += args
+      self
+    end
+
     # Allows you to change a previously set select statement.
     #
     #   Post.select(:title, :body)
@@ -1379,6 +1436,7 @@ def build_arel(aliases = nil)
         arel.group(*arel_columns(group_values.uniq)) unless group_values.empty?
 
         build_order(arel)
+        build_with(arel)
         build_select(arel)
 
         arel.optimizer_hints(*optimizer_hints_values) unless optimizer_hints_values.empty?
@@ -1514,6 +1572,32 @@ def build_select(arel)
         end
       end
 
+      def build_with(arel)
+        return if with_values.empty?
+
+        with_statements = with_values.map do |with_value|
+          raise ArgumentError, "Unsupported argument type: #{with_value} #{with_value.class}" unless with_value.is_a?(Hash)
+
+          build_with_value_from_hash(with_value)
+        end
+
+        arel.with(with_statements)
+      end
+
+      def build_with_value_from_hash(hash)
+        hash.map do |name, value|
+          expression =
+            case value
+            when Arel::Nodes::SqlLiteral then Arel::Nodes::Grouping.new(value)
+            when ActiveRecord::Relation then value.arel
+            when Arel::SelectManager then value
+            else
+              raise ArgumentError, "Unsupported argument type: `#{value}` #{value.class}"
+            end
+          Arel::Nodes::TableAlias.new(expression, name)
+        end
+      end
+
       def arel_columns(columns)
         columns.flat_map do |field|
           case field

activerecord/test/cases/relation/delete_all_test.rb

@@ -70,6 +70,7 @@ def test_delete_all_with_group_by_and_having
 
   def test_delete_all_with_unpermitted_relation_raises_error
     assert_raises(ActiveRecord::ActiveRecordError) { Author.distinct.delete_all }
+    assert_raises(ActiveRecord::ActiveRecordError) { Author.with(limited: Author.limit(2)).delete_all }
   end
 
   def test_delete_all_with_joins_and_where_part_is_hash

activerecord/test/cases/relation/merging_test.rb

@@ -393,4 +393,30 @@ class MergingDifferentRelationsTest < ActiveRecord::TestCase
 
     assert_equal dev.ratings, [rating_1]
   end
+
+  test "merging relation with common table expression" do
+    posts_with_tags = Post.with(posts_with_tags: Post.where("tags_count > 0")).from("posts_with_tags AS posts")
+    posts_with_comments = Post.where("legacy_comments_count > 0")
+    relation = posts_with_comments.merge(posts_with_tags)
+
+    assert_equal [1, 2, 7], relation.pluck(:id)
+  end
+
+  test "merging multiple relations with common table expression" do
+    posts_with_tags = Post.with(posts_with_tags: Post.where("tags_count > 0"))
+    posts_with_comments = Post.with(posts_with_comments: Post.where("legacy_comments_count > 0"))
+    relation = posts_with_comments.merge(posts_with_tags).joins("JOIN posts_with_tags pwt ON pwt.id = posts.id JOIN posts_with_comments pwc ON pwc.id = posts.id")
+
+    assert_equal [1, 2, 7], relation.pluck(:id)
+  end
+
+  test "relation merger leaves to database to decide what to do when multiple CTEs with same alias are passed" do
+    posts_with_tags = Post.with(popular_posts: Post.where("tags_count > 0"))
+    posts_with_comments = Post.with(popular_posts: Post.where("legacy_comments_count > 0"))
+    relation = posts_with_tags.merge(posts_with_comments).joins("JOIN popular_posts pp ON pp.id = posts.id")
+
+    assert_raises ActiveRecord::StatementInvalid do
+      relation.load
+    end
+  end
 end

activerecord/test/cases/relation/with_test.rb

@@ -0,0 +1,64 @@
+# frozen_string_literal: true
+
+require "cases/helper"
+require "models/comment"
+require "models/post"
+
+module ActiveRecord
+  class WithTest < ActiveRecord::TestCase
+    fixtures :comments
+    fixtures :posts
+
+    POSTS_WITH_TAGS = [1, 2, 7, 8, 9, 10, 11].freeze
+    POSTS_WITH_COMMENTS = [1, 2, 4, 5, 7].freeze
+    POSTS_WITH_MULTIPLE_COMMENTS = [1, 4, 5].freeze
+    POSTS_WITH_TAGS_AND_COMMENTS = (POSTS_WITH_COMMENTS & POSTS_WITH_TAGS).sort.freeze
+    POSTS_WITH_TAGS_AND_MULTIPLE_COMMENTS = (POSTS_WITH_MULTIPLE_COMMENTS & POSTS_WITH_TAGS).sort.freeze
+
+    def test_with_when_hash_is_passed_as_an_argument
+      relation = Post
+        .with(posts_with_comments: Post.where("legacy_comments_count > 0"))
+        .from("posts_with_comments AS posts")
+
+      assert_equal POSTS_WITH_COMMENTS, relation.order(:id).pluck(:id)
+    end
+
+    def test_with_when_hash_with_multiple_elements_of_different_type_is_passed_as_an_argument
+      cte_options = {
+        posts_with_tags: Post.arel_table.project(Arel.star).where(Post.arel_table[:tags_count].gt(0)),
+        posts_with_tags_and_comments: Arel.sql("SELECT * FROM posts_with_tags WHERE legacy_comments_count > 0"),
+        "posts_with_tags_and_multiple_comments" => Post.where("legacy_comments_count > 1").from("posts_with_tags_and_comments AS posts")
+      }
+      relation = Post.with(cte_options).from("posts_with_tags_and_multiple_comments AS posts")
+
+      assert_equal POSTS_WITH_TAGS_AND_MULTIPLE_COMMENTS, relation.order(:id).pluck(:id)
+    end
+
+    def test_multiple_with_calls
+      relation = Post
+        .with(posts_with_tags: Post.where("tags_count > 0"))
+        .from("posts_with_tags_and_comments AS posts")
+        .with(posts_with_tags_and_comments: Arel.sql("SELECT * FROM posts_with_tags WHERE legacy_comments_count > 0"))
+
+      assert_equal POSTS_WITH_TAGS_AND_COMMENTS, relation.order(:id).pluck(:id)
+    end
+
+    def test_count_after_with_call
+      relation = Post.with(posts_with_comments: Post.where("legacy_comments_count > 0"))
+
+      assert_equal Post.count, relation.count
+      assert_equal POSTS_WITH_COMMENTS.size, relation.from("posts_with_comments AS posts").count
+      assert_equal POSTS_WITH_COMMENTS.size, relation.joins("JOIN posts_with_comments ON posts_with_comments.id = posts.id").count
+    end
+
+    def test_with_when_called_from_active_record_scope
+      assert_equal POSTS_WITH_TAGS, Post.with_tags_cte.order(:id).pluck(:id)
+    end
+
+    def test_with_when_invalid_params_are_passed
+      assert_raise(ArgumentError) { Post.with }
+      assert_raise(ArgumentError) { Post.with(posts_with_tags: nil).load }
+      assert_raise(ArgumentError) { Post.with(posts_with_tags: [Post.where("tags_count > 0")]).load }
+    end
+  end
+end

activerecord/test/cases/relations_test.rb

@@ -2371,7 +2371,7 @@ def test_find_by_with_take_memoization
     assert_empty authors
   end
 
-  (ActiveRecord::Relation::MULTI_VALUE_METHODS - [:extending]).each do |method|
+  (ActiveRecord::Relation::MULTI_VALUE_METHODS - [:extending, :with]).each do |method|
     test "#{method} with blank value" do
       authors = Author.public_send(method, [""])
       assert_empty authors.public_send(:"#{method}_values")

activerecord/test/models/post.rb

@@ -61,6 +61,7 @@ def first_comment
 
   scope :with_comments, -> { preload(:comments) }
   scope :with_tags, -> { preload(:taggings) }
+  scope :with_tags_cte, -> { with(posts_with_tags: where("tags_count > 0")).from("posts_with_tags AS posts") }
 
   scope :tagged_with, ->(id) { joins(:taggings).where(taggings: { tag_id: id }) }
   scope :tagged_with_comment, ->(comment) { joins(:taggings).where(taggings: { comment: comment }) }