egis 1.1.0

data/lib/egis/table.rb

@@ -0,0 +1,163 @@
+# frozen_string_literal: true
+
+module Egis
+  ##
+  # Interface for Athena table manipulation.
+  #
+  # It is recommended to create table objects using {Egis::Database#table} method.
+  #
+  # @!attribute [r] database
+  #   @return [Egis::Database]
+  # @!attribute [r] name
+  #   @return [String] Athena database name
+  # @!attribute [r] schema
+  #   @return [Egis::TableSchema] table's schema object
+  #
+  class Table
+    DEFAULT_OPTIONS = {format: :tsv}.freeze
+
+    def initialize(database, name, schema, location, options: {},
+                   partitions_generator: Egis::PartitionsGenerator.new,
+                   table_ddl_generator: Egis::TableDDLGenerator.new,
+                   output_downloader: Egis::OutputDownloader.new,
+                   output_parser: Egis::OutputParser.new,
+                   table_data_wiper: Egis::TableDataWiper.new)
+      @database = database
+      @name = name
+      @schema = schema
+      @location = location
+      @options = DEFAULT_OPTIONS.merge(options)
+      @partitions_generator = partitions_generator
+      @table_ddl_generator = table_ddl_generator
+      @output_downloader = output_downloader
+      @output_parser = output_parser
+      @table_data_wiper = table_data_wiper
+    end
+
+    attr_reader :database, :name, :schema
+
+    ##
+    # Creates table in Athena.
+    #
+    # @return [void]
+
+    def create
+      create_table_sql = table_ddl_generator.create_table_sql(self, permissive: true)
+      database.execute_query(create_table_sql, async: false)
+    end
+
+    ##
+    # The same as {#create} but raising error when table with a given name already exists.
+    #
+    # @return [void]
+
+    def create!
+      create_table_sql = table_ddl_generator.create_table_sql(self, permissive: false)
+      database.execute_query(create_table_sql, async: false)
+    end
+
+    ##
+    # Creates partitions with all possible combinations of given partition values.
+    #
+    # @example
+    #   table.add_partitions(year: [2000, 2001], type: ['user'])
+    #
+    # @param [Hash] partitions
+    # @return [void]
+
+    def add_partitions(partitions)
+      load_partitions_query = partitions_generator.to_sql(name, partitions, permissive: true)
+      database.execute_query(load_partitions_query, async: false)
+    end
+
+    ##
+    # (see add_partitions)
+    # It raises error when a partition already exists.
+
+    def add_partitions!(partitions)
+      load_partitions_query = partitions_generator.to_sql(name, partitions, permissive: false)
+      database.execute_query(load_partitions_query, async: false)
+    end
+
+    ##
+    # Tells Athena to automatically discover table's partitions by scanning table's S3 location.
+    # This operation might take long time with big number of partitions. If that's the case, instead of this method use
+    # {#add_partitions} to define partitions manually.
+    #
+    # @return [void]
+
+    def discover_partitions
+      database.execute_query("MSCK REPAIR TABLE #{name};", async: false)
+    end
+
+    ##
+    # Insert data into the table. Mostly useful for testing purposes.
+    #
+    # @param [Array] rows Array of arrays with row values
+    # @return [void]
+
+    def upload_data(rows)
+      query = data_insert_query(rows)
+      database.execute_query(query, async: false)
+    end
+
+    ##
+    # Downloads table contents into memory. Mostly useful for testing purposes.
+    #
+    # @return [Array] Array of arrays with row values.
+
+    def download_data
+      result = database.execute_query("SELECT * FROM #{name};", async: false)
+      content = output_downloader.download(result.output_location)
+      output_parser.parse(content, column_types)
+    end
+
+    ##
+    # Removes table's content on S3. Optionally, you can limit files removed to specific partitions.
+    #
+    # @param [Hash] partitions Partitions values to remove. Follows the same argument format as {#add_partitions}.
+    # @return [void]
+
+    def wipe_data(partitions: nil)
+      table_data_wiper.wipe_table_data(self, partitions)
+    end
+
+    ##
+    # @return Table data format
+
+    def format
+      options.fetch(:format)
+    end
+
+    ##
+    # @return [String] table location URL
+
+    def location
+      Egis.mode.s3_path(@location)
+    end
+
+    private
+
+    attr_reader :options, :partitions_generator, :table_ddl_generator, :output_downloader, :output_parser,
+                :table_data_wiper
+
+    def column_serializers
+      @column_serializers ||= column_types.map { |type| Egis::Types.serializer(type) }
+    end
+
+    def column_types
+      (schema.columns + schema.partitions).map(&:type)
+    end
+
+    def data_insert_query(rows)
+      <<~SQL
+        INSERT INTO #{name} VALUES
+        #{rows.map { |row| row_values_statement(row) }.join(",\n")};
+      SQL
+    end
+
+    def row_values_statement(row)
+      "(#{row.zip(column_serializers).map { |value, serializer| serializer.literal(value) }.join(', ')})"
+    end
+  end
+end

data/lib/egis/table_data_wiper.rb

@@ -0,0 +1,51 @@
+# frozen_string_literal: true
+
+module Egis
+  # @!visibility private
+  class TableDataWiper
+    def initialize(s3_location_parser: Egis::S3LocationParser.new,
+                   s3_cleaner: Egis::S3Cleaner.new,
+                   cartesian_product_generator: Egis::CartesianProductGenerator.new)
+      @s3_location_parser = s3_location_parser
+      @s3_cleaner = s3_cleaner
+      @cartesian_product_generator = cartesian_product_generator
+    end
+
+    def wipe_table_data(table, partitions)
+      bucket, location = s3_location_parser.parse_url(table.location)
+
+      return s3_cleaner.delete(bucket, location) unless partitions
+
+      partition_values_to_remove = partition_values_to_remove(table, partitions)
+
+      validate_partition_values(partition_values_to_remove, partitions)
+
+      remove_partition_files(bucket, location, partition_values_to_remove)
+    end
+
+    private
+
+    attr_reader :s3_location_parser, :s3_cleaner, :cartesian_product_generator
+
+    def partition_values_to_remove(table, partitions)
+      table_partitions = table.schema.partitions.map(&:name)
+      given_partitions = partitions.keys
+
+      partitions_to_delete = table_partitions.take_while { |partition| given_partitions.include?(partition) }
+      partitions_to_delete.map { |partition_name| [partition_name, partitions.fetch(partition_name)] }.to_h
+    end
+
+    def validate_partition_values(removed_partition_values, partitions)
+      return unless removed_partition_values.empty? || removed_partition_values.values.any?(&:empty?)
+
+      raise Egis::Errors::PartitionError, "Incorrect partitions given: #{partitions}"
+    end
+
+    def remove_partition_files(bucket, location, partitions_with_values)
+      cartesian_product_generator.cartesian_product(partitions_with_values).each do |partition_value_set|
+        partition_prefix = partition_value_set.map { |name_value| name_value.join('=') }.join('/')
+        s3_cleaner.delete(bucket, "#{location}/#{partition_prefix}")
+      end
+    end
+  end
+end

data/lib/egis/table_ddl_generator.rb

@@ -0,0 +1,50 @@
+# frozen_string_literal: true
+
+module Egis
+  # @!visibility private
+  class TableDDLGenerator
+    def create_table_sql(table, permissive: false)
+      <<~SQL
+        CREATE EXTERNAL TABLE #{permissive_statement(permissive)}#{table.name} (
+          #{column_definition_sql(table.schema.columns)}
+        )
+        #{partition_statement(table.schema)}
+        #{format_statement(table.format)}
+        LOCATION '#{table.location}';
+      SQL
+    end
+
+    private
+
+    def permissive_statement(permissive_flag)
+      'IF NOT EXISTS ' if permissive_flag
+    end
+
+    def partition_statement(table_schema)
+      return if table_schema.partitions.empty?
+
+      <<~SQL
+        PARTITIONED BY (
+          #{column_definition_sql(table_schema.partitions)}
+        )
+      SQL
+    end
+
+    def column_definition_sql(columns)
+      columns.map { |column| "`#{column.name}` #{column.type}" }.join(",\n")
+    end
+
+    def format_statement(format)
+      case format
+      when :csv
+        "ROW FORMAT DELIMITED FIELDS TERMINATED BY ','"
+      when :tsv
+        "ROW FORMAT DELIMITED FIELDS TERMINATED BY '\\t'"
+      when :orc
+        'STORED AS ORC'
+      else
+        raise Errors::UnsupportedTableFormat, format.to_s
+      end
+    end
+  end
+end

data/lib/egis/table_schema.rb

@@ -0,0 +1,49 @@
+# frozen_string_literal: true
+
+module Egis
+  ##
+  # Provides DSL for defining table schemas.
+  #
+  # @example Table schema definition
+  #   schema = Egis::TableSchema.define do
+  #     column :id, :int
+  #     column :message, :string
+  #
+  #     partition :country, :string
+  #     partition :type, :int
+  #   end
+  #
+  # @!attribute [r] columns
+  #   @return [Egis::TableSchema::Column]
+  # @!attribute [r] partitions
+  #   @return [Egis::TableSchema::Column]
+  #
+  class TableSchema
+    ##
+    # @return [Egis::TableSchema]
+
+    def self.define(&block)
+      new(&block)
+    end
+
+    def initialize(&block)
+      @columns = []
+      @partitions = []
+      instance_eval(&block)
+    end
+
+    attr_reader :columns, :partitions
+
+    private
+
+    def column(name, type)
+      @columns << Column.new(name, type)
+    end
+
+    def partition(name, type)
+      @partitions << Column.new(name, type)
+    end
+
+    Column = Struct.new(:name, :type)
+  end
+end

data/lib/egis/testing/testing_mode.rb

@@ -0,0 +1,62 @@
+# frozen_string_literal: true
+
+module Egis
+  module Testing
+    # @!visibility private
+    class TestingMode
+      def initialize(test_id, s3_bucket,
+                     client: Egis::Client.new,
+                     output_downloader: Egis::OutputDownloader.new,
+                     s3_location_parser: Egis::S3LocationParser.new)
+        @test_id = test_id
+        @s3_bucket = s3_bucket
+        @dirty = false
+        @client = client
+        @output_downloader = output_downloader
+        @s3_location_parser = s3_location_parser
+      end
+
+      def s3_path(s3_url)
+        dirty!
+
+        bucket, key = s3_location_parser.parse_url(s3_url)
+
+        "s3://#{s3_bucket}/#{test_id}/#{bucket}/#{key}"
+      end
+
+      def database_name(name)
+        dirty!
+
+        "#{test_id}_#{name}"
+      end
+
+      def async(_async_flag)
+        dirty!
+
+        false
+      end
+
+      def cleanup
+        remove_test_databases if dirty?
+      end
+
+      private
+
+      attr_reader :test_id, :s3_bucket, :client, :output_downloader, :s3_location_parser
+
+      def remove_test_databases
+        result = client.execute_query("SHOW DATABASES LIKE '#{test_id}.*';", async: false)
+        query_result = output_downloader.download(result.output_location)
+        query_result.flatten.each { |database| client.database(database).drop }
+      end
+
+      def dirty!
+        @dirty = true
+      end
+
+      def dirty?
+        @dirty
+      end
+    end
+  end
+end

data/lib/egis/testing.rb

@@ -0,0 +1,48 @@
+# frozen_string_literal: true
+
+require 'securerandom'
+
+require 'egis/testing/testing_mode'
+
+module Egis # rubocop:disable Style/Documentation
+  # @!visibility private
+  module Testing
+  end
+
+  ##
+  # Egis testing mode.
+  # Every table and created within method's block is mapped to a "virtual" table space in your testing S3 bucket.
+  # Using it, you can insert test data to your production tables and they will be simulated within the testing bucket,
+  # not touching actual locations.
+  #
+  # @example RSpec configuration
+  #   # spec_helper.rb
+  #
+  #   require 'egis/testing'
+  #
+  #   Egis.configure do |config|
+  #     config.testing_s3_bucket = 'testing-bucket'
+  #   end
+  #
+  #   RSpec.configure do |config|
+  #     config.around(:each) do |example|
+  #       Egis.testing do
+  #         example.run
+  #       end
+  #     end
+  #   end
+  #
+  # @return [void]
+
+  def self.testing
+    test_id = SecureRandom.hex
+    test_mode = Egis::Testing::TestingMode.new(test_id, Egis.configuration.testing_s3_bucket)
+
+    previous_mode = Egis.mode
+    @mode = test_mode
+    yield
+  ensure
+    @mode = previous_mode
+    test_mode.cleanup
+  end
+end

data/lib/egis/types/boolean_serializer.rb

@@ -0,0 +1,53 @@
+# frozen_string_literal: true
+
+module Egis
+  module Types
+    # @!visibility private
+    class BooleanSerializer
+      TRUE_LITERAL = 'TRUE'
+      FALSE_LITERAL = 'FALSE'
+
+      TRUE_VALUE = 'true'
+      FALSE_VALUE = 'false'
+
+      def literal(value)
+        case value
+        when true
+          TRUE_LITERAL
+        when false
+          FALSE_LITERAL
+        else
+          illegal_value_error(value)
+        end
+      end
+
+      def dump(value)
+        case value
+        when true
+          TRUE_VALUE
+        when false
+          FALSE_VALUE
+        else
+          illegal_value_error(value)
+        end
+      end
+
+      def load(string)
+        case string
+        when TRUE_VALUE
+          true
+        when FALSE_VALUE
+          false
+        else
+          illegal_value_error(string)
+        end
+      end
+
+      private
+
+      def illegal_value_error(value)
+        raise Egis::TypeError, "Illegal value '#{value}' for type boolean"
+      end
+    end
+  end
+end

data/lib/egis/types/default_serializer.rb

@@ -0,0 +1,20 @@
+# frozen_string_literal: true
+
+module Egis
+  module Types
+    # @!visibility private
+    class DefaultSerializer
+      def literal(value)
+        "'#{value}'"
+      end
+
+      def dump(value)
+        value
+      end
+
+      def load(string)
+        string
+      end
+    end
+  end
+end

data/lib/egis/types/integer_serializer.rb

@@ -0,0 +1,20 @@
+# frozen_string_literal: true
+
+module Egis
+  module Types
+    # @!visibility private
+    class IntegerSerializer
+      def literal(integer)
+        integer.to_s
+      end
+
+      def dump(integer)
+        integer.to_s
+      end
+
+      def load(string)
+        string.to_i
+      end
+    end
+  end
+end

data/lib/egis/types/null_serializer.rb

@@ -0,0 +1,36 @@
+# frozen_string_literal: true
+
+module Egis
+  module Types
+    # @!visibility private
+    class NullSerializer
+      NULL_LITERAL = 'NULL'
+
+      def initialize(wrapped_serializer)
+        @wrapped_serializer = wrapped_serializer
+      end
+
+      def literal(value)
+        return NULL_LITERAL if value.nil?
+
+        wrapped_serializer.literal(value)
+      end
+
+      def dump(value)
+        return nil if value.nil?
+
+        wrapped_serializer.dump(value)
+      end
+
+      def load(string)
+        return nil if string.nil?
+
+        wrapped_serializer.load(string)
+      end
+
+      private
+
+      attr_reader :wrapped_serializer
+    end
+  end
+end

data/lib/egis/types/string_serializer.rb

@@ -0,0 +1,20 @@
+# frozen_string_literal: true
+
+module Egis
+  module Types
+    # @!visibility private
+    class StringSerializer
+      def literal(string)
+        "'#{string.gsub("'", "''")}'"
+      end
+
+      def dump(string)
+        string
+      end
+
+      def load(string)
+        string
+      end
+    end
+  end
+end

data/lib/egis/types/timestamp_serializer.rb

@@ -0,0 +1,22 @@
+# frozen_string_literal: true
+
+module Egis
+  module Types
+    # @!visibility private
+    class TimestampSerializer
+      ATHENA_TIME_FORMAT = '%Y-%m-%d %H:%M:%S'
+
+      def literal(time)
+        "timestamp '#{dump(time)}'"
+      end
+
+      def dump(time)
+        time.strftime(ATHENA_TIME_FORMAT)
+      end
+
+      def load(string)
+        Time.strptime(string, ATHENA_TIME_FORMAT)
+      end
+    end
+  end
+end

data/lib/egis/types.rb

@@ -0,0 +1,30 @@
+# frozen_string_literal: true
+
+require 'egis/types/boolean_serializer'
+require 'egis/types/default_serializer'
+require 'egis/types/integer_serializer'
+require 'egis/types/string_serializer'
+require 'egis/types/timestamp_serializer'
+require 'egis/types/null_serializer'
+
+module Egis
+  # @!visibility private
+  module Types
+    def self.serializer(type)
+      type_serializer = case type
+                        when :timestamp
+                          TimestampSerializer.new
+                        when :string
+                          StringSerializer.new
+                        when :int, :bigint
+                          IntegerSerializer.new
+                        when :boolean
+                          BooleanSerializer.new
+                        else
+                          raise Errors::TypeError, "Unsupported type: #{type}"
+                        end
+
+      NullSerializer.new(type_serializer)
+    end
+  end
+end

data/lib/egis/version.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+module Egis
+  VERSION = '1.1.0'
+end

data/lib/egis.rb

@@ -0,0 +1,62 @@
+# frozen_string_literal: true
+
+require 'egis/version'
+require 'egis/errors'
+require 'egis/configuration'
+require 'egis/types'
+require 'egis/query_status'
+require 'egis/aws_client_provider'
+require 'egis/s3_cleaner'
+require 'egis/output_downloader'
+require 'egis/output_parser'
+require 'egis/client'
+require 'egis/cartesian_product_generator'
+require 'egis/partitions_generator'
+require 'egis/table_data_wiper'
+require 'egis/table'
+require 'egis/database'
+require 'egis/query_output_location'
+require 'egis/table_ddl_generator'
+require 'egis/table_schema'
+require 'egis/standard_mode'
+require 'egis/s3_location_parser'
+
+##
+# Egis is configured using Egis.configure block.
+#
+# @example Configuration using AWS access key ID and secret
+#   Egis.configure do |config|
+#     config.aws_region = 'AWS region'
+#     config.aws_access_key_id = 'AWS key ID'
+#     config.aws_secret_access_key = 'AWS secret key'
+#     config.work_group = 'egis-integration-testing'
+#   end
+#
+# If you don't specify credentials they will be looked up in the default locations. For more information see
+# {https://docs.aws.amazon.com/sdk-for-ruby/v3/developer-guide/setup-config.html}
+#
+# @example Use specific credentials profile from `~/.aws/credentials`
+#   Egis.configure do |config|
+#     config.aws_profile = 'my-profile'
+#   end
+#
+# @yield [Egis::Configuration]
+# @return [void]
+#
+module Egis
+  class << self
+    def configure
+      yield(configuration)
+    end
+
+    # @!visibility private
+    def configuration
+      @configuration ||= Configuration.new
+    end
+
+    # @!visibility private
+    def mode
+      @mode ||= Egis::StandardMode.new
+    end
+  end
+end