1 files changed, 62 insertions, 0 deletions
diff --git a/lib/gitlab/sql/recursive_cte.rb b/lib/gitlab/sql/recursive_cte.rb
new file mode 100644
index 00000000000..5b1b03820a3
--- /dev/null
+++ b/lib/gitlab/sql/recursive_cte.rb
@@ -0,0 +1,62 @@
+module Gitlab
+  module SQL
+    # Class for easily building recursive CTE statements.
+    #
+    # Example:
+    #
+    #     cte = RecursiveCTE.new(:my_cte_name)
+    #     ns = Arel::Table.new(:namespaces)
+    #
+    #     cte << Namespace.
+    #       where(ns[:parent_id].eq(some_namespace_id))
+    #
+    #     cte << Namespace.
+    #       from([ns, cte.table]).
+    #       where(ns[:parent_id].eq(cte.table[:id]))
+    #
+    #     Namespace.with.
+    #       recursive(cte.to_arel).
+    #       from(cte.alias_to(ns))
+    class RecursiveCTE
+      attr_reader :table
+
+      # name - The name of the CTE as a String or Symbol.
+      def initialize(name)
+        @table = Arel::Table.new(name)
+        @queries = []
+      end
+
+      # Adds a query to the body of the CTE.
+      #
+      # relation - The relation object to add to the body of the CTE.
+      def <<(relation)
+        @queries << relation
+      end
+
+      # Returns the Arel relation for this CTE.
+      def to_arel
+        sql = Arel::Nodes::SqlLiteral.new(Union.new(@queries).to_sql)
+
+        Arel::Nodes::As.new(table, Arel::Nodes::Grouping.new(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.
+          recursive(to_arel).
+          from(alias_to(relation.model.arel_table))
+      end
+    end
+  end
+end