ath 0.1.5 → 0.1.6

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA1:
-  metadata.gz: f2e657482ca868e38e2823974acf2d99a7291bc0
-  data.tar.gz: fe8561a1ce77cbcee7decdb892ccbde7a0938b95
+  metadata.gz: c91de42f95791d3fe9be38882ac14ecf0c81618f
+  data.tar.gz: 8ecad5ef8c2818b371117ec83889ba709bd71282
 SHA512:
-  metadata.gz: 9b80a39f2864d34af1c7c301e64419419b36215a9f04ec48e6bcf7721ab57dd0c68247df901e64d83af5ba1300ca2ecd71116f282102620b241a65bc01b255d6
-  data.tar.gz: 718f99c07022cf0305cda6434715f768a13ca11e3e6283f23986d43b200901e9b12b9c87a11dfa408c9cdce319eb1bd0851521fdc56a29843e6a9c6135e43b45
+  metadata.gz: fde85f64b72240c7c5d0b0557d711fd623513071775aa0c5d18819545b52fee25081aeecc46b657f6eb9bb110ea54db81bfeb6e8f6b7ca329f3307ef954d820d
+  data.tar.gz: e4d68468a9e5f7ee113cc55dadce6ac431e32fa7c531bd2d7d513f09f6c340200a4e0e9f8554b078a5ac360e4640c0e8414d1aac2c12743b7b205d6328a890bc

data/lib/ath/driver.rb

@@ -1,11 +1,12 @@
 class Ath::Driver
   attr_accessor :database
 
-  def initialize(athena:, s3:, output_location:, database:)
+  def initialize(athena:, s3:, output_location:, database:, options: {})
     @athena = athena
     @s3 = s3
     @output_location = output_location
     @database = database
+    @options = options
   end
 
   def get_query_execution(query_execution_id:)
@@ -18,13 +19,27 @@ class Ath::Driver
   end
 
   def get_query_execution_result(query_execution_id:)
-    bucket, key = get_query_execution_result_output_location(query_execution_id: query_execution_id)
+    query_execution = @athena.get_query_execution(query_execution_id: query_execution_id).query_execution
+    output_location = query_execution.result_configuration.output_location
+    bucket, key = output_location.sub(%r{\As3://}, '').split('/', 2)
     tmp = Tempfile.create('ath')
 
-    if block_given?
-      @s3.get_object(bucket: bucket, key: key) do |chunk|
-        yield(chunk)
-        tmp.write(chunk)
+    if @options[:progress]
+      head = @s3.head_object(bucket: bucket, key: key)
+      download_progressbar = ProgressBar.create(title: 'Download', total: head.content_length, output: $stderr)
+
+      begin
+        @s3.get_object(bucket: bucket, key: key) do |chunk|
+          tmp.write(chunk)
+
+          begin
+            download_progressbar.progress += chunk.length
+          rescue ProgressBar::InvalidProgressError
+            # nothing to do
+          end
+        end
+      ensure
+        download_progressbar.clear
       end
     else
       @s3.get_object(bucket: bucket, key: key) do |chunk|
@@ -36,17 +51,6 @@ class Ath::Driver
     tmp
   end
 
-  def head_query_execution_result(query_execution_id:)
-    bucket, key = get_query_execution_result_output_location(query_execution_id: query_execution_id)
-    @s3.head_object(bucket: bucket, key: key)
-  end
-
-  def get_query_execution_result_output_location(query_execution_id:)
-    query_execution = @athena.get_query_execution(query_execution_id: query_execution_id).query_execution
-    output_location = query_execution.result_configuration.output_location
-    output_location.sub(%r{\As3://}, '').split('/', 2)
-  end
-
   def start_query_execution(query_string:)
     @athena.start_query_execution(
       query_string: query_string,

data/lib/ath/query.rb

@@ -42,27 +42,7 @@ class Ath::Query
     end
 
     if query_execution.status.state == 'SUCCEEDED'
-      head = @shell.driver.head_query_execution_result(query_execution_id: query_execution_id)
-      download_progressbar = nil
-
-      begin
-        if @shell.options[:progress]
-          download_progressbar = ProgressBar.create(
-            title: 'Download',
-            total: head.content_length,
-            output: $stderr)
-        end
-
-        @shell.driver.get_query_execution_result(query_execution_id: query_execution_id) do |chunk|
-          begin
-            download_progressbar.progress += chunk.length if download_progressbar
-          rescue ProgressBar::InvalidProgressError
-            # nothing to do
-          end
-        end
-      ensure
-        download_progressbar.clear if download_progressbar
-      end
+      @shell.driver.get_query_execution_result(query_execution_id: query_execution_id)
     else
       "QueryExecution #{query_execution_id}: #{query_execution.status.state_change_reason}"
     end

data/lib/ath/shell.rb

@@ -7,7 +7,7 @@ class Ath::Shell
   attr_accessor :pager
 
   def initialize(athena: Aws::Athena::Client.new, s3: Aws::S3::Client.new, output_location:, database: nil, options: {})
-    @driver = Ath::Driver.new(athena: athena, s3: s3, output_location: output_location, database: database)
+    @driver = Ath::Driver.new(athena: athena, s3: s3, output_location: output_location, database: database, options: options)
     @options = options
     @scanner = Ath::Scanner.new(shell: self)
   end

data/lib/ath/version.rb

@@ -1,3 +1,3 @@
 module Ath
-  VERSION = '0.1.5'
+  VERSION = '0.1.6'
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: ath
 version: !ruby/object:Gem::Version
-  version: 0.1.5
+  version: 0.1.6
 platform: ruby
 authors:
 - winebarrel