class Multiwoven::Integrations::Source::AwsAthena::Client

def query_execution(db, query)
  response = db.start_query_execution(
    query_string: query,
    query_execution_context: { database: @database },
    result_configuration: { output_location: @output_location }
  )
  query_execution_id = response[:query_execution_id]
  loop do
    response = db.get_query_execution(query_execution_id: query_execution_id)
    status = response.query_execution.status.state
    break if %w[SUCCEEDED FAILED CANCELLED].include?(status)
  end
  transform_results(db.get_query_results(query_execution_id: query_execution_id))
end