Commit 71a5dc94 authored by Alex Kalderimis's avatar Alex Kalderimis

Changes suggested during maintainer review

parent a693f3c2
...@@ -94,3 +94,10 @@ mapping = objects.to_h { |obj| [obj, bazzes[obj.id] } ...@@ -94,3 +94,10 @@ mapping = objects.to_h { |obj| [obj, bazzes[obj.id] }
obj.object_type.constantize obj.object_type.constantize
end end
``` ```
## Caveats
Note that this is a **very low level** tool, and operates on the raw column
values. Enumerations and state fields must be translated into their underlying
representations, for example, and nested associations are not supported. No
validations or hooks will be called.
...@@ -30,22 +30,57 @@ module Gitlab ...@@ -30,22 +30,57 @@ module Gitlab
# representations, for example, and no hooks will be called. # representations, for example, and no hooks will be called.
# #
module SetAll module SetAll
COMMA = ', ' LIST_SEPARATOR = ', '
class Setter class Setter
include Gitlab::Utils::StrongMemoize include Gitlab::Utils::StrongMemoize
attr_reader :table_name, :connection, :columns, :mapping
def initialize(model, 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 @table_name = model.table_name
@connection = model.connection @connection = model.connection
@columns = ([:id] + columns).map { |c| model.column_for_attribute(c) } @columns = self.class.column_definitions(model, columns)
@mapping = mapping @mapping = self.class.value_mapping(mapping)
end
def update!
if without_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
def self.column_definitions(model, columns)
raise ArgumentError, 'invalid columns' if columns.blank? || columns.any? { |c| !c.is_a?(Symbol) }
raise ArgumentError, 'cannot set ID' if columns.include?(:id)
([:id] | columns).map do |name|
definition = model.column_for_attribute(name)
raise ArgumentError, "Unknown column: #{name}" unless definition.type
definition
end
end
def self.value_mapping(mapping)
raise ArgumentError, 'invalid mapping' if mapping.blank?
raise ArgumentError, 'invalid mapping value' if mapping.any? { |_k, v| !v.is_a?(Hash) }
mapping
end
private
attr_reader :table_name, :connection, :columns, :mapping
def log_name
strong_memoize(:log_name) do
"SetAll #{table_name} #{columns.drop(1).map(&:name)}:#{mapping.size}"
end
end end
def params def params
...@@ -58,8 +93,8 @@ module Gitlab ...@@ -58,8 +93,8 @@ module Gitlab
# A workaround for https://github.com/rails/rails/issues/24893 # A workaround for https://github.com/rails/rails/issues/24893
# We need to detect if prepared statements have been disabled. # We need to detect if prepared statements have been disabled.
def no_prepared_statement? def without_prepared_statement?
strong_memoize(:no_prepared_statement) do strong_memoize(:without_prepared_statement) do
connection.send(:without_prepared_statement?, [1]) # rubocop: disable GitlabSecurity/PublicSend connection.send(:without_prepared_statement?, [1]) # rubocop: disable GitlabSecurity/PublicSend
end end
end end
...@@ -83,41 +118,46 @@ module Gitlab ...@@ -83,41 +118,46 @@ module Gitlab
end end
typed = true typed = true
"(#{binds.join(COMMA)})" "(#{list_of(binds)})"
end end
end end
def sql def list_of(list)
column_names = columns.map(&:name) list.join(LIST_SEPARATOR)
cte_columns = column_names.map do |c| end
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
def sql
<<~SQL <<~SQL
WITH cte(#{cte_columns.join(COMMA)}) AS (VALUES #{values.join(COMMA)}) WITH cte(#{list_of(cte_columns)}) AS (VALUES #{list_of(values)})
UPDATE #{table_name} SET #{updates.join(COMMA)} FROM cte WHERE cte_id = id UPDATE #{table_name} SET #{list_of(updates)} FROM cte WHERE cte_id = id
SQL SQL
end end
def update! def column_names
log_name = "SetAll #{table_name} #{columns.drop(1).map(&:name)}:#{mapping.size}" strong_memoize(:column_names) { columns.map(&:name) }
if no_prepared_statement? end
# A workaround for https://github.com/rails/rails/issues/24893
# When prepared statements are prevented (such as when using the def cte_columns
# query counter or in omnibus by default), we cannot call strong_memoize(:cte_columns) do
# `exec_update`, since that will discard the bindings. column_names.map do |c|
connection.send(:exec_no_cache, sql, log_name, params) # rubocop: disable GitlabSecurity/PublicSend connection.quote_column_name("cte_#{c}")
else end
connection.exec_update(sql, log_name, params) end
end
def updates
column_names.zip(cte_columns).drop(1).map do |dest, src|
"#{connection.quote_column_name(dest)} = cte.#{src}"
end end
end end
end end
def self.set_all(columns, mapping, &to_class) def self.set_all(columns, mapping, &to_class)
mapping.group_by { |k, v| block_given? ? to_class.call(k) : k.class }.each do |model, entries| raise ArgumentError if mapping.blank?
entries_by_class = mapping.group_by { |k, v| block_given? ? to_class.call(k) : k.class }
entries_by_class.each do |model, entries|
Setter.new(model, columns, entries).update! Setter.new(model, columns, entries).update!
end end
end end
......
...@@ -3,6 +3,52 @@ ...@@ -3,6 +3,52 @@
require 'spec_helper' require 'spec_helper'
RSpec.describe Gitlab::Database::SetAll do RSpec.describe Gitlab::Database::SetAll do
describe 'error states' do
let(:columns) { %i[title] }
let_it_be(:mapping) do
create_default(:user)
create_default(:project)
i_a, i_b = create_list(:issue, 2)
{
i_a => { title: 'Issue a' },
i_b => { title: 'Issue b' }
}
end
it 'does not raise errors on valid inputs' do
expect { described_class.set_all(columns, mapping) }.not_to raise_error
end
it 'expects a non-empty list of column names' do
expect { described_class.set_all([], mapping) }.to raise_error(ArgumentError)
end
it 'expects all columns to be symbols' do
expect { described_class.set_all([1], mapping) }.to raise_error(ArgumentError)
end
it 'expects all columns to be valid columns on the tables' do
expect { described_class.set_all([:foo], mapping) }.to raise_error(ArgumentError)
end
it 'refuses to set ID' do
expect { described_class.set_all([:id], mapping) }.to raise_error(ArgumentError)
end
it 'expects a non-empty mapping' do
expect { described_class.set_all(columns, []) }.to raise_error(ArgumentError)
end
it 'expects all map values to be Hash instances' do
bad_map = mapping.merge(build(:issue) => 2)
expect { described_class.set_all(columns, bad_map) }.to raise_error(ArgumentError)
end
end
it 'is possible to update all objects in a single query' do it 'is possible to update all objects in a single query' do
users = create_list(:user, 3) users = create_list(:user, 3)
mapping = users.zip(%w(foo bar baz)).to_h do |u, name| mapping = users.zip(%w(foo bar baz)).to_h do |u, name|
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment