athena-utils 0.1.0

checksums.yaml

@@ -0,0 +1,7 @@
+---
+SHA256:
+  metadata.gz: 3f4fe05a4b845153c14b1b655a590a0983b8264853c44289ffb1498247b526d8
+  data.tar.gz: 3685bc1f2542f2e02d6a31a7a7e2acf12ea5f78a13f6b526418f59465939f45e
+SHA512:
+  metadata.gz: 9b42c88bd431a5f3ee5fa5ee8ef889028f7943ebf2df6fec2b0aec982fdb90c0ffdc87338859147c2cff9f10541debcbec5e3f30b5d4963385d4faf00fe9c4e7
+  data.tar.gz: c3a2a4b4bd08c7aeda35552c7caf2abd352eb5b2c206135cc5d4853798948507078166b459bf22d42d33bfd79d9f67c44379a5fb6f61476e729030bbcbcbc0ec

data/.gitignore

@@ -0,0 +1,52 @@
+*.gem
+*.rbc
+/.config
+/coverage/
+/InstalledFiles
+/pkg/
+/spec/reports/
+/spec/examples.txt
+/test/tmp/
+/test/version_tmp/
+/tmp/
+
+# Used by dotenv library to load environment variables.
+# .env
+
+## Specific to RubyMotion:
+.dat*
+.repl_history
+build/
+*.bridgesupport
+build-iPhoneOS/
+build-iPhoneSimulator/
+
+## Specific to RubyMotion (use of CocoaPods):
+#
+# We recommend against adding the Pods directory to your .gitignore. However
+# you should judge for yourself, the pros and cons are mentioned at:
+# https://guides.cocoapods.org/using/using-cocoapods.html#should-i-check-the-pods-directory-into-source-control
+#
+# vendor/Pods/
+
+## Documentation cache and generated files:
+/.yardoc/
+/_yardoc/
+/doc/
+/rdoc/
+
+## Environment normalization:
+/.bundle/
+/vendor/bundle
+/lib/bundler/man/
+
+# for a library or gem, you might want to ignore these files since the code is
+# intended to run in multiple environments; otherwise, check them in:
+# Gemfile.lock
+# .ruby-version
+# .ruby-gemset
+
+# unless supporting rvm < 1.11.0 or doing something fancy, ignore this:
+.rvmrc
+
+*~

data/.ruby-gemset

@@ -0,0 +1 @@
+athena-utils

data/.ruby-version

@@ -0,0 +1 @@
+3.0.0

data/Gemfile

@@ -0,0 +1,7 @@
+# frozen_string_literal: true
+
+source 'http://rubygems.org'
+
+gem 'aws-sdk-athena'
+gem 'aws-sdk-s3'
+gem 'csv-utils'

data/Gemfile.lock

@@ -0,0 +1,37 @@
+GEM
+  remote: http://rubygems.org/
+  specs:
+    aws-eventstream (1.2.0)
+    aws-partitions (1.555.0)
+    aws-sdk-athena (1.49.0)
+      aws-sdk-core (~> 3, >= 3.126.0)
+      aws-sigv4 (~> 1.1)
+    aws-sdk-core (3.126.2)
+      aws-eventstream (~> 1, >= 1.0.2)
+      aws-partitions (~> 1, >= 1.525.0)
+      aws-sigv4 (~> 1.1)
+      jmespath (~> 1.0)
+    aws-sdk-kms (1.54.0)
+      aws-sdk-core (~> 3, >= 3.126.0)
+      aws-sigv4 (~> 1.1)
+    aws-sdk-s3 (1.112.0)
+      aws-sdk-core (~> 3, >= 3.126.0)
+      aws-sdk-kms (~> 1)
+      aws-sigv4 (~> 1.4)
+    aws-sigv4 (1.4.0)
+      aws-eventstream (~> 1, >= 1.0.2)
+    csv-utils (0.3.14)
+      inheritance-helper
+    inheritance-helper (0.2.5)
+    jmespath (1.6.0)
+
+PLATFORMS
+  x86_64-darwin-20
+
+DEPENDENCIES
+  aws-sdk-athena
+  aws-sdk-s3
+  csv-utils
+
+BUNDLED WITH
+   2.2.3

data/LICENSE.txt

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2019 Douglas Youch
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

data/README.md

@@ -0,0 +1 @@
+# AWS Athena Utils

data/athena-utils.gemspec

@@ -0,0 +1,19 @@
+# frozen_string_literal: true
+
+Gem::Specification.new do |s|
+  s.name        = 'athena-utils'
+  s.version     = '0.1.0'
+  s.licenses    = ['MIT']
+  s.summary     = 'Athena Utils'
+  s.description = 'Tools for querying AWS Athena'
+  s.authors     = ['Doug Youch']
+  s.email       = 'dougyouch@gmail.com'
+  s.homepage    = 'https://github.com/dougyouch/athena-utils'
+  s.files       = `git ls-files -z`.split("\x0").reject { |f| f.match(%r{^(test|spec|features)/}) }
+  s.bindir      = 'bin'
+  s.executables = s.files.grep(%r{^bin/}) { |f| File.basename(f) }
+
+  s.add_runtime_dependency 'aws-sdk-athena'
+  s.add_runtime_dependency 'aws-sdk-s3'
+  s.add_runtime_dependency 'csv-utils'
+end

data/bin/athena

@@ -0,0 +1,54 @@
+#!/usr/bin/env ruby
+
+require 'athena-utils'
+require 'optparse'
+
+options = {
+  database: nil,
+  output_location: nil,
+  query: nil,
+  save: nil
+}
+OptionParser.new do |opts|
+  opts.banner = "Usage: athena [options]"
+
+  opts.on("-d", "--database DATABASE", "Athena DB") do |v|
+    options[:database] = v
+  end
+
+  opts.on("-o", "--output-location OUTPUT_LOCATION", "S3 output location for athena queries") do |v|
+    options[:output_location] = v
+  end
+
+  opts.on("-q", "--query QUERY", "SQL Query") do |v|
+    options[:query] = v
+  end
+
+  opts.on('-s', '--save FILE', 'Save query results to file') do |v|
+    options[:save] = v
+  end
+end.parse!
+
+raise('must specify a database') unless options[:database]
+raise('must specify output location for athean queries') unless options[:output_location]
+
+@athena = AthenaUtils::AthenaClient.new(options[:database], options[:output_location])
+def athena
+  @athena
+end
+
+if options[:query]
+  @results = athena.query(options[:query])
+  def results
+    @results
+  end
+
+  if options[:save]
+    results.save(options[:save])
+    exit
+  end
+end
+
+require 'irb'
+IRB.start
+

data/lib/athena-utils.rb

@@ -0,0 +1,5 @@
+module AthenaUtils
+  class AthenaQueryError < StandardError; end
+  autoload :AthenaClient, 'athena_utils/athena_client'
+  autoload :AthenaQueryResults, 'athena_utils/athena_query_results'
+end

data/lib/athena_utils/athena_client.rb

@@ -0,0 +1,70 @@
+require 'aws-sdk-athena'
+
+module AthenaUtils
+  class AthenaClient
+    # 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)
+      @database = database
+      @output_location = output_location
+    end
+
+    def aws_athena_client
+      @aws_athena_client ||= create_aws_athena_client
+    end
+
+    def create_aws_athena_client
+      Aws::Athena::Client.new
+    end
+
+    def query(query)
+      query_execution_id = query_async(query)
+      wait([query_execution_id])[query_execution_id]
+    end
+
+    def query_async(query)
+      response = aws_athena_client.start_query_execution(
+        query_string: query,
+        query_execution_context: {
+          database: database
+        },
+        result_configuration: {
+          output_location: output_location
+        }
+      )
+
+      response.query_execution_id
+    end
+
+    def wait(query_execution_ids)
+      results = {}
+
+      while results.size != query_execution_ids.size
+        query_execution_ids.each do |query_execution_id|
+          next if results.key?(query_execution_id)
+
+          query_status = aws_athena_client.get_query_execution(
+            query_execution_id: query_execution_id
+          )
+
+          case query_status[:query_execution][:status][:state]
+          when 'SUCCEEDED'
+            results[query_execution_id] = AthenaQueryResults.new(query_status)
+          when 'RUNNING',
+               'QUEUED'
+          # no-op
+          else
+            raise(AthenaQueryError.new("Query failed #{query_status}"))
+          end
+        end
+
+        sleep(3) if results.size != query_execution_ids.size
+      end
+
+      results
+    end
+  end
+end

data/lib/athena_utils/athena_query_results.rb

@@ -0,0 +1,57 @@
+require 'aws-sdk-s3'
+require 'csv-utils'
+
+module AthenaUtils
+  class AthenaQueryResults
+    include Enumerable
+
+    attr_reader :query_status
+
+    def initialize(query_status)
+      @query_status = query_status
+    end
+
+    def s3_url
+      query_status.query_execution.result_configuration.output_location
+    end
+
+    def s3_object
+      uri = URI(s3_url)
+
+      aws_s3_client.get_object(
+        {
+          bucket: uri.host,
+          key: uri.path[1..-1]
+        }
+      )
+    end
+
+    def save(file)
+      uri = URI(s3_url)
+
+      aws_s3_client.get_object(
+        {
+          bucket: uri.host,
+          key: uri.path[1..-1],
+          response_target: file
+        }
+      )
+    end
+
+    def aws_s3_client
+      @aws_s3_client ||= create_aws_s3_client
+    end
+
+    def create_aws_s3_client
+      Aws::S3::Client.new
+    end
+
+    def csv_iterator
+      @csv_iterator ||= CSVUtils::CSVIterator.new(CSV.new(s3_object.body))
+    end
+
+    def each(&block)
+      csv_iterator.each(&block)
+    end
+  end
+end

data/scripts/console

@@ -0,0 +1,7 @@
+#!/usr/bin/env ruby
+# frozen_string_literal: true
+
+$LOAD_PATH << File.expand_path('../lib', __dir__)
+require 'athena-utils'
+require 'irb'
+IRB.start(__FILE__)

metadata

@@ -0,0 +1,98 @@
+--- !ruby/object:Gem::Specification
+name: athena-utils
+version: !ruby/object:Gem::Version
+  version: 0.1.0
+platform: ruby
+authors:
+- Doug Youch
+autorequire:
+bindir: bin
+cert_chain: []
+date: 2022-02-21 00:00:00.000000000 Z
+dependencies:
+- !ruby/object:Gem::Dependency
+  name: aws-sdk-athena
+  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'
+- !ruby/object:Gem::Dependency
+  name: aws-sdk-s3
+  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'
+- !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:
+- athena
+extensions: []
+extra_rdoc_files: []
+files:
+- ".gitignore"
+- ".ruby-gemset"
+- ".ruby-version"
+- Gemfile
+- Gemfile.lock
+- LICENSE.txt
+- README.md
+- athena-utils.gemspec
+- bin/athena
+- lib/athena-utils.rb
+- lib/athena_utils/athena_client.rb
+- lib/athena_utils/athena_query_results.rb
+- scripts/console
+homepage: https://github.com/dougyouch/athena-utils
+licenses:
+- MIT
+metadata: {}
+post_install_message:
+rdoc_options: []
+require_paths:
+- lib
+required_ruby_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '0'
+required_rubygems_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '0'
+requirements: []
+rubygems_version: 3.2.3
+signing_key:
+specification_version: 4
+summary: Athena Utils
+test_files: []