athena-cli 0.1.0

data/lib/amazon_athena/commands/show_table_properties.rb

@@ -0,0 +1,36 @@
+require_relative '../command'
+
+module AmazonAthena
+  module Commands
+    class ShowTableProperties < AmazonAthena::Command
+
+      def initialize(database_table)
+        @database_table = database_table
+      end
+
+      def statement
+        "SHOW TBLPROPERTIES #{@database_table};"
+      end
+
+      def run(connection)
+        result = connection.query(statement).raw_output
+
+        data = Hash[*result.split("\n").map {|line| line.split("\t")}.flatten]
+
+        data[:name] = @database_table
+
+        if type = data.delete('EXTERNAL')
+          data[:external] = type
+        end
+
+        if last_modified = data.delete('transient_lastDdlTime')
+          data[:last_modified] = Time.at(last_modified.to_i)
+        end
+
+        data
+      end
+
+    end
+  end
+end
+

data/lib/amazon_athena/commands/show_tables.rb

@@ -0,0 +1,23 @@
+require_relative '../command'
+
+module AmazonAthena
+  module Commands
+    class ShowTables < AmazonAthena::Command
+
+      def initialize(database_name)
+        @database_name = database_name.strip
+      end
+
+      def statement
+        "SHOW TABLES IN #{@database_name};"
+      end
+
+      def run(connection)
+        connection.query(statement).map {|row| row.tab_name }
+      rescue Exception => e
+        e.getCause()
+      end
+    end
+  end
+end
+

data/lib/amazon_athena/partition.rb

@@ -0,0 +1,21 @@
+module AmazonAthena
+  class Partition
+
+    def initialize(options: {}, location: nil)
+      @options = options
+      @location = location
+    end
+
+    def to_s
+      return nil if @options.empty?
+
+      # TODO: Sanitize and handle non-strings
+      opts = @options.map {|k,v| "#{k} = '#{v}'"}.join(", ")
+
+      sql = "PARTITION (#{opts})"
+      sql += " LOCATION '#{@location}'" if @location
+
+      sql
+    end
+  end
+end

data/lib/amazon_athena/transformer.rb

@@ -0,0 +1,19 @@
+module AmazonAthena
+  class Transformer
+
+    TABLE_NAME_PATTERN = /CREATE EXTERNAL TABLE `(?<name>\S+)`/
+    TABLE_LOCATION_PATTERN = /'(?<location>s3:\/\/\S+)'/
+
+    def self.transform_table(ddl, options = {})
+      if name = options[:name]
+        ddl[TABLE_NAME_PATTERN] = "CREATE EXTERNAL TABLE `#{name}`"
+      end
+
+      if location = options[:location]
+        ddl[TABLE_LOCATION_PATTERN] = "'s3://#{location}'"
+      end
+
+      ddl
+    end
+  end
+end

data/lib/amazon_athena/version.rb

@@ -0,0 +1,3 @@
+module AmazonAthena
+  VERSION = '0.1.0'
+end

data/lib/jdbc_helper/athena.rb

@@ -0,0 +1,38 @@
+require "jdbc-helper"
+require "jdbc_helper/resultset"
+
+module JDBCHelper
+  module Athena
+    extend Connector
+
+    DRIVER_NAME = "com.amazonaws.athena.jdbc.AthenaDriver".freeze
+    JDBCHelper::Constants::Connector::DEFAULT_PARAMETERS[:athena] = {
+      driver: "com.amazonaws.athena.jdbc.AthenaDriver"
+    }
+
+    def self.connect(key: nil, secret: nil, region: "us-east-1", s3_staging_dir: nil, extra_params: {}, &block)
+      connect_impl :athena, {
+        url: "jdbc:awsathena://athena.#{region}.amazonaws.com:443",
+        user: key || ENV["AWS_ACCESS_KEY"],
+        password: secret || ENV['AWS_SECRET_KEY'],
+        s3_staging_dir: s3_uri(s3_staging_dir || ENV["ATHENA_S3_STAGING_DIR"])
+      }, {}, &block
+    end
+
+    def self.configure_driver_path(class_path)
+      paths = ENV["CLASSPATH"].to_s.split(":")
+      paths << class_path
+
+      ENV["CLASSPATH"] = paths.uniq.join(":")
+    end
+
+    def self.s3_uri(path)
+      return path if path.to_s.start_with?("s3://")
+
+      path = "s3://#{path}"
+      path = path + "/" unless path.end_with?("/")
+
+      path
+    end
+  end
+end

data/lib/jdbc_helper/resultset.rb

@@ -0,0 +1,14 @@
+module JDBCHelper
+  class Connection
+    class ResultSet
+      def raw_output(column_index = 1)
+        lines = []
+        begin
+          lines << @rset.getString(column_index)
+        end while @rset.next
+
+        lines.map(&:rstrip).join("\n")
+      end
+    end
+  end
+end

data/log4j.xml

@@ -0,0 +1,18 @@
+<?xml version="1.0" encoding="UTF-8" ?>
+<!DOCTYPE log4j:configuration SYSTEM "log4j.dtd">
+
+<log4j:configuration xmlns:log4j="http://jakarta.apache.org/log4j/">
+  <appender name="console" class="org.apache.log4j.ConsoleAppender">
+    <param name="Target" value="System.out"/>
+    <layout class="org.apache.log4j.PatternLayout">
+      <param name="ConversionPattern" value="%-5p %c{1} - %m%n"/>
+    </layout>
+  </appender>
+
+  <root>
+    <priority value ="OFF" />
+    <appender-ref ref="console" />
+  </root>
+
+
+</log4j:configuration>

data/script/bootstrap

@@ -0,0 +1,5 @@
+#!/bin/sh
+
+set -e
+
+bundle install --quiet "$@"

data/script/package

@@ -0,0 +1,8 @@
+#!/usr/bin/env bash
+# Usage: script/gem
+# Updates the gemspec and builds a new gem in the pkg directory.
+
+mkdir -p pkg
+gem build *.gemspec
+mv *.gem pkg
+

data/script/release

@@ -0,0 +1,16 @@
+#!/usr/bin/env bash
+# Usage: script/release
+# Build the package, tag a commit, push it to origin, and then release the
+# package publicly.
+
+set -e
+
+version="$(script/package | grep Version: | awk '{print $2}')"
+[ -n "$version" ] || exit 1
+
+git commit --allow-empty -a -m "Release $version"
+git tag "v$version"
+git push origin
+git push origin "v$version"
+gem push pkg/*-${version}.gem
+

data/script/test

@@ -0,0 +1,6 @@
+#!/usr/bin/env bash
+# Usage: script/test
+# Runs the library's test suite.
+
+script/bootstrap
+bundle exec rake

data/test/lib/amazon_athena/commands/alter_table_add_partition_test.rb

@@ -0,0 +1,64 @@
+require 'minitest/autorun'
+require './lib/amazon_athena/commands/alter_table_add_partition'
+
+describe AmazonAthena::Commands::AlterTableAddPartition do
+
+  before do
+    @klass = AmazonAthena::Commands::AlterTableAddPartition
+  end
+
+  it "builds a ddl statement" do
+    partitions = []
+
+    partitions << AmazonAthena::Partition.new(
+      location: 's3://mystorage/path/to/INDIA_14_May_2014',
+      options: {
+        dt: '2014-05-14',
+        country: 'IN'
+      }
+    )
+
+    partitions << AmazonAthena::Partition.new(
+      location: 's3://mystorage/path/to/INDIA_15_May_2014',
+      options: {
+        dt: '2014-05-15',
+        country: 'IN'
+      }
+    )
+
+    cmd = @klass.new("mydb.mytable", partitions)
+    expected = <<~SQL
+    ALTER TABLE mydb.mytable ADD
+      PARTITION (dt = '2014-05-14', country = 'IN') LOCATION 's3://mystorage/path/to/INDIA_14_May_2014'
+      PARTITION (dt = '2014-05-15', country = 'IN') LOCATION 's3://mystorage/path/to/INDIA_15_May_2014';
+    SQL
+    assert_equal expected.strip, cmd.statement
+  end
+
+  it "executes a query" do
+    partitions = []
+
+    partitions << AmazonAthena::Partition.new(
+      location: 's3://mystorage/path/to/INDIA_14_May_2014',
+      options: {
+        dt: '2014-05-14',
+        country: 'IN'
+      }
+    )
+
+    cmd = @klass.new("mydb.mytable", partitions)
+    sql = <<~SQL
+    ALTER TABLE mydb.mytable ADD
+      PARTITION (dt = '2014-05-14', country = 'IN') LOCATION 's3://mystorage/path/to/INDIA_14_May_2014';
+    SQL
+
+    results = MiniTest::Mock.new
+    results.expect(:raw_output, nil)
+
+    conn = MiniTest::Mock.new
+    conn.expect(:query, results, [sql.strip])
+
+    cmd = @klass.new("mydb.mytable", partitions)
+    cmd.run(conn)
+  end
+end

data/test/lib/amazon_athena/commands/alter_table_drop_partition_test.rb

@@ -0,0 +1,61 @@
+require 'minitest/autorun'
+require './lib/amazon_athena/commands/alter_table_drop_partition'
+
+describe AmazonAthena::Commands::AlterTableDropPartition do
+
+  before do
+    @klass = AmazonAthena::Commands::AlterTableDropPartition
+  end
+
+  it "builds a ddl statement" do
+    partitions = []
+
+    partitions << AmazonAthena::Partition.new(
+      options: {
+        dt: '2014-05-14',
+        country: 'IN'
+      }
+    )
+
+    partitions << AmazonAthena::Partition.new(
+      options: {
+        dt: '2014-05-15',
+        country: 'IN'
+      }
+    )
+
+    cmd = @klass.new("mydb.mytable", partitions)
+    expected = <<~SQL
+    ALTER TABLE mydb.mytable DROP
+      PARTITION (dt = '2014-05-14', country = 'IN'),
+      PARTITION (dt = '2014-05-15', country = 'IN');
+    SQL
+    assert_equal expected.strip, cmd.statement
+  end
+
+  it "executes a query" do
+    partitions = []
+
+    partitions << AmazonAthena::Partition.new(
+      options: {
+        dt: '2014-05-14',
+        country: 'IN'
+      }
+    )
+
+    cmd = @klass.new("mydb.mytable", partitions)
+    sql = <<~SQL
+    ALTER TABLE mydb.mytable DROP
+      PARTITION (dt = '2014-05-14', country = 'IN');
+    SQL
+
+    results = MiniTest::Mock.new
+    results.expect(:raw_output, nil)
+
+    conn = MiniTest::Mock.new
+    conn.expect(:query, results, [sql.strip])
+
+    cmd = @klass.new("mydb.mytable", partitions)
+    cmd.run(conn)
+  end
+end

data/test/lib/amazon_athena/commands/create_database_test.rb

@@ -0,0 +1,27 @@
+require 'minitest/autorun'
+require './lib/amazon_athena/commands/create_database'
+
+describe AmazonAthena::Commands::CreateDatabase do
+
+  before do
+    @klass = AmazonAthena::Commands::CreateDatabase
+  end
+
+  it "provides a statement" do
+    cmd = @klass.new("mydb")
+
+    assert_equal "CREATE DATABASE IF NOT EXISTS mydb;", cmd.statement
+  end
+
+  it "executes a query" do
+    cmd = @klass.new("mydb")
+
+    results = MiniTest::Mock.new
+    results.expect(:raw_output, nil)
+
+    conn = MiniTest::Mock.new
+    conn.expect(:query, results, ["CREATE DATABASE IF NOT EXISTS mydb;"])
+
+    cmd.run(conn)
+  end
+end

data/test/lib/amazon_athena/commands/describe_table_test.rb

@@ -0,0 +1,23 @@
+require 'minitest/autorun'
+require './lib/amazon_athena/commands/describe_table'
+
+describe AmazonAthena::Commands::DescribeTable do
+
+  before do
+    @cmd = AmazonAthena::Commands::DescribeTable.new("mydb.mytable")
+  end
+
+  it "provides a db statement" do
+    assert_equal "DESCRIBE mydb.mytable;", @cmd.statement
+  end
+
+  it "executes a query" do
+    results = MiniTest::Mock.new
+    results.expect(:raw_output, nil)
+
+    conn = MiniTest::Mock.new
+    conn.expect(:query, results, ["DESCRIBE mydb.mytable;"])
+
+    @cmd.run(conn)
+  end
+end

data/test/lib/amazon_athena/commands/drop_database_test.rb

@@ -0,0 +1,22 @@
+require 'minitest/autorun'
+require './lib/amazon_athena/commands/drop_database'
+
+describe AmazonAthena::Commands::DropDatabase do
+
+  before do
+    @cmd = AmazonAthena::Commands::DropDatabase.new("mydb")
+  end
+
+  it "provides a db statement" do
+    assert_equal "DROP DATABASE IF EXISTS mydb;", @cmd.statement
+  end
+
+  it "executes a query" do
+    results = MiniTest::Mock.new
+
+    conn = MiniTest::Mock.new
+    conn.expect(:query, results, ["DROP DATABASE IF EXISTS mydb;"])
+
+    @cmd.run(conn)
+  end
+end

data/test/lib/amazon_athena/commands/drop_table_test.rb

@@ -0,0 +1,23 @@
+require 'minitest/autorun'
+require './lib/amazon_athena/commands/drop_table'
+
+describe AmazonAthena::Commands::DropTable do
+
+  before do
+    @cmd = AmazonAthena::Commands::DropTable.new("mydb.mytable")
+  end
+
+  it "provides a db statement" do
+    assert_equal "DROP TABLE mydb.mytable;", @cmd.statement
+  end
+
+  it "executes a query" do
+    results = MiniTest::Mock.new
+    results.expect(:raw_output, nil)
+
+    conn = MiniTest::Mock.new
+    conn.expect(:query, results, ["DROP TABLE mydb.mytable;"])
+
+    @cmd.run(conn)
+  end
+end