From f491dfecca26602f405628fe5ce8a13ec8b33c56 Mon Sep 17 00:00:00 2001
From: Sergei Maximov <s.b.maximov@gmail.com>
Date: Fri, 19 Feb 2021 14:50:31 +0300
Subject: [PATCH] Allow setting index predicate for upserts w/ partial unique
 indices

---
 lib/rom/sql/extensions/postgres/commands.rb |  7 ++++-
 spec/integration/commands/upsert_spec.rb    | 35 +++++++++++++++++++++
 2 files changed, 41 insertions(+), 1 deletion(-)

diff --git a/lib/rom/sql/extensions/postgres/commands.rb b/lib/rom/sql/extensions/postgres/commands.rb
index 3d6cdb6b7..4971a4c64 100644
--- a/lib/rom/sql/extensions/postgres/commands.rb
+++ b/lib/rom/sql/extensions/postgres/commands.rb
@@ -85,7 +85,7 @@ def update(tuple)
         class Upsert < SQL::Commands::Create
           adapter :sql
 
-          defines :constraint, :conflict_target, :update_statement, :update_where
+          defines :constraint, :conflict_target, :conflict_where, :update_statement, :update_where
 
           # @!attribute [r] constraint
           #  @return [Symbol] the name of the constraint expected to be violated
@@ -95,6 +95,10 @@ class Upsert < SQL::Commands::Create
           #  @return [Object] the column or expression to handle a violation on
           option :conflict_target, default: -> { self.class.conflict_target }
 
+          # @!attribute [r] conflict_where
+          #  @return [Object] the index filter, when using a partial index to determine uniqueness
+          option :conflict_where, default: -> { self.class.conflict_where }
+
           # @!attribute [r] update_statement
           #  @return [Object] the update statement which will be executed in case of a violation
           option :update_statement, default: -> { self.class.update_statement }
@@ -123,6 +127,7 @@ def upsert_options
             @upsert_options ||= {
               constraint: constraint,
               target: conflict_target,
+              conflict_where: conflict_where,
               update_where: update_where,
               update: update_statement
             }
diff --git a/spec/integration/commands/upsert_spec.rb b/spec/integration/commands/upsert_spec.rb
index 9b150a036..99cf4fe71 100644
--- a/spec/integration/commands/upsert_spec.rb
+++ b/spec/integration/commands/upsert_spec.rb
@@ -50,6 +50,41 @@
         it 'returns updated data' do
           expect(command.call(excluded)).to eql(id: 1, user_id: 2, title: 'task 1')
         end
+
+        context 'with index predicate' do
+          before do
+            conn.execute <<~SQL
+              ALTER TABLE tasks DROP CONSTRAINT tasks_title_key;
+
+              CREATE UNIQUE INDEX tasks_title_partial_index ON tasks (title)
+                            WHERE user_id = 1;
+            SQL
+          end
+
+          let(:command_config) do
+            -> do
+              conflict_target :title
+              conflict_where user_id: 1
+              update_statement user_id: 2
+            end
+          end
+
+          context 'when predicate matches' do
+            let(:excluded) { task }
+
+            it 'returns updated data', :aggregate_failures do
+              expect(command.call(excluded)).to eql(id: 1, user_id: 2, title: 'task 1')
+            end
+          end
+
+          context 'when predicate does not match' do
+            let(:excluded) { task.update(user_id: 2) }
+
+            it 'creates new task', :aggregate_failures do
+              expect(command.call(excluded)).to eql(id: 2, user_id: 2, title: 'task 1')
+            end
+          end
+        end
       end
 
       context 'with constraint name' do