module TinyPresto

# File lib/tiny-presto.rb, line 116
def self.ensure_stop
  presto = TinyPresto.instance
  presto.stop
end

# File lib/tiny-presto.rb, line 56
def self.prepare(table_name, table_data)
  # {'c1': [1, 2, 3], 'c2': ['a', 'b', 'c']}
  columns = table_data.keys
  records = []
  table_data.each do |_, rows|
    rows.each_with_index do |r, idx|
      if records[idx].nil?
        records << [r]
      else
        records[idx] << r
      end
    end
  end
  values_clause = []
  records.each do |record|
    values_clause << print_record(record)
  end
  query = "CREATE TABLE #{table_name} AS SELECT * FROM (values #{values_clause.join(',')}) t(#{columns.join(',')})"
  run_with_retry(query)
end

# File lib/tiny-presto.rb, line 44
def self.print_record(record)
  ret = record.map do |v|
    if v.is_a? Numeric
      v.to_s
    else
      # Non numeric value is interpreted as string
      "'#{v}'"
    end
  end
  "(#{ret.join(',')})"
end

# File lib/tiny-presto.rb, line 81
def self.run(sql)
  presto = TinyPresto.instance
  _, rows = presto.client.run(sql)
  rows
end

# File lib/tiny-presto.rb, line 91
def self.run_with_retry(sql, max_retry = 3)
  max_retry.times do
    return run(sql)
  rescue Presto::Client::PrestoQueryError => e
    # Cluster may be in the initialization phase.
    raise unless e.message.match?(/^No nodes available to run query/)

    sleep(1000)
    next
  end
end

# File lib/tiny-presto.rb, line 107
def self.verify(sql, expected_result)
  rows = run_with_retry(sql)
  rows == expected_result
end