<commit_before>module BeetleETL
  class AssignIds < Step

    def dependencies
      [TableDiff.step_name(table_name)].to_set
    end

    def run
      [
        Thread.new { assign_new_ids },
        Thread.new { map_existing_ids }
      ].each(&:join)
    end

    def assign_new_ids
      stage_table
        .where(
          import_run_id: run_id,
          transition: 'CREATE'
        )
        .update(
          id: Sequel.function(:NEXTVAL, "public.#{table_name}_id_seq")
        )
    end

    def map_existing_ids
      stage_table
        .from(stage_table_identifier, public_table_identifier)
        .where(
          stage__import_run_id: run_id,
          stage__transition: %w(KEEP UPDATE DELETE UNDELETE),
          stage__external_id: :public__external_id
        )
        .update(id: :public__id)
    end

    private

    def stage_table_identifier
      :"#{stage_schema}__#{table_name}___stage"
    end

    def stage_table
      database[stage_table_identifier]
    end

    def public_table_identifier
      :"#{table_name}___public"
    end

    def public_table
      database[public_table_identifier]
    end

  end
end<commit_msg>use SQL in AssignIds step<commit_after>module BeetleETL
  class AssignIds < Step

    def dependencies
      [TableDiff.step_name(table_name)].to_set
    end

    def run
      [
        Thread.new { assign_new_ids },
        Thread.new { map_existing_ids }
      ].each(&:join)
    end

    def assign_new_ids
      database.execute <<-SQL
        UPDATE #{stage_table_name}
        SET id = nextval('#{table_name}_id_seq')
        WHERE import_run_id = #{run_id}
        AND transition = 'CREATE'
      SQL
    end

    def map_existing_ids
      database.execute <<-SQL
        UPDATE #{stage_table_name} stage
        SET id = public.id
        FROM #{public_table_name} public
        WHERE stage.import_run_id = #{run_id}
        AND stage.transition IN ('KEEP', 'UPDATE', 'DELETE', 'UNDELETE')
        AND stage.external_id = public.external_id
      SQL
    end

    private

    def stage_table_name
      %Q("#{stage_schema}"."#{table_name}")
    end

    def public_table_name
      %Q("#{table_name}")
    end

  end
end
{'fraction_non_alphanumeric': 0.0861440291704649, 'fraction_numerical': 0.0, 'mean_word_length': 2.223201174743025, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6419125', 'n_tokens_mistral': 710, 'n_tokens_neox': 688, 'n_words': 143}