Skip to content

Commit

Permalink
Add support for Bulk Upserting records
Browse files Browse the repository at this point in the history
🍀🍀🍀🍀🍀🍀🍀🍀🍀🍀🍀🍀🍀🍀🍀🍀🍀🍀🍀🍀🍀🍀

It's alive! It's alive! Adds the ability to
insert large numbers of records all at once.

(don't worry, will rebase once I'm sane again)

🍀🍀🍀🍀🍀🍀🍀🍀🍀🍀🍀🍀🍀🍀🍀🍀🍀🍀🍀🍀🍀🍀

Co-authored-by: Alex Piechowski <alex@piechowski.io>
  • Loading branch information
robcole and grepsedawk committed Jan 13, 2022
1 parent 36e2f87 commit 278b28b
Show file tree
Hide file tree
Showing 7 changed files with 174 additions and 4 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
class AddUniqueConstraintToUsers::V20220113043033 < Avram::Migrator::Migration::V1
def migrate
create_index :users, [:name, :nickname], unique: true
end

def rollback
drop_index :users, [:name, :nickname]
end
end
19 changes: 19 additions & 0 deletions spec/avram/bulk_upsert_spec.cr
Original file line number Diff line number Diff line change
@@ -0,0 +1,19 @@
require "../spec_helper"

describe Avram::BulkUpsert do
describe "bulk upserts" do
# Ideally, compiler can catch this / this should be impossible..
context "when collections mismatch" do
end

context "when mixed new records and updated records" do
# Insert spec example.
it "inserts with a hash of String" do
# params = {:first_name => "Paul", :last_name => "Smith"}
# insert = Avram::Insert.new(table: :users, params: params)
# insert.statement.should eq "insert into users(first_name, last_name) values($1, $2) returning *"
# insert.args.should eq ["Paul", "Smith"]
end
end
end
end
31 changes: 31 additions & 0 deletions spec/avram/operations/save_operation_spec.cr
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,7 @@ end

private class UpsertUserOperation < User::SaveOperation
upsert_lookup_columns :name, :nickname
upsert_unique_on :name, :nickname
end

private class OverrideDefaults < ModelWithDefaultValues::SaveOperation
Expand Down Expand Up @@ -307,6 +308,36 @@ describe "Avram::SaveOperation" do
end
end

describe ".bulk_upsert" do
context "when the records are persisted" do
it "should upsert records" do
user = UserFactory.create do |u|
u.name("Test 1")
u.nickname("Test Nickname 1")
u.age(64)
u.year_born(1942)
u.joined_at(Time.utc)
end

record_args = (1..2).to_a.map do |i|
{
name: "Test #{i}",
nickname: "Test Nickname #{i}",
year_born: nil,
age: 42,
joined_at: Time.utc,
}
end

records = UpsertUserOperation.bulk_upsert(record_args)
records.map(&.year_born).flatten.uniq.should eq [nil]
end
end

context "when the records are persisted" do
end
end

describe "#errors" do
it "includes errors for all operation attributes" do
operation = SaveUser.new
Expand Down
6 changes: 3 additions & 3 deletions spec/avram/view_spec.cr
Original file line number Diff line number Diff line change
Expand Up @@ -12,9 +12,9 @@ describe "views" do
end

it "works without a primary key" do
UserFactory.new.nickname("Johnny").create
UserFactory.new.nickname("Johnny").create
UserFactory.new.nickname("Johnny").create
UserFactory.new.name("P1").nickname("Johnny").create
UserFactory.new.name("P2").nickname("Johnny").create
UserFactory.new.name("P3").nickname("Johnny").create
nickname_info = NicknameInfo::BaseQuery.first

nickname_info.nickname.should eq "Johnny"
Expand Down
82 changes: 82 additions & 0 deletions src/avram/bulk_upsert.cr
Original file line number Diff line number Diff line change
@@ -0,0 +1,82 @@
class Avram::BulkUpsert(T)
def initialize(@records : Array(T), @column_names : Array(Symbol))
@records = set_timestamps(records)
end

def statement
<<-SQL
INSERT INTO #{table}(#{fields})
(select * from unnest(#{value_placeholders}))
ON CONFLICT (#{conflicts}) DO UPDATE SET #{updates}
RETURNING #{returning}
SQL
end

def args
@records.map do |record|
record.changed_attributes.map(&.value)
end.transpose
end

private def conflicts
@column_names.join(", ")
end

private def set_timestamps(collection)
collection.map do |record|
record.created_at.value ||= Time.utc if record.responds_to?(:created_at)
record.updated_at.value = Time.utc if record.responds_to?(:updated_at)
record
end
end

private def table
@records.first.table_name
end

private def updates
update_keys = changed_attributes.flat_map(&.name)
(update_keys - [:created_at]).map do |column|
"#{column}=EXCLUDED.#{column}"
end.join(", ")
end

private def returning
T.column_names.join(", ")
end

private def changed_attributes
@records.first.changed_attributes
end

private def fields
changed_attributes.map do |key|
<<-TEXT
"#{key.name.to_s}"
TEXT
end.join(", ")
end

private def column_types
T.database_table_info.not_nil!.columns.map do |column_info|
[
column_info.column_name,
column_info.data_type,
]
end.to_h
end

private def cast(column)
"#{column_types[column.name.to_s]}[]"
end

private def cast(column : Avram::Attribute(Time))
"timestamptz[]"
end

private def value_placeholders
changed_attributes.map_with_index(1) do |k, index|
"$#{index}::#{cast(k)}"
end.join(", ")
end
end
12 changes: 12 additions & 0 deletions src/avram/save_operation.cr
Original file line number Diff line number Diff line change
Expand Up @@ -367,6 +367,10 @@ abstract class Avram::SaveOperation(T)
{{ T.constant(:PRIMARY_KEY_NAME).id }}.value.nil?
end

def changed_attributes
column_attributes.select(&.changed?)
end

private def insert_or_update
if persisted?
update record_id
Expand All @@ -379,6 +383,14 @@ abstract class Avram::SaveOperation(T)
@record.try &.id
end

def self.column_names
T.column_names
end

def self.database_table_info
T.database_table_info
end

def before_save; end

def after_save(_record : T); end
Expand Down
19 changes: 18 additions & 1 deletion src/avram/upsert.cr
Original file line number Diff line number Diff line change
Expand Up @@ -90,9 +90,26 @@ module Avram::Upsert
end
end

macro upsert_unique_on(*attribute_names)
def self.bulk_upsert(upserts)
operations = upserts.map do |upsert_args|
new(**upsert_args)
end

upsert = Avram::BulkUpsert(self).new(
operations,
{{ attribute_names }}.to_a
)

new.database.query upsert.statement, args: upsert.args do |rs|
T.from_rs(rs)
end
end
end

# :nodoc:
macro included
{% for method in ["upsert", "upsert!"] %}
{% for method in ["upsert", "upsert!", "bulk_upsert"] %}
# Performs a create or update depending on if there is a conflicting row in the database.
#
# See `Avram::Upsert.upsert_lookup_columns` for full documentation and examples.
Expand Down

0 comments on commit 278b28b

Please sign in to comment.