athena-utils 0.1.0 → 0.1.1

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 3f4fe05a4b845153c14b1b655a590a0983b8264853c44289ffb1498247b526d8
-  data.tar.gz: 3685bc1f2542f2e02d6a31a7a7e2acf12ea5f78a13f6b526418f59465939f45e
+  metadata.gz: f60738b7ecb0c8155fbdf4f85e0bc0371905840808af83eace67328b9a65177b
+  data.tar.gz: 8055dfa958fe0e4a3d8d7364a75e7a6ee500ea32316677472c80c20525c47d82
 SHA512:
-  metadata.gz: 9b42c88bd431a5f3ee5fa5ee8ef889028f7943ebf2df6fec2b0aec982fdb90c0ffdc87338859147c2cff9f10541debcbec5e3f30b5d4963385d4faf00fe9c4e7
-  data.tar.gz: c3a2a4b4bd08c7aeda35552c7caf2abd352eb5b2c206135cc5d4853798948507078166b459bf22d42d33bfd79d9f67c44379a5fb6f61476e729030bbcbcbc0ec
+  metadata.gz: aeccd09b1b52e06d5c4f7da44ae6b0538f2c779e959d22a2aa3b117709da72c463d409fa640af51208f1e5fc8fecb2cb8e891a45868f4ac9131f01202fbe5b24
+  data.tar.gz: 9485169a4b35602f55b7cf60321cc4629a2798b7380f6ec3d8edf4045c03debe20c362f520136fbd858726d65af15c1bf6b4c04cba16473c8dda4d39c6d1d7ce

data/Gemfile

@@ -4,4 +4,3 @@ source 'http://rubygems.org'
 
 gem 'aws-sdk-athena'
 gem 'aws-sdk-s3'
-gem 'csv-utils'

data/README.md

@@ -1 +1,3 @@
 # AWS Athena Utils
+
+A simple query wrapper around Aws::Athena::Client.

data/athena-utils.gemspec

@@ -2,7 +2,7 @@
 
 Gem::Specification.new do |s|
   s.name        = 'athena-utils'
-  s.version     = '0.1.0'
+  s.version     = '0.1.1'
   s.licenses    = ['MIT']
   s.summary     = 'Athena Utils'
   s.description = 'Tools for querying AWS Athena'
@@ -15,5 +15,4 @@ Gem::Specification.new do |s|
 
   s.add_runtime_dependency 'aws-sdk-athena'
   s.add_runtime_dependency 'aws-sdk-s3'
-  s.add_runtime_dependency 'csv-utils'
 end

data/bin/athena

@@ -20,7 +20,7 @@ OptionParser.new do |opts|
     options[:output_location] = v
   end
 
-  opts.on("-q", "--query QUERY", "SQL Query") do |v|
+  opts.on("-e", "--execute QUERY", "Execute SQL Query") do |v|
     options[:query] = v
   end
 

data/lib/athena_utils/athena_client.rb

@@ -2,14 +2,23 @@ require 'aws-sdk-athena'
 
 module AthenaUtils
   class AthenaClient
+    DEFAULT_WAIT_TIME = 3 # seconds
+
     # database is the name of the Athena DB
     # output_location is the full S3 path to store the results of Athena queries
     attr_reader :database,
                 :output_location
 
-    def initialize(database, output_location)
+    # wait_time is time to wait before checking query results again
+    attr_accessor :wait_time
+
+    attr_writer :aws_athena_client,
+                :aws_s3_client
+
+    def initialize(database, output_location, wait_time = DEFAULT_WAIT_TIME)
       @database = database
       @output_location = output_location
+      @wait_time = wait_time
     end
 
     def aws_athena_client
@@ -20,6 +29,14 @@ module AthenaUtils
       Aws::Athena::Client.new
     end
 
+    def aws_s3_client
+      @aws_s3_client ||= create_aws_s3_client
+    end
+
+    def create_aws_s3_client
+      Aws::S3::Client.new
+    end
+
     def query(query)
       query_execution_id = query_async(query)
       wait([query_execution_id])[query_execution_id]
@@ -52,16 +69,17 @@ module AthenaUtils
 
           case query_status[:query_execution][:status][:state]
           when 'SUCCEEDED'
-            results[query_execution_id] = AthenaQueryResults.new(query_status)
+            results[query_execution_id] = AthenaQueryResults.new(query_status, aws_s3_client)
           when 'RUNNING',
                'QUEUED'
           # no-op
+            next
           else
             raise(AthenaQueryError.new("Query failed #{query_status}"))
           end
         end
 
-        sleep(3) if results.size != query_execution_ids.size
+        sleep(wait_time) if results.size != query_execution_ids.size
       end
 
       results

data/lib/athena_utils/athena_query_results.rb

@@ -1,14 +1,16 @@
 require 'aws-sdk-s3'
-require 'csv-utils'
+require 'csv'
 
 module AthenaUtils
   class AthenaQueryResults
     include Enumerable
 
-    attr_reader :query_status
+    attr_reader :query_status,
+                :aws_s3_client
 
-    def initialize(query_status)
+    def initialize(query_status, aws_s3_client)
       @query_status = query_status
+      @aws_s3_client = aws_s3_client
     end
 
     def s3_url
@@ -38,20 +40,21 @@ module AthenaUtils
       )
     end
 
-    def aws_s3_client
-      @aws_s3_client ||= create_aws_s3_client
+    def csv
+      @csv ||= CSV.new(s3_object.body)
     end
 
-    def create_aws_s3_client
-      Aws::S3::Client.new
+    def headers
+      csv.rewind
+      csv.shift
     end
 
-    def csv_iterator
-      @csv_iterator ||= CSVUtils::CSVIterator.new(CSV.new(s3_object.body))
-    end
-
-    def each(&block)
-      csv_iterator.each(&block)
+    def each
+      csv.rewind
+      headers = csv.shift
+      while (row = csv.shift)
+        yield Hash[headers.zip(row)]
+      end
     end
   end
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: athena-utils
 version: !ruby/object:Gem::Version
-  version: 0.1.0
+  version: 0.1.1
 platform: ruby
 authors:
 - Doug Youch
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2022-02-21 00:00:00.000000000 Z
+date: 2022-02-22 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: aws-sdk-athena
@@ -38,20 +38,6 @@ dependencies:
     - - ">="
       - !ruby/object:Gem::Version
         version: '0'
-- !ruby/object:Gem::Dependency
-  name: csv-utils
-  requirement: !ruby/object:Gem::Requirement
-    requirements:
-    - - ">="
-      - !ruby/object:Gem::Version
-        version: '0'
-  type: :runtime
-  prerelease: false
-  version_requirements: !ruby/object:Gem::Requirement
-    requirements:
-    - - ">="
-      - !ruby/object:Gem::Version
-        version: '0'
 description: Tools for querying AWS Athena
 email: dougyouch@gmail.com
 executables: