Add Gitlab::SQL:CTE for easily building CTE statements

2 files changed, 92 insertions, 0 deletions

diff --git a/lib/gitlab/sql/cte.rb b/lib/gitlab/sql/cte.rb
new file mode 100644
index 00000000000..f357829ba3f
--- /dev/null
+++ b/lib/gitlab/sql/cte.rb
@@ -0,0 +1,50 @@
+module Gitlab
+  module SQL
+    # Class for easily building CTE statements.
+    #
+    # Example:
+    #
+    #     cte = CTE.new(:my_cte_name)
+    #     ns = Arel::Table.new(:namespaces)
+    #
+    #     cte << Namespace.
+    #       where(ns[:parent_id].eq(some_namespace_id))
+    #
+    #     Namespace
+    #       with(cte.to_arel).
+    #       from(cte.alias_to(ns))
+    class CTE
+      attr_reader :table, :query
+
+      # name - The name of the CTE as a String or Symbol.
+      def initialize(name, query)
+        @table = Arel::Table.new(name)
+        @query = query
+      end
+
+      # Returns the Arel relation for this CTE.
+      def to_arel
+        sql = Arel::Nodes::SqlLiteral.new("(#{query.to_sql})")
+
+        Arel::Nodes::As.new(table, sql)
+      end
+
+      # Returns an "AS" statement that aliases the CTE name as the given table
+      # name. This allows one to trick ActiveRecord into thinking it's selecting
+      # from an actual table, when in reality it's selecting from a CTE.
+      #
+      # alias_table - The Arel table to use as the alias.
+      def alias_to(alias_table)
+        Arel::Nodes::As.new(table, alias_table)
+      end
+
+      # Applies the CTE to the given relation, returning a new one that will
+      # query from it.
+      def apply_to(relation)
+        relation.except(:where)
+          .with(to_arel)
+          .from(alias_to(relation.model.arel_table))
+      end
+    end
+  end
+end
diff --git a/spec/lib/gitlab/sql/cte_spec.rb b/spec/lib/gitlab/sql/cte_spec.rb
new file mode 100644
index 00000000000..d6763c7b2e1
--- /dev/null
+++ b/spec/lib/gitlab/sql/cte_spec.rb
@@ -0,0 +1,42 @@
+require 'spec_helper'
+
+describe Gitlab::SQL::CTE, :postgresql do
+  describe '#to_arel' do
+    it 'generates an Arel relation for the CTE body' do
+      relation = User.where(id: 1)
+      cte = described_class.new(:cte_name, relation)
+      sql = cte.to_arel.to_sql
+      name = ActiveRecord::Base.connection.quote_table_name(:cte_name)
+
+      sql1 = ActiveRecord::Base.connection.unprepared_statement do
+        relation.except(:order).to_sql
+      end
+
+      expect(sql).to eq("#{name} AS (#{sql1})")
+    end
+  end
+
+  describe '#alias_to' do
+    it 'returns an alias for the CTE' do
+      cte = described_class.new(:cte_name, nil)
+      table = Arel::Table.new(:kittens)
+
+      source_name = ActiveRecord::Base.connection.quote_table_name(:cte_name)
+      alias_name = ActiveRecord::Base.connection.quote_table_name(:kittens)
+
+      expect(cte.alias_to(table).to_sql).to eq("#{source_name} AS #{alias_name}")
+    end
+  end
+
+  describe '#apply_to' do
+    it 'applies a CTE to an ActiveRecord::Relation' do
+      user = create(:user)
+      cte = described_class.new(:cte_name, User.where(id: user.id))
+
+      relation = cte.apply_to(User.all)
+
+      expect(relation.to_sql).to match(/WITH .+cte_name/)
+      expect(relation.to_a).to eq(User.where(id: user.id).to_a)
+    end
+  end
+end