aq 0.0.1 → 0.1.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA1:
-  metadata.gz: 42a8c4ed6011726e5d62b1926a5f20fba8e6e093
-  data.tar.gz: 6050ecc604778075a1a2cb523f51b4d694d5ba2a
+  metadata.gz: 044c3e56fb4aeb56f5ead7205d9919f4dfbe8efb
+  data.tar.gz: 8ed0ea0d85526f8439a127fc85abde8c52ae2a61
 SHA512:
-  metadata.gz: 3abd963cddbfbe9f1cd905985073183161102da1c1be47cf9a2cdf82643ffe6753d5cb418e9abe5ca379618a321ec875fee77243f6c1c5de0c7343e15069acc3
-  data.tar.gz: 67ad4c36da859a8586f02458ab0a0d92a8d52feda03a49a211e7be52774e7b614dd48d813dedc94b10da145981b115d3e96eb8e0995b4bc7b8fa5c6a5131cfd1
+  metadata.gz: 879e730fac2abee336ed7c8f8a888c849918bff440e808b744d6b08d1f690cf84323d86c0b99d540b23fc029bd9aa44e21bf5dd7f36730d6c116e021a7c23187
+  data.tar.gz: 3ad0511876952645be034aefb63841f261efd9516d0af2060f2fc6e95476b174afdb11699a390d3e9a52f68a802c4e0b6273f3dfc9da216652839458ffea4d11

data/.travis.yml

@@ -5,3 +5,4 @@ rvm:
   - 2.3.4
   - 2.4.2
 before_install: gem install bundler
+script: bundle exec rake test

data/README.md

@@ -12,12 +12,66 @@ $ gem install aq
 
 ## Usage
 
+All commands need `--bucket` option because Athena stores query result into S3.
+You can specify it by `AQ_DEFAULT_BUCKET` environment variable.
+
+### help
+
+Display command help
+
+```bash
+$ aq help
+$ aq help [COMMAND]
+```
+
+### ls
+
+Show databases or tables in specified database
+
+```bash
+$ aq ls
+$ aq ls my_database_name
+```
+
+### head
+
+Show records in specified table
+
+```bash
+$ aq head my_db.my_table
+```
+
+### mk
+
+Create database
+
+```bash
+$ aq mk my_database_name
+```
+
+### load
+
+Create table and load data
+
+```bash
+$ aq load my_db.my_table s3://my_bucket/my_object_key/ test/resource/schema.json --partitioning dt:string
+```
+
+### rm
+
+Drop database or table
+
+```bash
+$ aq rm my_db
+$ aq rm my_db.my_table
+```
+
 ### query
 
 Run query
 
 ```bash
-$ aq query 'SELECT * FROM "test"."test_logs" limit 10;' --bucket=aws-athena-query-results-xxxx-region
+$ aq query 'SELECT * FROM "test"."test_logs" limit 10;'
 ```
 
 ## Development

data/Rakefile

@@ -3,7 +3,7 @@ require 'rake/testtask'
 
 Rake::TestTask.new(:test) do |t|
   t.libs << 'test' << 'lib'
-  t.test_files = FileList['test/**/*_test.rb']
+  t.test_files = FileList['test/**/test_*.rb']
 end
 
 task :default => :test

data/aq.gemspec

@@ -22,10 +22,11 @@ Gem::Specification.new do |spec|
   spec.executables = `git ls-files -- bin/*`.split("\n").map{ |f| File.basename(f) }
   spec.require_paths = ["lib"]
 
-  spec.add_dependency "thor"
+  spec.add_dependency "thor", "~> 0.0"
   spec.add_dependency "aws-sdk-athena", "~> 1.0"
   spec.add_dependency "aws-sdk-s3", "~> 1.0"
   spec.add_dependency "kosi", "~> 1.0"
+  spec.add_dependency "highline", "~> 1.0"
 
   spec.add_development_dependency "bundler"
   spec.add_development_dependency "rake"

data/lib/aq/command.rb

@@ -1,25 +1,73 @@
 # -*- coding: utf-8 -*-
 
 require 'date'
+require 'highline/import'
 require 'thor'
 
+require 'aq/query'
+require 'aq/query_builder'
+require 'aq/schema'
+require 'aq/version'
+
 module Aq
   class AqCmd < Thor
     package_name 'aq'
+    class_option :bucket, desc: 'S3 bucket where the query result is stored. This param can also be specified by AQ_DEFAULT_BUCKET environment variable.', default: ENV['AQ_DEFAULT_BUCKET']
+    class_option :object_prefix, desc: 'S3 object prefix where the query result is stored', default: "Unsaved/#{Date.today.strftime('%Y/%m/%d')}"
+
+    desc "ls [DATABASE]", "Show databases or tables in specified database"
+    def ls(database=nil)
+      query = QueryBuilder.ls database
+      Aq::Query.new(options[:bucket], options[:object_prefix]).run(query)
+    end
+
+    desc "head TABLE_NAME", "Show records in specified table"
+    option :max_rows, desc: 'This number of rows are printed', default: 100, type: :numeric, aliases: '-n'
+    def head(table)
+      query = QueryBuilder.head table, options[:max_rows]
+      Aq::Query.new(options[:bucket], options[:object_prefix]).run(query)
+    end
+
+    desc "mk NAME", "Create database"
+    def mk(name)
+      query = QueryBuilder.mk name
+      Aq::Query.new(options[:bucket], options[:object_prefix]).run(query)
+    end
+
+    desc "load TABLE SOURCE SCHEMA", "Create table and load data"
+    option :source_format, desc: 'Specify source file data format. Now aq support only NEWLINE_DELIMITED_JSON.', default: 'NEWLINE_DELIMITED_JSON'
+    option :partitioning, desc: 'Specify partition key and type. ex. key1:type1,key2:type2,...', default: nil
+    def load(table, source, schema)
+      schema = SchemaLoader.new.load schema
+      query = QueryBuilder.load table, source, schema, options[:source_format], options[:partitioning]
+      Aq::Query.new(options[:bucket], options[:object_prefix]).run(query)
+    end
+
+    desc "rm NAME", "Drop database or table"
+    option :force, desc: 'Skip confirmation if this is set', type: :boolean, aliases: '-f'
+    def rm(name)
+      if options[:force]
+        answer = true
+      else
+        answer = agree("Would you remove #{name}? (y)es or (n)o")
+      end
+      unless answer
+        puts 'Canceled'
+        exit 1
+      end
+
+      query = QueryBuilder.rm name
+      Aq::Query.new(options[:bucket], options[:object_prefix]).run(query)
+    end
 
     desc "query QUERY", "Run the query"
-    option :bucket, desc: 'S3 bucket where the query result is stored.', required: true
-    option :object_prefix, desc: 'S3 object prefix where the query result is stored. Default is `Unsaved/%Y/%m/%d`'
     option :timeout, desc: 'Wait for execution of the query for this number of seconds', default: nil, type: :numeric
     def query(query)
-      require 'aq/query'
-      object_prefix = options[:object_prefix] || "Unsaved/#{Date.today.strftime("%Y/%m/%d")}"
-      Aq::Query.new(options[:bucket], object_prefix).run(query, options[:timeout])
+      Aq::Query.new(options[:bucket], options[:object_prefix]).run(query, options[:timeout])
     end
 
     desc "version", "Print the version"
     def version
-      require 'aq/version'
       puts "aq v#{Aq::VERSION}"
     end
   end

data/lib/aq/error.rb

@@ -0,0 +1,5 @@
+module Aq
+  class InvalidParameterError < StandardError; end
+  class NotSupportedError < StandardError; end
+  class NotImplementedError < StandardError; end
+end

data/lib/aq/query.rb

@@ -1,6 +1,7 @@
 # -*- coding: utf-8 -*-
 
 require 'aq/base'
+require 'aq/error'
 require 'aws-sdk-s3'
 require 'csv'
 require 'kosi'
@@ -10,6 +11,7 @@ module Aq
   class Query < Base
     def initialize(bucket, object_prefix)
       super()
+      raise InvalidParameterError.new '`bucket` must be specified.' if bucket.nil? || bucket.empty?
       @bucket = bucket
       @object_prefix = object_prefix
     end
@@ -79,7 +81,7 @@ module Aq
         kosi = Kosi::Table.new({header: csv.headers})
         print kosi.render(result)
       else
-        print body.read
+        print body.read + "\n"
       end
     end
 

data/lib/aq/query_builder.rb

@@ -0,0 +1,53 @@
+require 'aq/error'
+
+module Aq
+  class QueryBuilder
+    def self.ls(database)
+      if database.nil?
+        'SHOW DATABASES'
+      else
+        "SHOW TABLES IN #{database}"
+      end
+    end
+
+    def self.head(table, num)
+      raise InvalidParameterError.new 'The table name must be specified in the format `DATABSE.TABLE`' unless table.include? '.'
+      "SELECT * FROM #{table} LIMIT #{num}"
+    end
+
+    def self.mk(name)
+      if !name.include? '.'
+        "CREATE DATABASE IF NOT EXISTS #{name}"
+      else
+        raise InvalidParameterError.new 'Use `load` command if you create new table.'
+      end
+    end
+
+    def self.load(table, source, schema, source_format, patition)
+      raise InvalidParameterError.new '`SOURCE` must start with "s3://"' unless source.start_with? 's3://'
+      schema_state = schema.get_all.map do |s|
+        "`#{s[:name]}` #{s[:type]}"
+      end.join(',')
+      raise InvalidParameterError.new 'Now aq support only NEWLINE_DELIMITED_JSON.' unless %w(NEWLINE_DELIMITED_JSON).include? source_format
+      serde = 'org.apache.hive.hcatalog.data.JsonSerDe'
+      patition_state = patition.nil? ? '' : "PARTITIONED BY (#{patition.gsub(':', ' ')})"
+
+      <<"SQL"
+CREATE EXTERNAL TABLE IF NOT EXISTS #{table} (
+#{schema_state}
+)
+#{patition_state}
+ROW FORMAT SERDE "#{serde}"
+LOCATION "#{source}"
+SQL
+    end
+
+    def self.rm(name)
+      if !name.include? '.'
+        "DROP DATABASE IF EXISTS #{name}"
+      else
+        "DROP TABLE IF EXISTS #{name}"
+      end
+    end
+  end
+end

data/lib/aq/schema.rb

@@ -0,0 +1,89 @@
+# -*- coding: utf-8 -*-
+
+require 'json'
+require 'aq/error'
+
+module Aq
+  class Schema
+    def initialize
+      @schema = []
+    end
+
+    def append_column(name, type, mode=nil)
+      nullable = mode != 'required'
+      column = {name: name, type: convert_type_from_bq_to_athena(type), nullable: nullable}
+      @schema.push column
+    end
+
+    def get_all
+      @schema
+    end
+
+    private
+    def convert_type_from_bq_to_athena(type)
+=begin
+      Schema correspondence table
+      ----------+----------
+      BQ        | Athena
+      ----------+----------
+      STRING    | STRING
+      BYTES     | ?
+      INTEGER   | INT
+      FLOAT     | DOUBLE
+      BOOLEAN   | BOOLEAN
+      RECORD    | ARRAY or MAP
+      TIMESTAMP | TIMESTAMP
+      DATE      | DATE
+      TIME      | TIMESTAMP
+      DATETIME  | TIMESTAMP
+=end
+      type.upcase!
+      case type
+      when 'BYTES'
+        raise NotSupportedError.new '`BYTES` is not supported in Athena.'
+      when 'INTEGER'
+        'BIGINT'
+      when 'FLOAT'
+        'DOUBLE'
+      when 'RECORD'
+        raise NotImplementedError.new 'Sorry, `RECORD` is not supported yet in aq.'
+      when 'TIME', 'DATETIME'
+        'TIMESTAMP'
+      else
+        type
+      end
+    end
+  end
+
+  class SchemaLoader
+    def load(schema)
+      if File.exist? File.expand_path(schema)
+        load_from_file File.expand_path(schema)
+      else
+        load_from_string schema
+      end
+    end
+
+    private
+    def load_from_file(file_path)
+      schema = Schema.new
+      JSON.load(File.open(file_path).read).each do |c|
+        if c.has_key? 'mode'
+          schema.append_column c['name'], c['type'], c['mode']
+        else
+          schema.append_column c['name'], c['type']
+        end
+      end
+      schema
+    end
+
+    def load_from_string(str)
+      schema = Schema.new
+      str.split(',').each do |column|
+        c = column.split(':')
+        schema.append_column c[0], c[1]
+      end
+      schema
+    end
+  end
+end

data/lib/aq/version.rb

@@ -1,3 +1,3 @@
 module Aq
-  VERSION = '0.0.1'
+  VERSION = '0.1.0'
 end

data/test/aq/test_query_builder.rb

@@ -0,0 +1,98 @@
+require_relative '../test_helper'
+require 'aq/error'
+require 'aq/query_builder'
+require 'aq/schema'
+
+module Aq
+  class QueryBuilderTest < Test::Unit::TestCase
+    sub_test_case 'ls' do
+      test 'show databases' do
+        query = QueryBuilder.ls nil
+        expected = 'SHOW DATABASES'
+        assert_equal(expected, query)
+      end
+
+      test 'show tables' do
+        query = QueryBuilder.ls 'my_db'
+        expected = 'SHOW TABLES IN my_db'
+        assert_equal(expected, query)
+      end
+    end
+
+    sub_test_case 'head' do
+      test 'select' do
+        query = QueryBuilder.head 'my_db.my_table', 10
+        expected = 'SELECT * FROM my_db.my_table LIMIT 10'
+        assert_equal(expected, query)
+      end
+
+      test 'invalid table name' do
+        assert_raise InvalidParameterError do
+          QueryBuilder.head 'my_table', 10
+        end
+      end
+    end
+
+    sub_test_case 'mk' do
+      test 'create database' do
+        query = QueryBuilder.mk 'my_db'
+        expected = 'CREATE DATABASE IF NOT EXISTS my_db'
+        assert_equal(expected, query)
+      end
+
+      test 'create table' do
+        assert_raise InvalidParameterError do
+          QueryBuilder.mk 'my_db.my_table'
+        end
+      end
+    end
+
+    sub_test_case 'load' do
+      test 'create table with partition' do
+        schema = SchemaLoader.new.load 'test/resource/schema.json'
+        query = QueryBuilder.load(
+          'my_db.my_table', 's3://hoge/foo/', schema, 'NEWLINE_DELIMITED_JSON', 'dt:string'
+        )
+        expected = <<'SQL'
+CREATE EXTERNAL TABLE IF NOT EXISTS my_db.my_table (
+`str0` STRING,`str1` STRING,`str2` STRING
+)
+PARTITIONED BY (dt string)
+ROW FORMAT SERDE "org.apache.hive.hcatalog.data.JsonSerDe"
+LOCATION "s3://hoge/foo/"
+SQL
+        assert_equal(expected, query)
+      end
+
+      test 'create table without partition' do
+        schema = SchemaLoader.new.load 'test/resource/schema.json'
+        query = QueryBuilder.load(
+          'my_db.my_table', 's3://hoge/foo/', schema, 'NEWLINE_DELIMITED_JSON', nil
+        )
+        expected = <<'SQL'
+CREATE EXTERNAL TABLE IF NOT EXISTS my_db.my_table (
+`str0` STRING,`str1` STRING,`str2` STRING
+)
+
+ROW FORMAT SERDE "org.apache.hive.hcatalog.data.JsonSerDe"
+LOCATION "s3://hoge/foo/"
+SQL
+        assert_equal(expected, query)
+      end
+    end
+
+    sub_test_case 'rm' do
+      test 'drop database' do
+        query = QueryBuilder.rm 'my_db'
+        expected = 'DROP DATABASE IF EXISTS my_db'
+        assert_equal(expected, query)
+      end
+
+      test 'drop table' do
+        query = QueryBuilder.rm 'my_db.my_table'
+        expected = 'DROP TABLE IF EXISTS my_db.my_table'
+        assert_equal(expected, query)
+      end
+    end
+  end
+end

data/test/aq/test_schema.rb

@@ -0,0 +1,77 @@
+require_relative '../test_helper'
+require 'aq/error'
+require 'aq/schema'
+
+module Aq
+  class SchemaTest < Test::Unit::TestCase
+    setup do
+      @schema = Schema.new
+    end
+
+    test 'ok' do
+      @schema.append_column 'str0', 'string'
+      @schema.append_column 'str1', 'string', 'required'
+      @schema.append_column 'str2', 'string', 'nullable'
+      @schema.append_column 'num',  'integer'
+      @schema.append_column 'f',    'float'
+      @schema.append_column 'bool', 'boolean'
+      @schema.append_column 'ts',   'timestamp'
+      @schema.append_column 'd',    'date'
+      @schema.append_column 't',    'time'
+      @schema.append_column 'dt',   'datetime'
+
+      expected = [
+        {name: 'str0', type: 'STRING'   , nullable: true },
+        {name: 'str1', type: 'STRING'   , nullable: false},
+        {name: 'str2', type: 'STRING'   , nullable: true },
+        {name: 'num',  type: 'BIGINT'   , nullable: true },
+        {name: 'f',    type: 'DOUBLE'   , nullable: true },
+        {name: 'bool', type: 'BOOLEAN'  , nullable: true },
+        {name: 'ts',   type: 'TIMESTAMP', nullable: true },
+        {name: 'd',    type: 'DATE'     , nullable: true },
+        {name: 't',    type: 'TIMESTAMP', nullable: true },
+        {name: 'dt',   type: 'TIMESTAMP', nullable: true },
+      ]
+
+      assert_equal(expected, @schema.get_all)
+    end
+
+    test 'not supported type' do
+      assert_raise NotSupportedError do
+        @schema.append_column 'ng',  'bytes'
+      end
+    end
+
+    test 'not implemented type' do
+      assert_raise NotImplementedError do
+        @schema.append_column 'ng',  'record'
+      end
+    end
+  end
+
+  class SchemaLoaderTest < Test::Unit::TestCase
+    setup do
+      @loader = SchemaLoader.new
+    end
+
+    test 'from file' do
+      result = @loader.load('test/resource/schema.json').get_all
+      expected = [
+        {name: 'str0', type: 'STRING', nullable: false},
+        {name: 'str1', type: 'STRING', nullable: true },
+        {name: 'str2', type: 'STRING', nullable: true }
+      ]
+      assert_equal(expected, result)
+    end
+
+    test 'from string' do
+      result = @loader.load('str0:STRING,str1:STRING,str2:STRING').get_all
+      expected = [
+        {name: 'str0', type: 'STRING', nullable: true},
+        {name: 'str1', type: 'STRING', nullable: true },
+        {name: 'str2', type: 'STRING', nullable: true }
+      ]
+      assert_equal(expected, result)
+    end
+  end
+end

data/test/resource/schema.json

@@ -0,0 +1,5 @@
+[
+  {"name": "str0", "type": "string", "mode": "required"},
+  {"name": "str1", "type": "string", "mode": "nullable"},
+  {"name": "str2", "type": "string"}
+]

data/test/test_helper.rb

@@ -11,7 +11,6 @@ require 'test/unit'
 require 'test/unit/rr'
 
 $LOAD_PATH.unshift File.expand_path("../../lib", __FILE__)
-require 'aq'
 
 class Test::Unit::TestCase
 end

metadata

@@ -1,29 +1,29 @@
 --- !ruby/object:Gem::Specification
 name: aq
 version: !ruby/object:Gem::Version
-  version: 0.0.1
+  version: 0.1.0
 platform: ruby
 authors:
 - Yoshihiro MIYAI
 autorequire: 
 bindir: bin
 cert_chain: []
-date: 2017-10-31 00:00:00.000000000 Z
+date: 2017-11-13 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: thor
   requirement: !ruby/object:Gem::Requirement
     requirements:
-    - - ">="
+    - - "~>"
       - !ruby/object:Gem::Version
-        version: '0'
+        version: '0.0'
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
-    - - ">="
+    - - "~>"
       - !ruby/object:Gem::Version
-        version: '0'
+        version: '0.0'
 - !ruby/object:Gem::Dependency
   name: aws-sdk-athena
   requirement: !ruby/object:Gem::Requirement
@@ -66,6 +66,20 @@ dependencies:
     - - "~>"
       - !ruby/object:Gem::Version
         version: '1.0'
+- !ruby/object:Gem::Dependency
+  name: highline
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '1.0'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '1.0'
 - !ruby/object:Gem::Dependency
   name: bundler
   requirement: !ruby/object:Gem::Requirement
@@ -140,9 +154,14 @@ files:
 - bin/aq
 - lib/aq/base.rb
 - lib/aq/command.rb
+- lib/aq/error.rb
 - lib/aq/query.rb
+- lib/aq/query_builder.rb
+- lib/aq/schema.rb
 - lib/aq/version.rb
-- test/aq/aq_test.rb
+- test/aq/test_query_builder.rb
+- test/aq/test_schema.rb
+- test/resource/schema.json
 - test/test_helper.rb
 homepage: https://github.com/mia-0032/aq
 licenses:
@@ -169,5 +188,7 @@ signing_key:
 specification_version: 4
 summary: Command Line Tool for AWS Athena (bq command like)
 test_files:
-- test/aq/aq_test.rb
+- test/aq/test_query_builder.rb
+- test/aq/test_schema.rb
+- test/resource/schema.json
 - test/test_helper.rb

data/test/aq/aq_test.rb

@@ -1,11 +0,0 @@
-require_relative '../test_helper'
-
-class AqTest < Test::Unit::TestCase
-  def test_that_it_has_a_version_number
-    refute_nil ::Aq::VERSION
-  end
-
-  def test_it_does_something_useful
-    assert false
-  end
-end