Add Relation#on_conflict with a DSL for upserts

changelog.yml

@@ -1,6 +1,8 @@
 - version: 3.8.0
   summary:
   date: unreleased
+  added:
+  - "`Relation#on_conflict`, `#do_update` and `#do_nothing` for building INSERT ... ON CONFLICT statements, with a DSL for the DO UPDATE part. Create commands and changesets built on such a relation become upserts (@flash-gordon)"
   fixed:
   - "Use dataset to fetch record after insert so that read types are properly applied in Create command (@katafrakt)"
   changed:

lib/rom/sql/commands/create.rb

@@ -11,6 +11,9 @@ module Commands
       class Create < ROM::Commands::Create
         adapter :sql
 
+        # Lets a relation carrying ON CONFLICT be used via `relation.command(:create)`
+        restrictable true
+
         include ErrorWrapper
 
         use :associates

lib/rom/sql/errors.rb

@@ -19,6 +19,7 @@ module SQL
     MissingPrimaryKeyError    = Class.new(StandardError)
     MigrationError            = Class.new(StandardError)
     UnsupportedConversion     = Class.new(MigrationError)
+    UnsupportedFeatureError   = Class.new(StandardError)
 
     ERROR_MAP = {
       Sequel::DatabaseError => DatabaseError,

lib/rom/sql/relation/writing.rb

@@ -1,9 +1,106 @@
 # frozen_string_literal: true
 
+require 'rom/sql/upsert_dsl'
+
 module ROM
   module SQL
     class Relation < ROM::Relation
       module Writing
+        # Handle conflicts with unique constraints on insert
+        #
+        # Returns a relation carrying an ON CONFLICT clause, so any insert
+        # through it becomes an upsert: `insert`, `multi_insert`, `import`,
+        # `command(:create)` and changesets. Conflicts are ignored unless
+        # `do_update` is chained.
+        #
+        # @example ignore conflicting rows
+        #   users.on_conflict(:email).insert(name: 'Jane', email: 'jane@doe.org')
+        #
+        # @example infer a partial unique index
+        #   users.on_conflict(:email) { active.is(true) }.insert(...)
+        #
+        # @example use a named constraint
+        #   users.on_conflict(constraint: :users_email_key).insert(...)
+        #
+        # @param [Array<Symbol, SQL::Attribute>] target Columns of the unique index
+        # @param [Symbol] constraint Name of the constraint, takes precedence over target
+        #
+        # @yield Index predicate built with the restriction DSL
+        #
+        # @return [Relation]
+        #
+        # @api public
+        def on_conflict(*target, constraint: nil, &block)
+          opts = conflict_options.except(:target, :constraint, :conflict_where)
+          opts[:constraint] = constraint if constraint
+          opts[:target] = target.map { |t| t.is_a?(Attribute) ? t.name : t } unless target.empty?
+          opts[:conflict_where] = schema.canonical.restriction(&block) if block
+
+          new(insert_conflict(opts))
+        end
+
+        # Update the existing row on conflict
+        #
+        # Without arguments every column except the primary key and the
+        # conflict target is set from the excluded row, i.e. the one proposed
+        # for insertion. Listed columns restrict that, pairs set arbitrary
+        # values, and the block gives access to the upsert DSL.
+        #
+        # @example set columns from the excluded row
+        #   users.on_conflict(:email).do_update(:name, :updated_at).insert(...)
+        #
+        # @example set values and expressions
+        #   users.on_conflict(:email).do_update(removed_at: nil).insert(...)
+        #
+        # @example use the DSL
+        #   users.on_conflict(:email).do_update {
+        #     set(last_seen_at: greatest(last_seen_at, excluded[:last_seen_at]))
+        #       .where(updated_at < excluded[:updated_at])
+        #   }.multi_insert(tuples)
+        #
+        # @param [Array<Symbol>] columns Columns to set from the excluded row
+        # @param [Hash] pairs Column-value pairs
+        #
+        # @yield Block evaluated with the upsert DSL, returns `set(...)` or a hash
+        #
+        # @return [Relation]
+        #
+        # @see UpsertDSL
+        #
+        # @api public
+        def do_update(*columns, **pairs, &)
+          assignments, condition = conflict_update(columns, pairs, &)
+
+          opts = conflict_options.except(:update_where).merge(update: assignments)
+          opts[:update_where] = condition if condition
+
+          new(insert_conflict(opts))
+        end
+
+        # Ignore the conflicting row
+        #
+        # This is the default of `on_conflict`, the method exists to revert
+        # a previous `do_update`
+        #
+        # @return [Relation]
+        #
+        # @api public
+        def do_nothing
+          new(insert_conflict(conflict_options.except(:update, :update_where)))
+        end
+
+        # Row proposed for insertion in ON CONFLICT DO UPDATE
+        #
+        # @example
+        #   users.on_conflict(:email).do_update(name: users.excluded[:name])
+        #
+        # @return [SQL::Schema] Schema with attributes qualified with `excluded`
+        #
+        # @api public
+        def excluded
+          schema.qualified(:excluded)
+        end
+
         # Add upsert option (only PostgreSQL >= 9.5)
         # Uses internal Sequel implementation
         # Default - ON CONFLICT DO NOTHING
@@ -128,6 +225,53 @@ def import(other, options = EMPTY_HASH)
             )
           end
         end
+
+        private
+
+        # @api private
+        def insert_conflict(opts)
+          unless dataset.respond_to?(:insert_conflict)
+            raise UnsupportedFeatureError,
+                  "ON CONFLICT is not supported by #{dataset.db.database_type}"
+          end
+
+          dataset.insert_conflict(opts)
+        end
+
+        # Options of the ON CONFLICT clause set on the dataset
+        #
+        # @api private
+        def conflict_options
+          opts = dataset.opts[:insert_conflict] || dataset.opts[:insert_on_conflict]
+          opts.is_a?(Hash) ? opts : EMPTY_HASH
+        end
+
+        # Assignments and condition of the DO UPDATE clause
+        #
+        # @return [Array(Hash, Object)]
+        #
+        # @api private
+        def conflict_update(columns, pairs, &block)
+          assignments = columns.to_h { |column| [column, excluded[column]] }.merge(pairs)
+          condition = nil
+
+          if block
+            dsl_assignments, condition = UpsertDSL.new(schema).call(&block)
+            assignments = assignments.merge(dsl_assignments)
+          end
+
+          assignments = default_conflict_assignments if assignments.empty?
+
+          [assignments.to_h { |k, v| [k.is_a?(Attribute) ? k.name : k, v] }, condition]
+        end
+
+        # Every column except the primary key and the conflict target
+        #
+        # @api private
+        def default_conflict_assignments
+          keys = schema.primary_key.map(&:name) + Array(conflict_options[:target])
+          (schema.map(&:name) - keys).to_h { |name| [name, excluded[name]] }
+        end
       end
     end
   end

lib/rom/sql/upsert_dsl.rb

@@ -0,0 +1,126 @@
+# frozen_string_literal: true
+
+require 'rom/sql/restriction_dsl'
+
+module ROM
+  module SQL
+    # DSL for the DO UPDATE part of INSERT ... ON CONFLICT statements
+    #
+    # Attributes refer to the existing row, `excluded` gives access to
+    # the row proposed for insertion, like PostgreSQL's EXCLUDED pseudo-table.
+    # Functions are resolved through the virtual row, as in the restriction DSL.
+    #
+    # @example
+    #   users.on_conflict(:email).do_update {
+    #     set(name: excluded[:name], updated_at: excluded[:updated_at])
+    #       .where(updated_at < excluded[:updated_at])
+    #   }
+    #
+    # @api public
+    class UpsertDSL < RestrictionDSL
+      # Immutable description of the update: assignments and an optional condition
+      #
+      # Every method returns a new instance, so the parts are chained
+      #
+      # @api public
+      class Update
+        # @!attribute [r] schema
+        #   @return [SQL::Schema]
+        attr_reader :schema
+
+        # @!attribute [r] assignments
+        #   @return [Hash] Column-value pairs for the SET clause
+        attr_reader :assignments
+
+        # @!attribute [r] condition
+        #   @return [Object, nil] Condition for the WHERE clause
+        attr_reader :condition
+
+        # @api private
+        def initialize(schema, assignments = EMPTY_HASH, condition = nil)
+          @schema = schema
+          @assignments = assignments
+          @condition = condition
+          freeze
+        end
+
+        # Add assignments, later ones take precedence
+        #
+        # @param [Hash] pairs Column-value pairs
+        #
+        # @return [Update]
+        #
+        # @api public
+        def set(**pairs)
+          self.class.new(schema, assignments.merge(pairs), condition)
+        end
+
+        # Restrict the update to rows matching the condition
+        #
+        # @param [Hash, Object] condition
+        #
+        # @return [Update]
+        #
+        # @api public
+        def where(condition)
+          self.class.new(schema, assignments, qualify(condition))
+        end
+
+        private
+
+        # Bare column names in DO UPDATE ... WHERE are ambiguous
+        # between the table and EXCLUDED
+        #
+        # @api private
+        def qualify(condition)
+          return condition unless condition.is_a?(Hash)
+
+          condition.to_h { |key, value|
+            [key.is_a?(Symbol) && schema.key?(key) ? schema[key].qualified : key, value]
+          }
+        end
+      end
+
+      # Row proposed for insertion
+      #
+      # @return [SQL::Schema] Schema with attributes qualified with `excluded`
+      #
+      # @api public
+      def excluded
+        @excluded ||= schema.qualified(:excluded)
+      end
+
+      # Start the update with assignments
+      #
+      # @see Update#set
+      #
+      # @api public
+      def set(**pairs)
+        Update.new(schema).set(**pairs)
+      end
+
+      # Start the update with a condition
+      #
+      # @see Update#where
+      #
+      # @api public
+      def where(condition)
+        Update.new(schema).where(condition)
+      end
+
+      # @return [Array(Hash, Object)] Assignments and condition
+      #
+      # @api private
+      def call(&)
+        result = super
+
+        case result
+        when Update then [result.assignments, result.condition]
+        when ::Hash then [result, nil]
+        else
+          ::Kernel.raise ::ArgumentError, "expected set(...) or a hash, got #{result.inspect}"
+        end
+      end
+    end
+  end
+end

spec/integration/commands/create_spec.rb

@@ -350,4 +350,36 @@ def self.[](input)
       ])
     end
   end
+
+  describe 'on a relation with conflict handling' do
+    let(:task) { { title: 'task 1' } }
+
+    seed do
+      create_user.call(name: 'Jane')
+      create_task.call(task)
+    end
+
+    it 'returns nothing when the conflict is ignored' do
+      expect(tasks.on_conflict(:title).command(:create).call(task)).to eql([])
+    end
+
+    it 'returns the updated tuple' do
+      command = tasks.on_conflict(:title).do_update(:user_id).command(:create)
+
+      expect(command.call(task.merge(user_id: 1))).to eql([{ id: 1, user_id: 1, title: 'task 1' }])
+    end
+
+    it 'returns inserted and updated tuples' do
+      command = tasks.on_conflict(:title).do_update(:user_id).command(:create)
+
+      expect(command.call([task.merge(user_id: 1), title: 'task 2', user_id: 1])).to match([
+        { id: 1, user_id: 1, title: 'task 1' },
+        hash_including(user_id: 1, title: 'task 2')
+      ])
+    end
+
+    it 'leaves the plain command failing on conflict' do
+      expect { tasks.command(:create).call(task) }.to raise_error(ROM::SQL::UniqueConstraintError)
+    end
+  end
 end