Class: AthenaUtils::AthenaQueryResults

Inherits:
Object
  • Object
show all
Includes:
Enumerable
Defined in:
lib/athena_utils/athena_query_results.rb

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(query_status, aws_s3_client) ⇒ AthenaQueryResults

Returns a new instance of AthenaQueryResults.



11
12
13
14
 
# File 'lib/athena_utils/athena_query_results.rb', line 11

def initialize(query_status, aws_s3_client)
  @query_status = query_status
  @aws_s3_client = aws_s3_client
end

Instance Attribute Details

#aws_s3_client ⇒ Object (readonly)

Returns the value of attribute aws_s3_client.



8
9
10
 
# File 'lib/athena_utils/athena_query_results.rb', line 8

def aws_s3_client
  @aws_s3_client
end

#query_status ⇒ Object (readonly)

Returns the value of attribute query_status.



8
9
10
 
# File 'lib/athena_utils/athena_query_results.rb', line 8

def query_status
  @query_status
end

Instance Method Details

#csv ⇒ Object 



53
54
55
 
# File 'lib/athena_utils/athena_query_results.rb', line 53

def csv
  @csv ||= CSV.new(s3_object.body)
end

#each ⇒ Object 



62
63
64
65
66
67
68
 
# File 'lib/athena_utils/athena_query_results.rb', line 62

def each
  csv.rewind
  headers = csv.shift
  while (row = csv.shift)
    yield Hash[headers.zip(row)]
  end
end

#headers ⇒ Object 



57
58
59
60
 
# File 'lib/athena_utils/athena_query_results.rb', line 57

def headers
  csv.rewind
  csv.shift
end

#s3_object ⇒ Object 



20
21
22
23
24
25
26
27
28
29
 
# File 'lib/athena_utils/athena_query_results.rb', line 20

def s3_object
  uri = URI(s3_url)

  aws_s3_client.get_object(
    {
      bucket: uri.host,
      key: uri.path[1..-1]
    }
  )
end

#s3_url ⇒ Object 



16
17
18
 
# File 'lib/athena_utils/athena_query_results.rb', line 16

def s3_url
  query_status.query_execution.result_configuration.output_location
end

#save(file) ⇒ Object 



31
32
33
34
35
36
37
38
39
40
41
 
# File 'lib/athena_utils/athena_query_results.rb', line 31

def save(file)
  uri = URI(s3_url)

  aws_s3_client.get_object(
    {
      bucket: uri.host,
      key: uri.path[1..-1],
      response_target: file
    }
  )
end

#save_in_jsonl(file) ⇒ Object 



43
44
45
46
47
48
49
50
51
 
# File 'lib/athena_utils/athena_query_results.rb', line 43

def save_in_jsonl(file)
  require 'json'

  File.open(file, 'wb') do |out|
    each do |row|
      out.print(row.to_json + "\n")
    end
  end
end