diff --git a/app/models/manager_refresh/save_collection/saver/concurrent_safe_batch.rb b/app/models/manager_refresh/save_collection/saver/concurrent_safe_batch.rb index edd0f53ef13..f2c09358bc7 100644 --- a/app/models/manager_refresh/save_collection/saver/concurrent_safe_batch.rb +++ b/app/models/manager_refresh/save_collection/saver/concurrent_safe_batch.rb @@ -118,21 +118,35 @@ def create_records!(inventory_collection, all_attribute_keys, batch, attributes_ return if hashes.blank? - ActiveRecord::Base.connection.execute( + result = ActiveRecord::Base.connection.execute( build_insert_query(inventory_collection, all_attribute_keys, hashes) ) if inventory_collection.dependees.present? # We need to get primary keys of the created objects, but only if there are dependees that would use them - map_ids_to_inventory_objects(inventory_collection, indexed_inventory_objects, hashes) + map_ids_to_inventory_objects(inventory_collection, indexed_inventory_objects, all_attribute_keys, hashes, result) end end - def map_ids_to_inventory_objects(inventory_collection, indexed_inventory_objects, hashes) - inventory_collection.model_class.where( - build_multi_selection_query(inventory_collection, hashes) - ).select(inventory_collection.unique_index_columns + [:id]).each do |inserted_record| - inventory_object = indexed_inventory_objects[inventory_collection.unique_index_columns.map { |x| inserted_record.public_send(x) }] - inventory_object.id = inserted_record.id if inventory_object + def map_ids_to_inventory_objects(inventory_collection, indexed_inventory_objects, all_attribute_keys, hashes, result) + # The remote_data_timestamp is adding a WHERE condition to ON CONFLICT UPDATE. As a result, the RETURNING + # clause is not guaranteed to return all ids of the inserted/updated records in the result. In that case + # we test if the number of results matches the expected batch size. Then if the counts do not match, the only + # safe option is to query all the data from the DB, using the unique_indexes. The batch size will also not match + # for every remainders(a last batch in a stream of batches) + if !supports_remote_data_timestamp?(all_attribute_keys) || result.count == batch_size + result.each do |inserted_record| + key = inventory_collection.unique_index_columns.map { |x| inserted_record[x.to_s] } + inventory_object = indexed_inventory_objects[key] + inventory_object.id = inserted_record["id"] if inventory_object + end + else + inventory_collection.model_class.where( + build_multi_selection_query(inventory_collection, hashes) + ).select(inventory_collection.unique_index_columns + [:id]).each do |inserted_record| + key = inventory_collection.unique_index_columns.map { |x| inserted_record.public_send(x) } + inventory_object = indexed_inventory_objects[key] + inventory_object.id = inserted_record.id if inventory_object + end end end end diff --git a/app/models/manager_refresh/save_collection/saver/sql_helper.rb b/app/models/manager_refresh/save_collection/saver/sql_helper.rb index 462b896514d..6cac5e10b1e 100644 --- a/app/models/manager_refresh/save_collection/saver/sql_helper.rb +++ b/app/models/manager_refresh/save_collection/saver/sql_helper.rb @@ -38,6 +38,13 @@ def build_insert_query(inventory_collection, all_attribute_keys, hashes) WHERE EXCLUDED.remote_data_timestamp IS NULL OR (EXCLUDED.remote_data_timestamp > #{table_name}.remote_data_timestamp) } end + + if inventory_collection.dependees.present? + insert_query += %{ + RETURNING id,#{inventory_collection.unique_index_columns.join(",")} + } + end + insert_query end