diff options
-rw-r--r-- | lib/gitlab/sql/cte.rb | 50 | ||||
-rw-r--r-- | spec/lib/gitlab/sql/cte_spec.rb | 42 |
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 |