Adds Arel::Nodes::Cte for use in WITH expressions

SelectManager#with currently accepts As and TableAlias nodes.
Neither of these support materialization hints for the query
planner. Both Postgres and SQLite support such hints.

This commit adds a Cte node that does support materialization
hints. It continues to support As and TableAlias nodes by
translating them into Cte nodes.

activerecord/CHANGELOG.md

@@ -1,3 +1,26 @@
+*   Add `Arel::Nodes::Cte` for use in `WITH` expressions
+
+    `Cte` nodes also support explicit `MATERIALIZED` or `NOT MATERIALIZED` hints for the query planner.
+
+    ```ruby
+    posts = Arel::Table.new(:posts)
+    comments = Arel::Table.new(:comments)
+
+    good_comments_query = comments.project(Arel.star).where(comments[:rating].gt(7))
+    cte = Arel::Nodes::Cte.new(:good_comments, good_comments_query, materialized: true)
+
+    query = posts.
+      project(Arel.star).
+      with(cte).
+      where(posts[:id].in(cte.to_table.project(:post_id))).
+
+    puts query.to_sql
+
+    # WITH "good_comments" AS MATERIALIZED (SELECT * FROM "comments" WHERE "comments"."rating" > 7) SELECT * FROM "posts" WHERE "posts"."id" IN (SELECT post_id FROM "good_comments")
+    ```
+
+    *Jon Zeppieri*
+
 *   Fix mutation detection for serialized attributes backed by binary columns.
 
     *Jean Boussier*

activerecord/lib/arel/nodes.rb

@@ -39,6 +39,7 @@
 require "arel/nodes/over"
 require "arel/nodes/matches"
 require "arel/nodes/regexp"
+require "arel/nodes/cte"
 
 # nary
 require "arel/nodes/and"

activerecord/lib/arel/nodes/binary.rb

@@ -39,6 +39,12 @@ def fetch_attribute
       end
     end
 
+    class As < Binary
+      def to_cte
+        Arel::Nodes::Cte.new(left.name, right)
+      end
+    end
+
     class Between < Binary; include FetchAttribute; end
 
     class GreaterThan < Binary
@@ -112,7 +118,6 @@ def fetch_attribute(&block)
     end
 
     %w{
-      As
       Assignment
       Join
       Union

activerecord/lib/arel/nodes/cte.rb

@@ -0,0 +1,36 @@
+# frozen_string_literal: true
+
+module Arel # :nodoc: all
+  module Nodes
+    class Cte < Arel::Nodes::Binary
+      alias :name :left
+      alias :relation :right
+      attr_reader :materialized
+
+      def initialize(name, relation, materialized: nil)
+        super(name, relation)
+        @materialized = materialized
+      end
+
+      def hash
+        [name, relation, materialized].hash
+      end
+
+      def eql?(other)
+        self.class == other.class &&
+          self.name == other.name &&
+          self.relation == other.relation &&
+          self.materialized == other.materialized
+      end
+      alias :== :eql?
+
+      def to_cte
+        self
+      end
+
+      def to_table
+        Arel::Table.new(name)
+      end
+    end
+  end
+end

activerecord/lib/arel/nodes/table_alias.rb

@@ -26,6 +26,10 @@ def type_for_attribute(name)
       def able_to_type_cast?
         relation.respond_to?(:able_to_type_cast?) && relation.able_to_type_cast?
       end
+
+      def to_cte
+        Arel::Nodes::Cte.new(name, relation)
+      end
     end
   end
 end

activerecord/lib/arel/visitors/mysql.rb

@@ -64,6 +64,12 @@ def visit_Arel_Nodes_NullsFirst(o, collector)
           visit o.expr, collector
         end
 
+        def visit_Arel_Nodes_Cte(o, collector)
+          collector << quote_table_name(o.name)
+          collector << " AS "
+          visit o.relation, collector
+        end
+
         # In the simple case, MySQL allows us to place JOINs directly into the UPDATE
         # query. However, this does not allow for LIMIT, OFFSET and ORDER. To support
         # these, we must use a subquery.

activerecord/lib/arel/visitors/to_sql.rb

@@ -734,6 +734,20 @@ def visit_Arel_Nodes_UnqualifiedColumn(o, collector)
           collector << quote_column_name(o.name)
         end
 
+        def visit_Arel_Nodes_Cte(o, collector)
+          collector << quote_table_name(o.name)
+          collector << " AS "
+
+          case o.materialized
+          when true
+            collector << "MATERIALIZED "
+          when false
+            collector << "NOT MATERIALIZED "
+          end
+
+          visit o.relation, collector
+        end
+
         def visit_Arel_Attributes_Attribute(o, collector)
           join_name = o.relation.table_alias || o.relation.name
           collector << quote_table_name(join_name) << "." << quote_column_name(o.name)
@@ -995,19 +1009,7 @@ def is_distinct_from(o, collector)
         def collect_ctes(children, collector)
           children.each_with_index do |child, i|
             collector << ", " unless i == 0
-
-            case child
-            when Arel::Nodes::As
-              name = child.left.name
-              relation = child.right
-            when Arel::Nodes::TableAlias
-              name = child.name
-              relation = child.relation
-            end
-
-            collector << quote_table_name(name)
-            collector << " AS "
-            visit relation, collector
+            visit child.to_cte, collector
           end
 
           collector

activerecord/test/cases/arel/nodes/as_test.rb

@@ -31,6 +31,18 @@ module Nodes
           assert_equal 2, array.uniq.size
         end
       end
+
+      describe "#to_cte" do
+        it "returns a Cte node using the LHS's name and the RHS as the relation" do
+          table = Table.new(:users)
+          as_node = As.new(table, "foo")
+          cte_node = as_node.to_cte
+
+          assert_kind_of Arel::Nodes::Cte, cte_node
+          assert_equal as_node.left.name, cte_node.name
+          assert_equal as_node.right, cte_node.relation
+        end
+      end
     end
   end
 end

activerecord/test/cases/arel/nodes/cte_test.rb

@@ -0,0 +1,46 @@
+# frozen_string_literal: true
+
+require_relative "../helper"
+
+module Arel
+  module Nodes
+    describe "Cte" do
+      describe "equality" do
+        it "is equal with equal ivars" do
+          array = [
+            Cte.new("foo", "bar", materialized: true),
+            Cte.new("foo", "bar", materialized: true)
+          ]
+
+          assert_equal 1, array.uniq.size
+        end
+
+        it "is not equal with unequal ivars" do
+          array = [
+            Cte.new("foo", "bar", materialized: true),
+            Cte.new("foo", "bar")
+          ]
+
+          assert_equal 2, array.uniq.size
+        end
+      end
+
+      describe "#to_cte" do
+        it "returns self" do
+          cte = Cte.new("foo", "bar")
+
+          assert_equal cte.to_cte, cte
+        end
+      end
+
+      describe "#to_table" do
+        it "returns an Arel::Table using the Cte's name" do
+          table = Cte.new("foo", "bar").to_table
+
+          assert_kind_of Arel::Table, table
+          assert_equal "foo", table.name
+        end
+      end
+    end
+  end
+end

activerecord/test/cases/arel/nodes/table_alias_test.rb

@@ -24,6 +24,18 @@ module Nodes
           assert_equal 2, array.uniq.size
         end
       end
+
+      describe "#to_cte" do
+        it "returns a Cte node using the TableAlias's name and relation" do
+          relation = Table.new(:users).project(Arel.star)
+          table_alias = TableAlias.new(relation, :foo)
+          cte = table_alias.to_cte
+
+          assert_kind_of Arel::Nodes::Cte, cte
+          assert_equal :foo, cte.name
+          assert_equal relation, cte.relation
+        end
+      end
     end
   end
 end

activerecord/test/cases/arel/visitors/mysql_test.rb

@@ -161,6 +161,24 @@ def compile(node)
           }
         end
       end
+
+      describe "Nodes::Cte" do
+        it "ignores MATERIALIZED modifiers" do
+          cte = Nodes::Cte.new("foo", Table.new(:bar).project(Arel.star), materialized: true)
+
+          _(compile(cte)).must_be_like %{
+            "foo" AS (SELECT * FROM "bar")
+          }
+        end
+
+        it "ignores NOT MATERIALIZED modifiers" do
+          cte = Nodes::Cte.new("foo", Table.new(:bar).project(Arel.star), materialized: false)
+
+          _(compile(cte)).must_be_like %{
+            "foo" AS (SELECT * FROM "bar")
+          }
+        end
+      end
     end
   end
 end

activerecord/test/cases/arel/visitors/to_sql_test.rb

@@ -862,6 +862,19 @@ def dispatch
             WITH expr1 AS (SELECT * FROM "bar"), expr2 AS (SELECT * FROM "baz") SELECT * FROM expr2
           }
         end
+
+        it "handles Cte nodes" do
+          cte = Arel::Nodes::Cte.new("expr1", Table.new(:bar).project(Arel.star))
+          manager = Table.new(:foo).
+            project(Arel.star).
+            with(cte).
+            from(cte.to_table).
+            where(cte.to_table[:score].gt(5))
+
+          _(compile(manager.ast)).must_be_like %{
+            WITH "expr1" AS (SELECT * FROM "bar") SELECT * FROM "expr1" WHERE "expr1"."score" > 5
+          }
+        end
       end
 
       describe "Nodes::WithRecursive" do
@@ -876,6 +889,32 @@ def dispatch
         end
       end
 
+      describe "Nodes::Cte" do
+        it "handles CTEs with no MATERIALIZED modifier" do
+          cte = Nodes::Cte.new("foo", Table.new(:bar).project(Arel.star))
+
+          _(compile(cte)).must_be_like %{
+            "foo" AS (SELECT * FROM "bar")
+          }
+        end
+
+        it "handles CTEs with a MATERIALIZED modifier" do
+          cte = Nodes::Cte.new("foo", Table.new(:bar).project(Arel.star), materialized: true)
+
+          _(compile(cte)).must_be_like %{
+            "foo" AS MATERIALIZED (SELECT * FROM "bar")
+          }
+        end
+
+        it "handles CTEs with a NOT MATERIALIZED modifier" do
+          cte = Nodes::Cte.new("foo", Table.new(:bar).project(Arel.star), materialized: false)
+
+          _(compile(cte)).must_be_like %{
+            "foo" AS NOT MATERIALIZED (SELECT * FROM "bar")
+          }
+        end
+      end
+
       describe "Nodes::Fragments" do
         it "joins subexpressions" do
           sql = Arel.sql("SELECT foo, bar") + Arel.sql(" FROM customers")