Add support for Bulk Upserting records

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

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 <[email protected]>

Author: robcole

db/migrations/20220113043033_add_unique_constraint_to_users.cr

@@ -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

spec/avram/bulk_upsert_spec.cr

@@ -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

spec/avram/operations/save_operation_spec.cr

@@ -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
@@ -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

spec/avram/view_spec.cr

@@ -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"

src/avram/bulk_upsert.cr

@@ -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

src/avram/save_operation.cr

@@ -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
@@ -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

src/avram/upsert.cr

@@ -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.