athena-cli 0.1.0

data/lib/amazon_athena/client.rb

@@ -0,0 +1,114 @@
+require "forwardable"
+require "jdbc_helper/athena"
+require_relative "commands"
+
+module AmazonAthena
+  class Client
+
+    extend Forwardable
+
+    def_delegators :connection, :query, :execute
+
+    def initialize(key: nil, secret: nil, region: "us-east-1", s3_staging_dir: nil)
+      @key = key || ENV["AWS_ACCESS_KEY"]
+      @secret = secret || ENV['AWS_SECRET_KEY']
+      @region = region
+      @s3_staging_dir = s3_staging_dir || ENV["ATHENA_S3_STAGING_DIR"]
+    end
+
+    def databases
+      cmd = AmazonAthena::Commands::ShowDatabases.new
+
+      run(cmd)
+    end
+
+    def database_drop(database)
+      cmd = AmazonAthena::Commands::DropDatabase.new(database)
+
+      run(cmd)
+    end
+
+    def database_create(name:, location: nil, comment: nil, properties: {})
+      cmd = AmazonAthena::Commands::CreateDatabase.new(
+        name: name,
+        location: location,
+        comment: comment,
+        properties: properties
+      )
+
+      run(cmd)
+    end
+
+    def tables(database)
+      cmd = AmazonAthena::Commands::ShowTables.new(database)
+
+      run(cmd)
+    end
+
+    def table_drop(database_table)
+      cmd = AmazonAthena::Commands::DropTable.new(database)
+
+      run(cmd)
+    end
+
+    def table_columns(database_table)
+      cmd = AmazonAthena::Commands::ShowColumns.new(database_table)
+
+      run(cmd)
+    end
+
+    def table_show_create(database_table)
+      cmd = AmazonAthena::Commands::ShowCreateTable.new(database_table)
+
+      run(cmd)
+    end
+
+    def table_describe(database_table)
+      cmd = AmazonAthena::Commands::DescribeTable.new(database_table)
+
+      run(cmd)
+    end
+
+    def table_repair(database_table)
+      cmd = AmazonAthena::Commands::RepairTable.new(database_table)
+      run(cmd)
+
+      partitions(database_table)
+    end
+
+    def table_properties(database_table)
+      cmd = AmazonAthena::Commands::ShowTableProperties.new(database_table)
+
+      run(cmd)
+    end
+
+    def partitions(table)
+      cmd = AmazonAthena::Commands::ShowPartitions.new(database_table)
+
+      run(cmd)
+    end
+
+    def partitions_drop(database_table, partitions_expression)
+      cmd = AmazonAthena::Commands::DropPartition.new(database_table, partitions)
+
+      run(cmd)
+    end
+
+    def connection
+      return @connection if defined?(@connection) && !@connection.closed?
+
+      @connection = JDBCHelper::Athena.connect(
+        key: @key,
+        secret: @secret,
+        region: @region,
+        s3_staging_dir: @s3_staging_dir
+      )
+    end
+
+    def run(cmd, preview = false)
+      return cmd.preview if preview
+
+      cmd.run(connection)
+    end
+  end
+end

data/lib/amazon_athena/command.rb

@@ -0,0 +1,16 @@
+module AmazonAthena
+  class Command
+
+    def statement
+      raise "Not implemented"
+    end
+
+    def preview
+      statement
+    end
+
+    def run(connection)
+      raise "Not implemented"
+    end
+  end
+end

data/lib/amazon_athena/commands.rb

@@ -0,0 +1,20 @@
+require_relative "command"
+require_relative "commands/alter_table_add_partition"
+require_relative "commands/alter_table_drop_partition"
+require_relative "commands/create_database"
+require_relative "commands/create_table"
+require_relative "commands/describe_table"
+require_relative "commands/drop_database"
+require_relative "commands/drop_table"
+require_relative "commands/repair_table"
+require_relative "commands/show_columns"
+require_relative "commands/show_create_table"
+require_relative "commands/show_databases"
+require_relative "commands/show_partitions"
+require_relative "commands/show_table_properties"
+require_relative "commands/show_tables"
+
+module AmazonAthena
+  module Commands
+  end
+end

data/lib/amazon_athena/commands/alter_table_add_partition.rb

@@ -0,0 +1,29 @@
+require_relative '../command'
+require_relative '../partition'
+
+module AmazonAthena
+  module Commands
+    class AlterTableAddPartition < AmazonAthena::Command
+
+      def initialize(database_table, partitions)
+        @database_table = database_table
+        @partitions = partitions
+      end
+
+      def partition_clauses
+        @partitions.map {|p| "  #{p}"}.join("\n")
+      end
+
+      def statement
+        "ALTER TABLE #{@database_table} ADD\n#{partition_clauses};"
+      end
+
+      def run(connection)
+        # TODO: Map fields directly
+        connection.query(statement).raw_output
+      end
+
+    end
+  end
+end
+

data/lib/amazon_athena/commands/alter_table_drop_partition.rb

@@ -0,0 +1,29 @@
+require_relative '../command'
+require_relative '../partition'
+
+module AmazonAthena
+  module Commands
+    class AlterTableDropPartition < AmazonAthena::Command
+
+      def initialize(database_table, partitions)
+        @database_table = database_table
+        @partitions = partitions
+      end
+
+      def partition_clauses
+        @partitions.map {|p| "  #{p}"}.join(",\n")
+      end
+
+      def statement
+        "ALTER TABLE #{@database_table} DROP\n#{partition_clauses};"
+      end
+
+      def run(connection)
+        # TODO: Map fields directly
+        connection.query(statement).raw_output
+      end
+
+    end
+  end
+end
+

data/lib/amazon_athena/commands/create_database.rb

@@ -0,0 +1,23 @@
+require_relative '../command'
+
+module AmazonAthena
+  module Commands
+    class CreateDatabase < AmazonAthena::Command
+
+      def initialize(name)
+        @name = name
+      end
+
+      def statement
+        "CREATE DATABASE IF NOT EXISTS #{@name};"
+      end
+
+      def run(connection)
+        connection.query(statement)
+
+        return
+      end
+    end
+  end
+end
+

data/lib/amazon_athena/commands/create_table.rb

@@ -0,0 +1,23 @@
+require_relative '../command'
+
+module AmazonAthena
+  module Commands
+    class CreateTable < AmazonAthena::Command
+
+      def initialize(ddl)
+        @ddl = ddl
+      end
+
+      def statement
+        @ddl
+      end
+
+      def run(connection)
+        connection.query(statement)
+
+        return
+      end
+    end
+  end
+end
+

data/lib/amazon_athena/commands/describe_table.rb

@@ -0,0 +1,21 @@
+require_relative '../command'
+
+module AmazonAthena
+  module Commands
+    class DescribeTable < AmazonAthena::Command
+
+      def initialize(database_table)
+        @database_table = database_table
+      end
+
+      def statement
+        "DESCRIBE #{@database_table};"
+      end
+
+      def run(connection)
+        connection.query(statement).raw_output
+      end
+    end
+  end
+end
+

data/lib/amazon_athena/commands/drop_database.rb

@@ -0,0 +1,23 @@
+require_relative '../command'
+
+module AmazonAthena
+  module Commands
+    class DropDatabase < AmazonAthena::Command
+
+      def initialize(database_name)
+        @database_name = database_name
+      end
+
+      def statement
+        "DROP DATABASE IF EXISTS #{@database_name};"
+      end
+
+      def run(connection)
+        connection.query(statement)
+
+        return
+      end
+    end
+  end
+end
+

data/lib/amazon_athena/commands/drop_table.rb

@@ -0,0 +1,23 @@
+require_relative '../command'
+
+module AmazonAthena
+  module Commands
+    class DropTable < AmazonAthena::Command
+
+      def initialize(database_table)
+        @database_table = database_table
+      end
+
+      def statement
+        "DROP TABLE #{@database_table};"
+      end
+
+      def run(connection)
+        connection.query(statement)
+
+        return
+      end
+    end
+  end
+end
+

data/lib/amazon_athena/commands/repair_table.rb

@@ -0,0 +1,23 @@
+require_relative '../command'
+
+module AmazonAthena
+  module Commands
+    class RepairTable < AmazonAthena::Command
+
+      def initialize(database_table)
+        @database_table = database_table
+      end
+
+      def statement
+        "MSCK REPAIR TABLE #{@database_table};"
+      end
+
+      def run(connection)
+        connection.query(statement)
+
+        return
+      end
+    end
+  end
+end
+

data/lib/amazon_athena/commands/show_columns.rb

@@ -0,0 +1,22 @@
+require_relative '../command'
+
+module AmazonAthena
+  module Commands
+    class ShowColumns < AmazonAthena::Command
+
+      def initialize(database_table)
+        @database_table = database_table
+      end
+
+      def statement
+        "SHOW COLUMNS IN #{@database_table};"
+      end
+
+      def run(connection)
+        connection.query(statement).map {|row| row.field.strip }
+      end
+
+    end
+  end
+end
+

data/lib/amazon_athena/commands/show_create_table.rb

@@ -0,0 +1,22 @@
+require_relative '../command'
+
+module AmazonAthena
+  module Commands
+    class ShowCreateTable < AmazonAthena::Command
+
+      def initialize(database_table)
+        @database_table = database_table
+      end
+
+      def statement
+        "SHOW CREATE TABLE #{@database_table};"
+      end
+
+      def run(connection)
+        connection.query(statement).raw_output
+      end
+
+    end
+  end
+end
+

data/lib/amazon_athena/commands/show_databases.rb

@@ -0,0 +1,17 @@
+require_relative '../command'
+
+module AmazonAthena
+  module Commands
+    class ShowDatabases < AmazonAthena::Command
+
+      def statement
+        "SHOW DATABASES;"
+      end
+
+      def run(connection)
+        connection.query(statement).map {|row| row.database_name }
+      end
+    end
+  end
+end
+

data/lib/amazon_athena/commands/show_partitions.rb

@@ -0,0 +1,28 @@
+require_relative '../command'
+
+module AmazonAthena
+  module Commands
+    class ShowPartitions < AmazonAthena::Command
+
+      def initialize(database_table)
+        @database_table = database_table
+      end
+
+      def statement
+        "SHOW PARTITIONS #{@database_table};"
+      end
+
+      def run(connection)
+        # TODO: Map fields directly
+        connection.query(statement).raw_output
+      rescue Exception => e
+        case e.message
+        when /not a partitioned table/
+          "Error: #{@database_table} is not a partitioned table."
+        end
+      end
+
+    end
+  end
+end
+