diff --git a/lib/pg_query/parse.rb b/lib/pg_query/parse.rb index 4b6364e5..704bfebe 100644 --- a/lib/pg_query/parse.rb +++ b/lib/pg_query/parse.rb @@ -123,7 +123,7 @@ def load_objects! # rubocop:disable Metrics/CyclomaticComplexity (statement.select_stmt.from_clause || []).each do |item| from_clause_items << { item: item, type: :select } end - when :SETOP_UNION + when :SETOP_UNION, :SETOP_EXCEPT, :SETOP_INTERSECT statements << PgQuery::Node.new(select_stmt: statement.select_stmt.larg) if statement.select_stmt.larg statements << PgQuery::Node.new(select_stmt: statement.select_stmt.rarg) if statement.select_stmt.rarg end diff --git a/spec/lib/parse_spec.rb b/spec/lib/parse_spec.rb index f5e3bda6..e73a1ee4 100644 --- a/spec/lib/parse_spec.rb +++ b/spec/lib/parse_spec.rb @@ -1427,7 +1427,7 @@ expect(query.tables).to match_array ['foo', 'join_a', 'join_b', 'sub_a', 'sub_b', 'sub_c', 'sub_d', 'sub_e', 'sub_f'] end - it 'does not list CTEs as tables after a union select' do + it 'does not list CTEs as tables after a UNION select' do query = described_class.parse(<<-SQL) with cte_a as ( select * from table_a @@ -1445,6 +1445,42 @@ expect(query.cte_names).to match_array(['cte_a', 'cte_b']) end + it 'does not list CTEs as tables after a EXCEPT select' do + query = described_class.parse(<<-SQL) + with cte_a as ( + select * from table_a + ), cte_b as ( + select * from table_b + ) + + select id from table_c + left join cte_b on + table_c.id = cte_b.c_id + except + select * from cte_a + SQL + expect(query.tables).to match_array(['table_a', 'table_b', 'table_c']) + expect(query.cte_names).to match_array(['cte_a', 'cte_b']) + end + + it 'does not list CTEs as tables after a INTERSECT select' do + query = described_class.parse(<<-SQL) + with cte_a as ( + select * from table_a + ), cte_b as ( + select * from table_b + ) + + select id from table_c + left join cte_b on + table_c.id = cte_b.c_id + intersect + select * from cte_a + SQL + expect(query.tables).to match_array(['table_a', 'table_b', 'table_c']) + expect(query.cte_names).to match_array(['cte_a', 'cte_b']) + end + it 'finds tables inside subselects in MIN/MAX and COALESCE functions' do query = described_class.parse(<<-SQL) SELECT GREATEST( @@ -1724,5 +1760,27 @@ expect(query.ddl_tables).to eq(['foo']) expect(query.select_tables).to eq(['bar', 'baz']) end + + it 'finds tables in the subquery with EXCEPT' do + query = described_class.parse(<<-SQL) + CREATE TABLE foo AS + SELECT id FROM bar EXCEPT SELECT id from baz; + SQL + + expect(query.tables).to eq(['foo', 'bar', 'baz']) + expect(query.ddl_tables).to eq(['foo']) + expect(query.select_tables).to eq(['bar', 'baz']) + end + + it 'finds tables in the subquery with INTERSECT' do + query = described_class.parse(<<-SQL) + CREATE TABLE foo AS + SELECT id FROM bar INTERSECT SELECT id from baz; + SQL + + expect(query.tables).to eq(['foo', 'bar', 'baz']) + expect(query.ddl_tables).to eq(['foo']) + expect(query.select_tables).to eq(['bar', 'baz']) + end end end diff --git a/spec/lib/truncate_spec.rb b/spec/lib/truncate_spec.rb index d944fe32..58651c3b 100644 --- a/spec/lib/truncate_spec.rb +++ b/spec/lib/truncate_spec.rb @@ -21,6 +21,11 @@ expect(described_class.parse(query).truncate(32)).to eq 'INSERT INTO x (...) VALUES (?)' end + it 'omits comments' do + query = 'SELECT $1 /* application:test */' + expect(described_class.parse(query).truncate(100)).to eq 'SELECT $1' + end + it 'performs a simple truncation if necessary' do query = 'SELECT * FROM t' expect(described_class.parse(query).truncate(10)).to eq 'SELECT ...'