<commit_before>require 'tmpdir'

require 'salesforce_bulk_query/utils'


module SalesforceBulkQuery
  # Represents a Salesforce api batch. Batch contains a single subquery.
  # Many batches are contained in a Job.
  class Batch
    def initialize(options)
      @sobject = options[:sobject]
      @soql = options[:soql]
      @job_id = options[:job_id]
      @connection = options[:connection]
      @start = options[:start]
      @stop = options[:stop]
      @logger = options[:logger]
      @@directory_path ||= Dir.mktmpdir
      @filename = nil
    end

    attr_reader :soql, :start, :stop, :filename, :fail_message, :batch_id, :csv_record_count

    # Do the api request
    def create
      path = "job/#{@job_id}/batch/"

      response_parsed = @connection.post_xml(path, @soql, {:csv_content_type => true})

      @batch_id = response_parsed['id'][0]
    end

    # check status of the batch
    # if it fails, don't throw an error now, let the job above collect all fails and raise it at once
    def check_status
      succeeded = nil
      failed = nil

      # get the status of the batch
      # https://www.salesforce.com/us/developer/docs/api_asynch/Content/asynch_api_batches_get_info.htm
      status_path = "job/#{@job_id}/batch/#{@batch_id}"
      status_response = @connection.get_xml(status_path)

      # interpret the status
      @status = status_response['state'][0]

      # https://www.salesforce.com/us/developer/docs/api_asynch/Content/asynch_api_batches_interpret_status.htm
      case @status
        when 'Failed'
          failed = true
          @fail_message = status_response['stateMessage']
        when 'InProgress', 'Queued'
          succeeded = false
        when 'Completed'
          succeeded = true
          failed = false
        else
          fail "Something weird happened, #{@batch_id} has status #{@status}."
      end

      if succeeded
        # request to get the result id
        # https://www.salesforce.com/us/developer/docs/api_asynch/Content/asynch_api_batches_get_results.htm
        path = "job/#{@job_id}/batch/#{@batch_id}/result"

        response_parsed = @connection.get_xml(path)

        @result_id = response_parsed["result"] ? response_parsed["result"][0] : nil
      end

      return {
        :failed => failed,
        :fail_message => @fail_message,
        :succeeded => succeeded,
        :result_id => @result_id
      }
    end

    def get_filename
      return "#{@sobject}_#{@batch_id}_#{@start}-#{@stop}.csv"
    end

    def get_result(options={})
      # if it was already downloaded, no one should ask about it
      if @filename
        raise "This batch was already downloaded once: #{@filename}, #{@batch_id}"
      end

      directory_path = options[:directory_path]
      skip_verification = options[:skip_verification]

      # request to get the actual results
      path = "job/#{@job_id}/batch/#{@batch_id}/result/#{@result_id}"

      if !@result_id
        raise "batch #{@batch_id} not finished yet, trying to get result: #{path}"
      end

      directory_path ||= @@directory_path

      # write it to a file
      @filename = File.join(directory_path, get_filename)
      @connection.get_to_file(path, @filename)

      # Verify the number of downloaded records is roughly the same as<commit_msg>slightly different filenames<commit_after>require 'tmpdir'

require 'salesforce_bulk_query/utils'


module SalesforceBulkQuery
  # Represents a Salesforce api batch. Batch contains a single subquery.
  # Many batches are contained in a Job.
  class Batch
    def initialize(options)
      @sobject = options[:sobject]
      @soql = options[:soql]
      @job_id = options[:job_id]
      @connection = options[:connection]
      @start = options[:start]
      @stop = options[:stop]
      @logger = options[:logger]
      @@directory_path ||= Dir.mktmpdir
      @filename = nil
    end

    attr_reader :soql, :start, :stop, :filename, :fail_message, :batch_id, :csv_record_count

    # Do the api request
    def create
      path = "job/#{@job_id}/batch/"

      response_parsed = @connection.post_xml(path, @soql, {:csv_content_type => true})

      @batch_id = response_parsed['id'][0]
    end

    # check status of the batch
    # if it fails, don't throw an error now, let the job above collect all fails and raise it at once
    def check_status
      succeeded = nil
      failed = nil

      # get the status of the batch
      # https://www.salesforce.com/us/developer/docs/api_asynch/Content/asynch_api_batches_get_info.htm
      status_path = "job/#{@job_id}/batch/#{@batch_id}"
      status_response = @connection.get_xml(status_path)

      # interpret the status
      @status = status_response['state'][0]

      # https://www.salesforce.com/us/developer/docs/api_asynch/Content/asynch_api_batches_interpret_status.htm
      case @status
        when 'Failed'
          failed = true
          @fail_message = status_response['stateMessage']
        when 'InProgress', 'Queued'
          succeeded = false
        when 'Completed'
          succeeded = true
          failed = false
        else
          fail "Something weird happened, #{@batch_id} has status #{@status}."
      end

      if succeeded
        # request to get the result id
        # https://www.salesforce.com/us/developer/docs/api_asynch/Content/asynch_api_batches_get_results.htm
        path = "job/#{@job_id}/batch/#{@batch_id}/result"

        response_parsed = @connection.get_xml(path)

        @result_id = response_parsed["result"] ? response_parsed["result"][0] : nil
      end

      return {
        :failed => failed,
        :fail_message => @fail_message,
        :succeeded => succeeded,
        :result_id => @result_id
      }
    end

    def get_filename
      return "#{@sobject}_#{@start}_#{@stop}_#{@batch_id}.csv"
    end

    def get_result(options={})
      # if it was already downloaded, no one should ask about it
      if @filename
        raise "This batch was already downloaded once: #{@filename}, #{@batch_id}"
      end

      directory_path = options[:directory_path]
      skip_verification = options[:skip_verification]

      # request to get the actual results
      path = "job/#{@job_id}/batch/#{@batch_id}/result/#{@result_id}"

      if !@result_id
        raise "batch #{@batch_id} not finished yet, trying to get result: #{path}"
      end

      directory_path ||= @@directory_path

      # write it to a file
      @filename = File.join(directory_path, get_filename)
      @connection.get_to_file(path, @filename)

      # Verify the number of downloaded records is roughly the same as
{'fraction_non_alphanumeric': 0.11007096481956817, 'fraction_numerical': 0.0009059338668277216, 'mean_word_length': 2.7551020408163267, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 13, 'https://': 6, 'lorem ipsum': 0, 'www.': 6, 'xml': 6}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3848961', 'n_tokens_mistral': 2055, 'n_tokens_neox': 1999, 'n_words': 509}