Add Gitlab::Database::SetAll helper

This helps write queries that use CTE's to update several rows at once,
setting arbitrary columns to arbitrary values.

app/models/concerns/relative_positioning.rb

@@ -103,31 +103,20 @@ module RelativePositioning
       indexed = (at_end ? objects : objects.reverse).each_with_index
 
       # Some classes are polymorphic, and not all siblings are in the same table.
-      by_model = indexed.group_by { |pair| pair.first.class }
+      by_model = indexed.group_by { |pair| pair.first.model_class }
       lower_bound, upper_bound = at_end ? [position, MAX_POSITION] : [MIN_POSITION, position]
 
       by_model.each do |model, pairs|
         model.transaction do
           pairs.each_slice(100) do |batch|
-            # These are known to be integers, one from the DB, and the other
-            # calculated by us, and thus safe to interpolate
-            values = batch.map do |obj, i|
+            mapping = batch.to_h.transform_values! do |i|
               desired_pos = position + delta * (i + 1)
-              pos = desired_pos.clamp(lower_bound, upper_bound)
-              obj.relative_position = pos
-              "(#{obj.id}, #{pos})"
-            end.join(', ')
-
-            model.connection.exec_query(<<~SQL, "UPDATE #{model.table_name} positions")
-              WITH cte(cte_id, new_pos) AS (
-               SELECT *
-               FROM (VALUES #{values}) as t (id, pos)
-              )
-              UPDATE #{model.table_name}
-              SET relative_position = cte.new_pos
-              FROM cte
-              WHERE cte_id = id
-            SQL
+              { relative_position: desired_pos.clamp(lower_bound, upper_bound) }
+            end
+
+            ::Gitlab::Database::SetAll::Setter
+              .new(model, [:relative_position], mapping)
+              .update!
           end
         end
       end
@@ -206,4 +195,10 @@ module RelativePositioning
   def reset_relative_position
     reset.relative_position
   end
+
+  # Override if the model class needs a more complicated computation (e.g. the
+  # object is a member of a union).
+  def model_class
+    self.class
+  end
 end

ee/app/models/concerns/epic_tree_sorting.rb

@@ -106,5 +106,14 @@ module EpicTreeSorting
 
       [type, id]
     end
+
+    override :model_class
+    def model_class
+      type = try(:object_type)
+
+      return type.camelcase.constantize if type
+
+      super
+    end
   end
 end

lib/gitlab/database/set_all.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    # Constructs queries of the form:
+    #
+    #   with cte(a, b, c) as (
+    #     select * from (values (:x, :y, :z), (:q, :r, :s)) as t
+    #     )
+    #   update table set b = cte.b, c = cte.c where a = cte.a
+    #
+    # Which is useful if you want to update a set of records in a single query
+    # but cannot express the update as a calculation (i.e. you have arbitrary
+    # updates to perform).
+    #
+    # The requirements are that the table must have an ID column used to
+    # identify the rows to be updated.
+    module SetAll
+      COMMA = ', '
+
+      class Setter
+        include Gitlab::Utils::StrongMemoize
+
+        attr_reader :table_name, :connection, :columns, :mapping
+
+        def initialize(model, columns, mapping)
+          raise ArgumentError if columns.blank? || columns.any? { |c| !c.is_a?(Symbol) }
+          raise ArgumentError if mapping.nil? || mapping.empty?
+          raise ArgumentError if mapping.any? { |_k, v| !v.is_a?(Hash) }
+
+          @table_name = model.table_name
+          @connection = model.connection
+          @columns = ([:id] + columns).map { |c| [c, model.column_for_attribute(c)] }
+          @mapping = mapping
+        end
+
+        def params
+          mapping.flat_map do |k, v|
+            obj_id = k.try(:id) || k
+            v = v.merge(id: obj_id)
+            columns.map { |c| query_attribute(c, k, v) }
+          end
+        end
+
+        # A workaround for https://github.com/rails/rails/issues/24893
+        # We need to detect if prepared statements have been disabled.
+        def no_prepared_statement?
+          strong_memoize(:no_prepared_statement) do
+            connection.send(:without_prepared_statement?, [1]) # rubocop: disable GitlabSecurity/PublicSend
+          end
+        end
+
+        def query_attribute(column, key, values)
+          column_name = column.first
+          value = values[column_name]
+          key[column_name] = value if key.try(:id) # optimistic update
+          ActiveRecord::Relation::QueryAttribute.from_user(nil, value, ActiveModel::Type.lookup(column.second.type))
+        end
+
+        def values
+          counter = 0
+          typed = false
+
+          mapping.map do |k, v|
+            binds = columns.map do |c|
+              bind = "$#{counter += 1}"
+              # PG is not great at inferring types - help it for the first row.
+              bind += "::#{c.second.sql_type}" unless typed
+              bind
+            end
+            typed = true
+
+            "(#{binds.join(COMMA)})"
+          end
+        end
+
+        def sql
+          column_names = columns.map(&:first)
+          cte_columns = column_names.map do |c|
+            connection.quote_column_name("cte_#{c}")
+          end
+          updates = column_names.zip(cte_columns).drop(1).map do |dest, src|
+            "#{connection.quote_column_name(dest)} = cte.#{src}"
+          end
+
+          <<~SQL
+            WITH cte(#{cte_columns.join(COMMA)}) AS (VALUES #{values.join(COMMA)})
+            UPDATE #{table_name} SET #{updates.join(COMMA)} FROM cte WHERE cte_id = id
+          SQL
+        end
+
+        def update!
+          log_name = "SetAll #{table_name} #{columns.drop(1).map(&:first)}:#{mapping.size}"
+          if no_prepared_statement?
+            # A workaround for https://github.com/rails/rails/issues/24893
+            # When prepared statements are prevented (such as when using the
+            # query counter or in omnibus by default), we cannot call
+            # `exec_update`, since that will discard the bindings.
+            connection.send(:exec_no_cache, sql, log_name, params) # rubocop: disable GitlabSecurity/PublicSend
+          else
+            connection.exec_update(sql, log_name, params)
+          end
+        end
+      end
+
+      def self.set_all(columns, mapping)
+        mapping.group_by { |k, v| k.class }.each do |model, entries|
+          Setter.new(model, columns, entries).update!
+        end
+      end
+    end
+  end
+end

spec/lib/gitlab/database/set_all_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::SetAll do
+  it 'is possible to update all objects in a single query' do
+    users = create_list(:user, 3)
+    mapping = users.zip(%w(foo bar baz)).to_h do |u, name|
+      [u, { username: name, admin: true }]
+    end
+
+    expect do
+      described_class.set_all(%i[username admin], mapping)
+    end.not_to exceed_query_limit(1)
+
+    # We have optimistically updated the values
+    expect(users).to all(be_admin)
+    expect(users.map(&:username)).to eq(%w(foo bar baz))
+
+    users.each(&:reset)
+
+    # The values are correct on reset
+    expect(users).to all(be_admin)
+    expect(users.map(&:username)).to eq(%w(foo bar baz))
+  end
+
+  it 'is possible to update heterogeneous sets' do
+    create_default(:user)
+    create_default(:project)
+
+    mr_a = create(:merge_request)
+    i_a, i_b = create_list(:issue, 2)
+
+    mapping = {
+      mr_a => { title: 'MR a' },
+      i_a  => { title: 'Issue a' },
+      i_b  => { title: 'Issue b' }
+    }
+
+    expect do
+      described_class.set_all(%i[title], mapping)
+    end.not_to exceed_query_limit(2)
+
+    expect([mr_a, i_a, i_b].map { |x| x.reset.title })
+      .to eq(['MR a', 'Issue a', 'Issue b'])
+  end
+
+  it 'works when not using a query-counter' do
+    create_default(:user)
+    create_default(:project)
+
+    i_a, i_b = create_list(:issue, 2)
+
+    mapping = {
+      i_a  => { title: 'Issue a' },
+      i_b  => { title: 'Issue b' }
+    }
+
+    described_class.set_all(%i[title], mapping)
+
+    expect([i_a, i_b].map { |x| x.reset.title })
+      .to eq(['Issue a', 'Issue b'])
+  end
+end

spec/lib/gitlab/relative_positioning/mover_spec.rb

@@ -37,18 +37,11 @@ RSpec.describe RelativePositioning::Mover do
   end
 
   def set_positions(positions)
-    vals = issues.zip(positions).map do |issue, pos|
-      issue.relative_position = pos
-      "(#{issue.id}, #{pos})"
-    end.join(', ')
-
-    Issue.connection.exec_query(<<~SQL, 'set-positions')
-      WITH cte(cte_id, new_pos) AS (
-       SELECT * FROM (VALUES #{vals}) as t (id, pos)
-      )
-      UPDATE issues SET relative_position = new_pos FROM cte WHERE id = cte_id
-      ;
-    SQL
+    mapping = issues.zip(positions).to_h do |issue, pos|
+      [issue, { relative_position: pos }]
+    end
+
+    ::Gitlab::Database::SetAll.set_all([:relative_position], mapping)
   end
 
   def ids_in_position_order

spec/support/shared_examples/models/relative_positioning_shared_examples.rb

@@ -152,9 +152,26 @@ RSpec.shared_examples 'a class that supports relative positioning' do
       expect(bunch.map(&:relative_position)).to all(be < nils.map(&:relative_position).min)
     end
 
+    it 'manages to move nulls found in the relative scope' do
+      nils = create_items_with_positions([nil] * 4)
+
+      described_class.move_nulls_to_end(sibling_query.to_a)
+      positions = nils.map { |item| item.reset.relative_position }
+
+      expect(positions).to all(be_present)
+      expect(positions).to all(be_valid_position)
+    end
+
+    it 'can move many nulls' do
+      nils = create_items_with_positions([nil] * 101)
+
+      described_class.move_nulls_to_end(nils)
+
+      expect(nils.map(&:relative_position)).to all(be_valid_position)
+    end
+
     it 'does not have an N+1 issue' do
       create_items_with_positions(10..12)
-
       a, b, c, d, e, f, *xs = create_items_with_positions([nil] * 10)
 
       baseline = ActiveRecord::QueryRecorder.new do