Parse CTEs and nested selects in INSERT/UPDATE

lib/pg_query/parse.rb

@@ -97,18 +97,23 @@ def load_tables_and_aliases! # rubocop:disable Metrics/CyclomaticComplexity
             statements << statement[SELECT_STMT]['rarg'] if statement[SELECT_STMT]['rarg']
           end
 
-          # CTEs
-          with_clause = statement[SELECT_STMT]['withClause']
-          if with_clause
-            with_clause[WITH_CLAUSE]['ctes'].each do |item|
-              next unless item[COMMON_TABLE_EXPR]
-              @cte_names << item[COMMON_TABLE_EXPR]['ctename']
-              statements << item[COMMON_TABLE_EXPR]['ctequery']
-            end
+          if (with_clause = statement[SELECT_STMT]['withClause'])
+            cte_statements, cte_names = statements_and_cte_names_for_with_clause(with_clause)
+            @cte_names.concat(cte_names)
+            statements.concat(cte_statements)
           end
         # The following statements modify the contents of a table
         when INSERT_STMT, UPDATE_STMT, DELETE_STMT
-          from_clause_items << { item: statement.values[0]['relation'], type: :dml }
+          value = statement.values[0]
+          from_clause_items << { item: value['relation'], type: :dml }
+          statements << value['selectStmt'] if value.key?('selectStmt')
+          statements << value['withClause'] if value.key?('withClause')
+
+          if (with_clause = value['withClause'])
+            cte_statements, cte_names = statements_and_cte_names_for_with_clause(with_clause)
+            @cte_names.concat(cte_names)
+            statements.concat(cte_statements)
+          end
         when COPY_STMT
           from_clause_items << { item: statement.values[0]['relation'], type: :dml } if statement.values[0]['relation']
           statements << statement.values[0]['query']
@@ -221,5 +226,19 @@ def load_tables_and_aliases! # rubocop:disable Metrics/CyclomaticComplexity
     end
 
     @tables.uniq!
+    @cte_names.uniq!
+  end
+
+  def statements_and_cte_names_for_with_clause(with_clause)
+    statements = []
+    cte_names = []
+
+    with_clause[WITH_CLAUSE]['ctes'].each do |item|
+      next unless item[COMMON_TABLE_EXPR]
+      cte_names << item[COMMON_TABLE_EXPR]['ctename']
+      statements << item[COMMON_TABLE_EXPR]['ctequery']
+    end
+
+    [statements, cte_names]
   end
 end

spec/lib/parse_spec.rb

@@ -812,6 +812,64 @@
     expect(query.cte_names).to match_array(['cte_a', 'cte_b'])
   end
 
+  describe 'parsing INSERT' do
+    it 'finds the table inserted into' do
+      query = described_class.parse(<<-SQL)
+        insert into users(pk, name) values (1, 'bob');
+      SQL
+      expect(query.warnings).to be_empty
+      expect(query.tables).to eq(['users'])
+    end
+
+    it 'finds tables in being selected from for insert' do
+      query = described_class.parse(<<-SQL)
+        insert into users(pk, name) select pk, name from other_users;
+      SQL
+      expect(query.warnings).to be_empty
+      expect(query.tables).to match_array(['users', 'other_users'])
+    end
+
+    it 'finds tables in a CTE' do
+      query = described_class.parse(<<-SQL)
+        with cte as (
+          select pk, name from other_users
+        )
+        insert into users(pk, name) select * from cte;
+      SQL
+      expect(query.warnings).to be_empty
+      expect(query.tables).to match_array(['users', 'other_users'])
+    end
+  end
+
+  describe 'parsing UPDATE' do
+    it 'finds the table updateed into' do
+      query = described_class.parse(<<-SQL)
+        update users set name = 'bob';
+      SQL
+      expect(query.warnings).to be_empty
+      expect(query.tables).to eq(['users'])
+    end
+
+    it 'finds tables in a sub-select' do
+      query = described_class.parse(<<-SQL)
+        update users set name = (select name from other_users limit 1);
+      SQL
+      expect(query.warnings).to be_empty
+      expect(query.tables).to match_array(['users', 'other_users'])
+    end
+
+    it 'finds tables in a CTE' do
+      query = described_class.parse(<<-SQL)
+        with cte as (
+          select name from other_users limit 1
+        )
+        update users set name = (select name from cte);
+      SQL
+      expect(query.warnings).to be_empty
+      expect(query.tables).to match_array(['users', 'other_users'])
+    end
+  end
+
   it 'handles DROP TYPE' do
     query = described_class.parse("DROP TYPE IF EXISTS repack.pk_something")
     expect(query.warnings).to eq []