<commit_before>module Embulk
  class InputPresto < InputPlugin
    require 'presto-client'

    Plugin.register_input("presto", self)

    def self.transaction(config, &control)
      task = {
        "host" => config.param("host", :string, default: "localhost"),
        "port" => config.param("port", :integer, default: 8080),
        "schema" => config.param("schema", :string, default: "default"),
        "catalog" => config.param("catalog", :string, default: "native"),
        "query" => config.param("query", :string),
        "user" => config.param("user", :string, default: "embulk"),
        "columns" => config.param("columns", :array, default: [])
      }

      columns = task['columns'].each_with_index.map do |c, i|
        Column.new(i, c["name"], c["type"].to_sym)
      end

      resume(task, columns, 1, &control)
    end

    def self.resume(task, columns, count, &control)
      task_reports = yield(task, columns, count)

      next_config_diff = {}
      return next_config_diff
    end

    def init
      @client = Presto::Client.new(
        server: "#{task['host']}:#{task['port']}",
        catalog: task['catalog'],
        user: task['user'],
        schema: task['schema']
      )
      @query = task["query"]

      Embulk.logger.info "SQL: #{@query}"
    end

    def run
      size = 0
      @client.query(@query) do |q|
        q.each_row {|row|
          page_builder.add(row)
        }
        size = q.rows.size
      end

      page_builder.finish

      task_report = { size: size }
      return task_report
    end
  end
end<commit_msg>fix convert_value for type.<commit_after>module Embulk
  class InputPresto < InputPlugin
    require 'presto-client'

    Plugin.register_input("presto", self)

    def self.transaction(config, &control)
      task = {
        "host" => config.param("host", :string, default: "localhost"),
        "port" => config.param("port", :integer, default: 8080),
        "schema" => config.param("schema", :string, default: "default"),
        "catalog" => config.param("catalog", :string, default: "native"),
        "query" => config.param("query", :string),
        "user" => config.param("user", :string, default: "embulk"),
        "columns" => config.param("columns", :array, default: [])
      }

      columns = task['columns'].each_with_index.map do |c, i|
        Column.new(i, c["name"], c["type"].to_sym)
      end

      resume(task, columns, 1, &control)
    end

    def self.resume(task, columns, count, &control)
      task_reports = yield(task, columns, count)

      next_config_diff = {}
      return next_config_diff
    end

    def init
      @client = Presto::Client.new(
        server: "#{task['host']}:#{task['port']}",
        catalog: task['catalog'],
        user: task['user'],
        schema: task['schema']
      )
      @query = task["query"]

      Embulk.logger.info "SQL: #{@query}"
    end

    def run
      size = 0
      @client.query(@query) do |q|
        q.each_row {|row|
          converted_values = row.map.with_index { |value,i| convert_value(value, schema[i]) }
          page_builder.add(converted_values)
        }
        size = q.rows.size
      end

      page_builder.finish

      task_report = { size: size }
      return task_report
    end

    def convert_value(value, field)
      return nil if value.nil?
      case field["type"]
      when :string
        value
      when :long
        value.to_i
      when :double
        value.to_f
      when :boolean
        if value.is_a?(TrueClass) || value.is_a?(FalseClass)
          value
        else
          downcased_val = value.downcase
          case downcased_val
          when 'true' then true
          when 'false' then false
          when '1' then true
          when '0' then false
          else nil
          end
        end
      when :timestamp
        Time.parse(value)
      when :json
        value
      else
        raise "Unsupported type #{field['type']}"
      end
    end
  end
end
{'fraction_non_alphanumeric': 0.14271381991472284, 'fraction_numerical': 0.003511412089290193, 'mean_word_length': 2.234387672343877, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 17, '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': '4196752', 'n_tokens_mistral': 1249, 'n_tokens_neox': 1221, 'n_words': 293}