lib/rom/sql/extensions/postgres/commands.rb in rom-sql-2.1.0 vs lib/rom/sql/extensions/postgres/commands.rb in rom-sql-2.2.0

@@ -3,44 +3,61 @@
 
 module ROM
   module SQL
     module Postgres
       module Commands
+        # @api private
+        module Core
+          private
+
+          # Common dataset used by create/update commands
+          #
+          # @return [Sequel::Dataset]
+          #
+          # @api private
+          def returning_dataset
+            relation.dataset.returning(*relation.columns)
+          end
+        end
+
         module Create
+          include Core
+
           # Executes insert statement and returns inserted tuples
           #
           # @api private
           def insert(tuples)
             dataset = tuples.flat_map do |tuple|
-              relation.dataset.returning.insert(tuple)
+              returning_dataset.insert(tuple)
             end
 
             wrap_dataset(dataset)
           end
 
           # Executes multi_insert statement and returns inserted tuples
           #
           # @api private
           def multi_insert(tuples)
-            relation.dataset.returning.multi_insert(tuples)
+            returning_dataset.multi_insert(tuples)
           end
 
           # Executes upsert statement (INSERT with ON CONFLICT clause)
           # and returns inserted/updated tuples
           #
           # @api private
           def upsert(tuple, opts = EMPTY_HASH)
-            relation.dataset.returning.insert_conflict(opts).insert(tuple)
+            returning_dataset.insert_conflict(opts).insert(tuple)
           end
         end
 
         module Update
+          include Core
+
           # Executes update statement and returns updated tuples
           #
           # @api private
           def update(tuple)
-            dataset = relation.dataset.returning.update(tuple)
-            wrap_dataset(dataset)
+            wrap_dataset(returning_dataset.update(tuple))
           end
         end
 
         # Upsert command
         #