<commit_before>module ManagerRefresh::SaveCollection
  module Saver
    class ConcurrentSafeBatch < ManagerRefresh::SaveCollection::Saver::Base
      private

      def record_key(record, key)
        send(record_key_method, record, key)
      end

      def ar_record_key(record, key)
        record.public_send(key)
      end

      def pure_sql_record_key(record, key)
        record[select_keys_indexes[key]]
      end

      def batch_iterator(association)
        if pure_sql_records_fetching
          # Building fast iterator doing pure SQL query and therefore avoiding redundant creation of AR objects. The
          # iterator responds to find_in_batches, so it acts like the AR relation. For targeted refresh, the association
          # can already be ApplicationRecordIterator, so we will skip that.
          pure_sql_iterator = lambda do |&block|
            primary_key_offset = nil
            loop do
              relation    = association.select(*select_keys)
                                       .reorder("#{primary_key} ASC")
                                       .limit(batch_size)
              # Using rails way of comparing primary key instead of offset
              relation    = relation.where(arel_primary_key.gt(primary_key_offset)) if primary_key_offset
              records     = get_connection.query(relation.to_sql)
              last_record = records.last
              block.call(records)

              break if records.size < batch_size
              primary_key_offset = record_key(last_record, primary_key)
            end
          end

          ManagerRefresh::ApplicationRecordIterator.new(:iterator => pure_sql_iterator)
        else
          # Normal Rails relation where we can call find_in_batches
          association
        end
      end

      def save!(association)
        attributes_index        = {}
        inventory_objects_index = {}
        all_attribute_keys      = Set.new + inventory_collection.batch_extra_attributes

        inventory_collection.each do |inventory_object|
          attributes = inventory_object.attributes_with_keys(inventory_collection, all_attribute_keys)
          index      = build_stringified_reference(attributes, unique_index_keys)

          # Interesting fact: not building attributes_index and using only inventory_objects_index doesn't do much
          # of a difference, since the most objects inside are shared.
          attributes_index[index]        = attributes
          inventory_objects_index[index] = inventory_object
        end

        all_attribute_keys << :created_at if supports_created_at?
        all_attribute_keys << :updated_at if supports_updated_at?
        all_attribute_keys << :created_on if supports_created_on?
        all_attribute_keys << :updated_on if supports_updated_on?

        _log.debug("Processing #{inventory_collection} of size #{inventory_collection.size}...")

        unless inventory_collection.create_only?
          update_or_destroy_records!(batch_iterator(association), inventory_objects_index, attributes_index, all_attribute_keys)
        end

        unless inventory_collection.create_only?
          unless inventory_collection.custom_reconnect_block.nil?
            inventory_collection.custom_reconnect_block.call(inventory_collection, inventory_objects_index, attributes_index)
          end
        end

        all_attribute_keys << :type if supports_sti?
        # Records that were not found in the DB but sent for saving, we will be creating these in the DB.
        if inventory_collection.create_allowed?
          inventory_objects_index.each_slice(batch_size_for_persisting) do |batch|
            create_records!(all_attribute_keys, batch, attributes_index)
          end
        end
        _log.debug("Processing #{inventory_collection}, "\
                   "created=#{inventory_collection.created_records.count}, "\
                   "updated=#{inventory_collection.updated_records.count}, "\
                   "deleted=#{inventory_collection.deleted_records.count}...Complete")
      rescue => e<commit_msg>Add code for creating skeletal records from skeletal index<commit_after>module ManagerRefresh::SaveCollection
  module Saver
    class ConcurrentSafeBatch < ManagerRefresh::SaveCollection::Saver::Base
      private

      def record_key(record, key)
        send(record_key_method, record, key)
      end

      def ar_record_key(record, key)
        record.public_send(key)
      end

      def pure_sql_record_key(record, key)
        record[select_keys_indexes[key]]
      end

      def batch_iterator(association)
        if pure_sql_records_fetching
          # Building fast iterator doing pure SQL query and therefore avoiding redundant creation of AR objects. The
          # iterator responds to find_in_batches, so it acts like the AR relation. For targeted refresh, the association
          # can already be ApplicationRecordIterator, so we will skip that.
          pure_sql_iterator = lambda do |&block|
            primary_key_offset = nil
            loop do
              relation    = association.select(*select_keys)
                                       .reorder("#{primary_key} ASC")
                                       .limit(batch_size)
              # Using rails way of comparing primary key instead of offset
              relation    = relation.where(arel_primary_key.gt(primary_key_offset)) if primary_key_offset
              records     = get_connection.query(relation.to_sql)
              last_record = records.last
              block.call(records)

              break if records.size < batch_size
              primary_key_offset = record_key(last_record, primary_key)
            end
          end

          ManagerRefresh::ApplicationRecordIterator.new(:iterator => pure_sql_iterator)
        else
          # Normal Rails relation where we can call find_in_batches
          association
        end
      end

      def save!(association)
        attributes_index        = {}
        inventory_objects_index = {}
        all_attribute_keys      = Set.new + inventory_collection.batch_extra_attributes

        inventory_collection.each do |inventory_object|
          attributes = inventory_object.attributes_with_keys(inventory_collection, all_attribute_keys)
          index      = build_stringified_reference(attributes, unique_index_keys)

          # Interesting fact: not building attributes_index and using only inventory_objects_index doesn't do much
          # of a difference, since the most objects inside are shared.
          attributes_index[index]        = attributes
          inventory_objects_index[index] = inventory_object
        end

        all_attribute_keys << :created_at if supports_created_at?
        all_attribute_keys << :updated_at if supports_updated_at?
        all_attribute_keys << :created_on if supports_created_on?
        all_attribute_keys << :updated_on if supports_updated_on?

        _log.debug("Processing #{inventory_collection} of size #{inventory_collection.size}...")

        unless inventory_collection.create_only?
          update_or_destroy_records!(batch_iterator(association), inventory_objects_index, attributes_index, all_attribute_keys)
        end

        unless inventory_collection.create_only?
          unless inventory_collection.custom_reconnect_block.nil?
            inventory_collection.custom_reconnect_block.call(inventory_collection, inventory_objects_index, attributes_index)
          end
        end

        all_attribute_keys << :type if supports_sti?
        # Records that were not found in the DB but sent for saving, we will be creating these in the DB.
        if inventory_collection.create_allowed?
          inventory_objects_index.each_slice(batch_size_for_persisting) do |batch|
            create_records!(all_attribute_keys, batch, attributes_index)
          end

          # Let the GC clean this up
          inventory_objects_index = nil
          attributes_index = nil

          if inventory_collection.parallel_safe?
            # We will create also remaining skeletal records
            skeletal_attributes_index        = {}
            skeletal_inventory_objects_index = {}

            inventory_collection.skeletal_primary_index.each do |_, inventory_object|
              attributes = inventory_object.attributes_with_keys(inventory_collection, all_attribute_keys)
              index      = build_stringified_reference(attributes, unique_index_keys)

              skeletal_attributes_index[index]        = attributes
              skeletal_inventory_objects_index[index] = inventory_object
            end

            skeletal_inventory_objects_index.each_slice(batch_size_for_persisting) do |batch|
              create_records!(all_attribute_keys, batch, skeletal_attributes_index)
            end
          end
        end
        _log.debug("Processing #{inventory_collection}, "\
                   "created=#{inventory_collection.created_records.count}, "\
                   "updated=#{inventory_collection.updated_records.count}, "\
                   "deleted=#{inventory_collection.deleted_records.count}...Complete")
      rescue => e
{'fraction_non_alphanumeric': 0.058573761567773545, 'fraction_numerical': 0.0, 'mean_word_length': 2.546718146718147, 'pattern_counts': {'":': 0, '<': 27, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4046570', 'n_tokens_mistral': 2296, 'n_tokens_neox': 2242, 'n_words': 542}