egis 1.1.0

checksums.yaml

@@ -0,0 +1,7 @@
+---
+SHA256:
+  metadata.gz: 1d2255a76aef464d0d8faee0c39be753e1a928133a25de5c4a86f553e51e92ca
+  data.tar.gz: f18ae28053651576ccf941bd42f5a9bf40db32fed3c1c2025e1bc987844199f8
+SHA512:
+  metadata.gz: 9b143dbc650f8c02ba39f1404a2d9b4ebe36c49182a36b947a1d091079b1c94839320c2d05dd9a0db1516d82c9b03694244e4fa467898302e888addb234e1291
+  data.tar.gz: a1577d30cbfd63632dbd55052adf0f863f7a5f571e8a8f8a210649446acc786764d8bec6a15be7cd91742f4f7681bfc1d379a9099934f62b8a30379d12a4bd77

data/egis.gemspec

@@ -0,0 +1,30 @@
+# frozen_string_literal: true
+
+lib = File.expand_path('lib', __dir__)
+$LOAD_PATH.unshift(lib) unless $LOAD_PATH.include?(lib)
+require 'egis/version'
+
+Gem::Specification.new do |spec|
+  spec.name          = 'egis'
+  spec.version       = Egis::VERSION
+  spec.authors       = ['Agnieszka Czereba', 'Marek Mateja']
+  spec.email         = %w[agnieszka.czereba@u2i.com marek.mateja@u2i.com]
+
+  spec.summary       = 'A handy wrapper for AWS Athena Ruby SDK.'
+  spec.homepage      = 'https://github.com/u2i/egis'
+  spec.license       = 'MIT'
+
+  spec.metadata['homepage_uri'] = spec.homepage
+  spec.metadata['source_code_uri'] = spec.homepage
+  spec.metadata['changelog_uri'] = 'https://github.com/u2i/egis/blob/master/CHANGELOG.md'
+
+  # Specify which files should be added to the gem when it is released.
+  # The `git ls-files -z` loads the files in the RubyGem that have been added into git.
+  spec.files = Dir.chdir(File.expand_path(__dir__)) do
+    `git ls-files -z egis.gemspec lib/`.split("\x0")
+  end
+  spec.require_paths = ['lib']
+
+  spec.add_dependency 'aws-sdk-athena', '~> 1.0'
+  spec.add_dependency 'aws-sdk-s3', '~> 1.0'
+end

data/lib/egis/aws_client_provider.rb

@@ -0,0 +1,30 @@
+# frozen_string_literal: true
+
+require 'aws-sdk-s3'
+require 'aws-sdk-athena'
+
+module Egis
+  # @!visibility private
+  class AwsClientProvider
+    def s3_client
+      Aws::S3::Client.new(client_config)
+    end
+
+    def athena_client
+      Aws::Athena::Client.new(client_config)
+    end
+
+    private
+
+    def client_config
+      configuration = Egis.configuration
+
+      config = {}
+      config[:region] = configuration.aws_region if configuration.aws_region
+      config[:access_key_id] = configuration.aws_access_key_id if configuration.aws_access_key_id
+      config[:secret_access_key] = configuration.aws_secret_access_key if configuration.aws_secret_access_key
+      config[:profile] = configuration.aws_profile if configuration.aws_profile
+      config
+    end
+  end
+end

data/lib/egis/cartesian_product_generator.rb

@@ -0,0 +1,17 @@
+# frozen_string_literal: true
+
+module Egis
+  # @!visibility private
+  class CartesianProductGenerator
+    def cartesian_product(values_by_key)
+      keys = values_by_key.keys
+      values = values_by_key.values
+
+      head, *tail = values
+
+      return keys.zip(head) unless tail
+
+      head.product(*tail).map { |vals| keys.zip(vals) }
+    end
+  end
+end

data/lib/egis/client.rb

@@ -0,0 +1,142 @@
+# frozen_string_literal: true
+
+module Egis
+  ##
+  # The most fundamental {Egis} class. Provides an interface for executing Athena queries.
+  #
+  # See configuration instructions {Egis.configure}.
+  #
+  # @see Egis.configure
+  #
+  # @example Create client and execute asynchronous query
+  #   client = Egis::Client.new
+  #   status = client.execute_query('SELECT * FROM my_table;')
+  #
+  #   while status.in_progress?
+  #     # do something useful
+  #     # ...
+  #     status = client.query_status(status.id)
+  #   end
+  #
+  #   status.output_location.url # s3://my-bucket/result/path
+  #
+  # @example Execute synchronous query and fetch results
+  #   status = client.execute_query('SELECT MAX(time), MIN(id) FROM my_table;', async: false)
+  #   status.fetch_result(schema: [:timestamp, :int]) # [[2020-05-04 11:19:03 +0200, 7]]
+  #
+  class Client
+    QUERY_STATUS_MAPPING = {
+      'QUEUED' => Egis::QueryStatus::QUEUED,
+      'RUNNING' => Egis::QueryStatus::RUNNING,
+      'SUCCEEDED' => Egis::QueryStatus::FINISHED,
+      'FAILED' => Egis::QueryStatus::FAILED,
+      'CANCELLED' => Egis::QueryStatus::CANCELLED
+    }.freeze
+
+    DEFAULT_QUERY_STATUS_BACKOFF = ->(attempt) { 1.5**attempt - 1 }
+
+    private_constant :QUERY_STATUS_MAPPING, :DEFAULT_QUERY_STATUS_BACKOFF
+
+    def initialize(aws_client_provider: Egis::AwsClientProvider.new, s3_location_parser: Egis::S3LocationParser.new)
+      @aws_athena_client = aws_client_provider.athena_client
+      @s3_location_parser = s3_location_parser
+      @query_status_backoff = Egis.configuration.query_status_backoff || DEFAULT_QUERY_STATUS_BACKOFF
+    end
+
+    ##
+    # Creates {Egis::Database} object with a given name. Executing it doesn't create Athena database yet.
+    #
+    # @param [String] database_name
+    # @return [Egis::Database]
+
+    def database(database_name)
+      Database.new(database_name, client: self)
+    end
+
+    ##
+    # Executes Athena query. By default, queries are being executed asynchronously.
+    #
+    # @param [String] query SQL query to execute
+    # @param [Boolean] async Decide whether you want to run query asynchronously or block execution until it finishes
+    # @param [String] work_group Change Athena work group the query will be executed in.
+    # @param [String] database Run query in the context of a specific database (implicit table references are expected
+    #   to be in given database).
+    # @param [String] output_location S3 url of the desired output location. By default, Athena uses location defined in
+    #   by workgroup.
+    # @return [Egis::QueryStatus]
+
+    def execute_query(query, work_group: nil, database: nil, output_location: nil, async: true)
+      query_execution_id = aws_athena_client.start_query_execution(
+        query_execution_params(query, work_group, database, output_location)
+      ).query_execution_id
+
+      return query_status(query_execution_id) if Egis.mode.async(async)
+
+      query_status = wait_for_query_to_finish(query_execution_id)
+
+      raise Egis::Errors::QueryExecutionError, query_status.message unless query_status.finished?
+
+      query_status
+    end
+
+    ##
+    # Check the status of asynchronous query execution.
+    #
+    # @param [String] query_id Query id from {Egis::QueryStatus} returned by {#execute_query} method
+    # @return [Egis::QueryStatus]
+
+    def query_status(query_id)
+      resp = aws_athena_client.get_query_execution(query_execution_id: query_id)
+
+      query_execution = resp.query_execution
+
+      Egis::QueryStatus.new(
+        query_execution.query_execution_id,
+        QUERY_STATUS_MAPPING.fetch(query_execution.status.state),
+        query_execution.status.state_change_reason,
+        parse_output_location(query_execution)
+      )
+    end
+
+    private
+
+    attr_reader :aws_athena_client, :s3_location_parser, :query_status_backoff
+
+    def query_execution_params(query, work_group, database, output_location)
+      work_group_params = work_group || Egis.configuration.work_group
+
+      params = {query_string: query}
+      params[:work_group] = work_group_params if work_group_params
+      params[:query_execution_context] = {database: database_name(database)} if database
+      params[:result_configuration] = {output_location: translate_path(output_location)} if output_location
+      params
+    end
+
+    def wait_for_query_to_finish(query_execution_id)
+      attempt = 1
+      loop do
+        sleep(query_status_backoff.call(attempt))
+        status = query_status(query_execution_id)
+        return status unless status.queued? || status.running?
+
+        attempt += 1
+      end
+    end
+
+    def parse_output_location(query_execution)
+      url = query_execution.result_configuration.output_location
+
+      bucket, path = s3_location_parser.parse_url(url)
+
+      QueryOutputLocation.new(url, bucket, path)
+    end
+
+    def translate_path(s3_url)
+      Egis.mode.s3_path(s3_url)
+    end
+
+    def database_name(name)
+      Egis.mode.database_name(name)
+    end
+  end
+end

data/lib/egis/configuration.rb

@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+module Egis
+  # @!visibility private
+  class Configuration
+    attr_accessor :work_group, :aws_region, :aws_access_key_id, :aws_secret_access_key, :aws_profile,
+                  :query_status_backoff, :testing_s3_bucket
+  end
+end

data/lib/egis/database.rb

@@ -0,0 +1,102 @@
+# frozen_string_literal: true
+
+module Egis
+  ##
+  # Interface for database manipulation and querying.
+  #
+  # Extends the interface of {Egis::Client} but all the queries scheduled using {Egis::Database} are executed
+  # within the database's context. SQL table references without explicit database will implicitly refer to
+  # the database they are executed from.
+  #
+  # It is recommended to create database objects using {Egis::Client#database} method.
+  #
+  class Database
+    def initialize(database_name, client: Egis::Client.new, output_downloader: Egis::OutputDownloader.new)
+      @client = client
+      @database_name = database_name
+      @output_downloader = output_downloader
+    end
+
+    ##
+    # Creates {Egis::Table} object. Executing it doesn't create Athena table yet.
+    #
+    # @param [String] table_name
+    # @param [Egis::TableSchema] table_schema
+    # @param [String] table_location S3 URL with table location (e.g. `s3://s3_bucket/table/location/`)
+    # @param [:tsv, :csv, :orc] format Table format (defaults to :tsv)
+    # @return [Egis::Table]
+
+    def table(table_name, table_schema, table_location, **options)
+      Table.new(self, table_name, table_schema, table_location, options: options)
+    end
+
+    ##
+    # Creates database in Athena.
+    #
+    # @return [void]
+
+    def create
+      client.execute_query("CREATE DATABASE IF NOT EXISTS #{translate_name(database_name)};", async: false)
+    end
+
+    ##
+    # The same as {#create} but raising error if it already exists.
+    #
+    # @return [void]
+
+    def create!
+      client.execute_query("CREATE DATABASE #{translate_name(database_name)};", async: false)
+    end
+
+    ##
+    # Removes database in Athena.
+    #
+    # @return [void]
+
+    def drop
+      client.execute_query("DROP DATABASE IF EXISTS #{translate_name(database_name)} CASCADE;", async: false)
+    end
+
+    ##
+    # The same as {#drop} but raising error if it the database does not exist.
+    #
+    # @return [void]
+
+    def drop!
+      client.execute_query("DROP DATABASE #{translate_name(database_name)} CASCADE;", async: false)
+    end
+
+    ##
+    # (see Egis::Client#execute_query)
+
+    def execute_query(query, **options)
+      client.execute_query(query, **{database: database_name, **options})
+    end
+
+    ##
+    # (see Egis::Client#query_status)
+
+    def query_status(query_id)
+      client.query_status(query_id)
+    end
+
+    ##
+    # Checks whether database with such name exists in Athena.
+    #
+    # @return [Boolean]
+
+    def exists?
+      query_status = client.execute_query("SHOW DATABASES LIKE '#{database_name}';", async: false)
+      parsed_result = output_downloader.download(query_status.output_location)
+      parsed_result.flatten.include?(database_name)
+    end
+
+    private
+
+    attr_reader :client, :database_name, :output_downloader
+
+    def translate_name(name)
+      Egis.mode.database_name(name)
+    end
+  end
+end

data/lib/egis/errors.rb

@@ -0,0 +1,12 @@
+# frozen_string_literal: true
+
+module Egis
+  module Errors
+    class Error < StandardError; end
+
+    class UnsupportedTableFormat < Error; end
+    class QueryExecutionError < Error; end
+    class PartitionError < Error; end
+    class TypeError < Error; end
+  end
+end

data/lib/egis/output_downloader.rb

@@ -0,0 +1,21 @@
+# frozen_string_literal: true
+
+require 'csv'
+
+module Egis
+  # @!visibility private
+  class OutputDownloader
+    def initialize(aws_client_provider: Egis::AwsClientProvider.new)
+      @s3_client = aws_client_provider.s3_client
+    end
+
+    def download(output_location)
+      query_result = s3_client.get_object(bucket: output_location.bucket, key: output_location.key)
+      CSV.parse(query_result.body.read)
+    end
+
+    private
+
+    attr_reader :s3_client
+  end
+end

data/lib/egis/output_parser.rb

@@ -0,0 +1,24 @@
+# frozen_string_literal: true
+
+module Egis
+  # @!visibility private
+  class OutputParser
+    def parse(output, types)
+      header, *content = output
+
+      serializers = serializers(header, types)
+
+      content.map do |row|
+        row.zip(serializers).map do |string, serializer|
+          serializer.load(string)
+        end
+      end
+    end
+
+    private
+
+    def serializers(row, types)
+      row.zip(types).map { |_, type| type ? Types.serializer(type) : Types::DefaultSerializer.new }
+    end
+  end
+end

data/lib/egis/partitions_generator.rb

@@ -0,0 +1,55 @@
+# frozen_string_literal: true
+
+module Egis
+  # @!visibility private
+  class PartitionsGenerator
+    def initialize(cartesian_product_generator: Egis::CartesianProductGenerator.new)
+      @cartesian_product_generator = cartesian_product_generator
+    end
+
+    def to_sql(table_name, values_by_partition, permissive: false)
+      validate_partition_values(values_by_partition)
+
+      <<~SQL
+        ALTER TABLE #{table_name} ADD #{permissive_statement(permissive)}
+          #{partitions_definition(values_by_partition)};
+      SQL
+    end
+
+    private
+
+    attr_reader :cartesian_product_generator
+
+    def validate_partition_values(values_by_partition)
+      raise Errors::PartitionError, 'Partition value(s) missing' if partition_values_missing?(values_by_partition)
+    end
+
+    def partition_values_missing?(values_by_partition)
+      values_by_partition.nil? || values_by_partition.empty? || values_by_partition.values.any?(&:empty?)
+    end
+
+    def permissive_statement(permissive)
+      'IF NOT EXISTS' if permissive
+    end
+
+    def partitions_definition(values_by_partition)
+      cartesian_product_generator.cartesian_product(values_by_partition).
+        map { |partition_values_combination| partition_values_clause(partition_values_combination) }.
+        join("\n")
+    end
+
+    def partition_values_clause(partition_values_combination)
+      "PARTITION (#{partition_values(partition_values_combination).join(', ')})"
+    end
+
+    def partition_values(partition_values_combination)
+      partition_values_combination.map do |partition_name, value|
+        if value.is_a?(String)
+          "#{partition_name} = '#{value}'"
+        else
+          "#{partition_name} = #{value}"
+        end
+      end
+    end
+  end
+end

data/lib/egis/query_output_location.rb

@@ -0,0 +1,13 @@
+# frozen_string_literal: true
+
+module Egis
+  ##
+  # @!attribute [r] url
+  #   @return [String] Query output file's URL
+  # @!attribute [r] bucket
+  #   @return [String] Query output's S3 bucket
+  # @!attribute [r] key
+  #   @return [String] Query output's S3 path
+
+  QueryOutputLocation = Struct.new(:url, :bucket, :key)
+end

data/lib/egis/query_status.rb

@@ -0,0 +1,76 @@
+# frozen_string_literal: true
+
+module Egis
+  ##
+  # @!attribute [r] id
+  #   @return [String] Athena query execution ID
+  # @!attribute [r] status
+  #   @return [:queued, :running, :finished, :failed, :cancelled]
+  # @!attribute [r] message
+  #   @return [String]
+  # @!attribute [r] output_location
+  #   @return [Egis::OutputLocation]
+  #
+  class QueryStatus
+    QUEUED = :queued
+    RUNNING = :running
+    FINISHED = :finished
+    FAILED = :failed
+    CANCELLED = :cancelled
+
+    STATUSES = [QUEUED, RUNNING, FINISHED, FAILED, CANCELLED].freeze
+
+    attr_reader :id, :status, :message, :output_location
+
+    def initialize(id, status, message, output_location,
+                   output_downloader: Egis::OutputDownloader.new,
+                   output_parser: Egis::OutputParser.new)
+      raise ArgumentError, "Unsupported status #{status}" unless STATUSES.include?(status)
+
+      @id = id
+      @status = status
+      @message = message
+      @output_location = output_location
+      @output_downloader = output_downloader
+      @output_parser = output_parser
+    end
+
+    def finished?
+      status == FINISHED
+    end
+
+    def failed?
+      status == FAILED
+    end
+
+    def queued?
+      status == QUEUED
+    end
+
+    def running?
+      status == RUNNING
+    end
+
+    def in_progress?
+      [RUNNING, QUEUED].include?(status)
+    end
+
+    ##
+    # Download query result.
+    #
+    # By default, Egis will just parse output CSV and return array of string arrays. Additionally, you
+    # can pass expected query result column types to parse them into Ruby objects accordingly.
+    #
+    # @param [Array] schema Array with expected query column types
+    # @return [Array] Array of row values
+
+    def fetch_result(schema: [])
+      output = output_downloader.download(output_location)
+      output_parser.parse(output, schema)
+    end
+
+    private
+
+    attr_reader :output_downloader, :output_parser
+  end
+end

data/lib/egis/s3_cleaner.rb

@@ -0,0 +1,22 @@
+# frozen_string_literal: true
+
+module Egis
+  # @!visibility private
+  class S3Cleaner
+    def initialize(aws_client_provider: Egis::AwsClientProvider.new)
+      @s3_client = aws_client_provider.s3_client
+    end
+
+    def delete(bucket, prefix)
+      prefix_contents = s3_client.list_objects_v2(bucket: bucket, prefix: prefix).contents
+      return if prefix_contents.empty?
+
+      objects_to_remove = prefix_contents.map { |content| {key: content.key} }
+      s3_client.delete_objects(bucket: bucket, delete: {objects: objects_to_remove})
+    end
+
+    private
+
+    attr_reader :s3_client
+  end
+end

data/lib/egis/s3_location_parser.rb

@@ -0,0 +1,14 @@
+# frozen_string_literal: true
+
+module Egis
+  # @!visibility private
+  class S3LocationParser
+    S3_URL_PATTERN = %r{^s3://(?<bucket>\S+?)/(?<key>\S+)$}.freeze
+
+    def parse_url(url)
+      matched_data = S3_URL_PATTERN.match(url)
+
+      [matched_data['bucket'], matched_data['key']]
+    end
+  end
+end

data/lib/egis/standard_mode.rb

@@ -0,0 +1,18 @@
+# frozen_string_literal: true
+
+module Egis
+  # @!visibility private
+  class StandardMode
+    def s3_path(s3_url)
+      s3_url
+    end
+
+    def database_name(name)
+      name
+    end
+
+    def async(async_flag)
+      async_flag
+    end
+  end
+end