Track tables in EXCEPT and INTERSECT queries

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

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

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 ...'