Truncate: Correctly handle UPDATE and ON CONFLICT target lists

These target list fields support MultiAssignRef nodes, which are not supported in regular SELECT target lists - thus requiring special handling.
pganalyze · Apr 4, 2022 · 53e39cb · 53e39cb
1 parent 825e15f
commit 53e39cb
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 7 deletions.
diff --git a/lib/pg_query/truncate.rb b/lib/pg_query/truncate.rb
@@ -64,7 +64,11 @@ def find_possible_truncations # rubocop:disable Metrics/CyclomaticComplexity
         case k
         when :target_list
           next unless node.is_a?(PgQuery::SelectStmt) || node.is_a?(PgQuery::UpdateStmt) || node.is_a?(PgQuery::OnConflictClause)
-          length = PgQuery.deparse_stmt(PgQuery::SelectStmt.new(k => v.to_a, op: :SETOP_NONE)).size - 7 # 'SELECT '.size
+          length = if node.is_a?(PgQuery::SelectStmt)
+                     select_target_list_len(v)
+                   else # UpdateStmt / OnConflictClause
+                     update_target_list_len(v)
+                   end
           truncations << PossibleTruncation.new(location, :target_list, length, true)
         when :where_clause
           next unless node.is_a?(PgQuery::SelectStmt) || node.is_a?(PgQuery::UpdateStmt) || node.is_a?(PgQuery::DeleteStmt) ||
@@ -82,18 +86,23 @@ def find_possible_truncations # rubocop:disable Metrics/CyclomaticComplexity
           truncations << PossibleTruncation.new(location, :ctequery, length, false)
         when :cols
           next unless node.is_a?(PgQuery::InsertStmt)
-          length = PgQuery.deparse_stmt(
-            PgQuery::InsertStmt.new(
-              relation: PgQuery::RangeVar.new(relname: 'x', inh: true),
-              cols: v.to_a
-            )
-          ).size - 31 # "INSERT INTO x () DEFAULT VALUES".size
+          length = cols_len(v)
           truncations << PossibleTruncation.new(location, :cols, length, true)
         end
       end
 
       truncations
     end
+
+    def select_target_list_len(target_list)
+      deparsed_len = PgQuery.deparse_stmt(
+        PgQuery::SelectStmt.new(
+          target_list: target_list.to_a, op: :SETOP_NONE
+        )
+      ).size
+      deparsed_len - 7 # 'SELECT '.size
+    end
+
     def select_values_lists_len(values_lists)
       deparsed_len = PgQuery.deparse_stmt(
         PgQuery::SelectStmt.new(
@@ -103,5 +112,24 @@ def select_values_lists_len(values_lists)
       deparsed_len - 7 # 'SELECT '.size
     end
 
+    def update_target_list_len(target_list)
+      deparsed_len = PgQuery.deparse_stmt(
+        PgQuery::UpdateStmt.new(
+          target_list: target_list.to_a,
+          relation: PgQuery::RangeVar.new(relname: 'x', inh: true)
+        )
+      ).size
+      deparsed_len - 13 # 'UPDATE x SET '.size
+    end
+
+    def cols_len(cols)
+      deparsed_len = PgQuery.deparse_stmt(
+        PgQuery::InsertStmt.new(
+          relation: PgQuery::RangeVar.new(relname: 'x', inh: true),
+          cols: cols.to_a
+        )
+      ).size
+      deparsed_len - 31 # "INSERT INTO x () DEFAULT VALUES".size
+    end
   end
 end
diff --git a/spec/lib/truncate_spec.rb b/spec/lib/truncate_spec.rb
@@ -45,6 +45,10 @@
     query = 'INSERT INTO y(a) VALUES(1) ON CONFLICT DO UPDATE SET a = 123456789'
     expect(described_class.parse(query).truncate(66)).to eq 'INSERT INTO y (a) VALUES (...) ON CONFLICT DO UPDATE SET ... = ...'
   end
+
+  it 'handles complex ON CONFLICT target lists' do
+    query = 'INSERT INTO foo (a, b, c, d) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18,$19,$20,$21,$22,$23,$24,$25,$26,$27,$28,$29) ON CONFLICT (id) DO UPDATE SET (a, b, c, d) = (excluded.a,excluded.b,excluded.c,case when foo.d = excluded.d then excluded.d end)'
+    expect(described_class.parse(query).truncate(100)).to eq 'INSERT INTO foo (a, b, c, d) VALUES (...) ON CONFLICT (id) DO UPDATE SET ... = ...'
   end
 
   it 'handles GRANT access privileges' do