require 'gli' require 'amazon_athena' require 'fileutils' require 'table_print' module AmazonAthena module CLI def self.athena AmazonAthena::Client.new \ key: self.access_key, secret: self.access_secret, s3_staging_dir: self.staging_folder end extend GLI::App program_desc 'CLI for Amazon Athena' version AmazonAthena::VERSION switch [:p, :preview], :default_value => false, :desc => "Output the SQL statement instead of running", :negatable => false flag ["key"], :default_value => ENV["AWS_ACCESS_KEY"], :mask => true, :desc => "AWS Access Key" flag ["secret"], :default_value => ENV["AWS_SECRET_KEY"], :mask => true, :desc => "AWS Secret Key" flag ["staging-dir"], :default_value => ENV["ATHENA_S3_STAGING_DIR"], :desc => "S3 bucket for staging results" desc 'Check for required AWS setup' command :doctor do |c| c.action do |global_options,options,args| puts "Good to go!" end end desc 'Manage Athena schemas' command :schema do |c| c.desc "Dump schemas to a local folder" c.command :dump do |add| add.flag :path, :required => true, :desc => "Local path for saving schemas" add.action do |global_options,options,args| path = options[:path] unless File.exists?(path) exit_now! "Path #{path} not found" end databases = athena.run(AmazonAthena::Commands::ShowDatabases.new) preview = !!global_options[:p] databases.each do |db| folder = File.join(path, db) prefix = File.exists?(folder) ? "Exists" : "Creating" render "#{prefix} #{folder}" unless preview FileUtils.mkdir_p(folder) end tables = athena.run(AmazonAthena::Commands::ShowTables.new(db)) tables.each do |table| database_table = [db, table].join(".") file = File.join(folder, table) + ".sql" render "Writing #{file}" if !preview sql = athena.run(AmazonAthena::Commands::ShowCreateTable.new(database_table)) File.open(file, 'w') { |f| f.write(sql) } end end end # puts global_options, options, args end end end desc 'Manage Athena databases' command :database do |c| c.desc "List databases" c.command :list do |add| add.action do |global_options,options,args| cmd = AmazonAthena::Commands::ShowDatabases.new render athena.run(cmd, global_options[:p]) end end c.desc "Create a new database" c.arg_name "database" c.command :create do |add| add.action do |global_options,options,args| cmd = AmazonAthena::Commands::CreateDatabase.new(args.first) render athena.run(cmd, global_options[:p]) end end c.desc "Drop an existing database" c.arg_name "database" c.command :drop do |add| add.action do |global_options,options,args| cmd = AmazonAthena::Commands::DropDatabase.new(args.first) render athena.run(cmd, global_options[:p]) end end c.default_command :list end desc 'Manage tables in Athena databases' command :table do |c| c.desc "Create table" c.arg_name "[file]" c.command :create do |add| add.flag [:l, :location], :desc => "S3 location - s3://bucket/path/" add.flag [:n, :name], :desc => "Fully qualifed name as: database.table" add.action do |global_options,options,args| ddl = if filename = args.first exit_now! "File not found" unless File.exists?(filename) File.read(filename) elsif !STDIN.tty? $stdin.read end exit_now! "Must supply path to file or pass via STDIN" if ddl.to_s.empty? opts = { :name => options[:name], :location => options[:location] } ddl = AmazonAthena::Transformer.transform_table(ddl, opts) cmd = AmazonAthena::Commands::CreateTable.new(ddl) render athena.run(cmd, global_options[:p]) end end c.desc "Describe a table" c.arg_name "database.table" c.command :describe do |add| add.action do |global_options,options,args| cmd = AmazonAthena::Commands::DescribeTable.new(args.first) render athena.run(cmd, global_options[:p]) end end c.desc "List tables in a database" c.arg_name "database" c.command :list do |add| add.action do |global_options,options,args| cmd = AmazonAthena::Commands::ShowTables.new(args.first) render athena.run(cmd, global_options[:p]) end end c.desc "Show table properties" c.arg_name "database.table" c.command :properties do |add| add.action do |global_options,options,args| cmd = AmazonAthena::Commands::ShowTableProperties.new(args.first) render athena.run(cmd, global_options[:p]) end end c.desc "Show table create DDL" c.arg_name "database.table" c.command :show do |add| add.action do |global_options,options,args| cmd = AmazonAthena::Commands::ShowCreateTable.new(args.first) render athena.run(cmd, global_options[:p]) end end c.desc "Refresh table partitions" c.arg_name "database.table" c.command :repair do |add| add.action do |global_options,options,args| commands = [ AmazonAthena::Commands::RepairTable.new(args.first), AmazonAthena::Commands::ShowPartitions.new(args.first), ] commands.each do |cmd| render athena.run(cmd, global_options[:p]) end end end c.desc "Drop table" c.arg_name "database.table" c.command :drop do |add| add.action do |global_options,options,args| cmd = AmazonAthena::Commands::DropTable.new(args.first) render athena.run(cmd, global_options[:p]) end end end desc 'Manage columns in Athena tables' command :column do |c| c.desc "List columns in a table" c.arg_name "database.table" c.command :list do |add| add.action do |global_options,options,args| cmd = AmazonAthena::Commands::ShowColumns.new(args.first) render athena.run(cmd, global_options[:p]) end end c.default_command :list end desc 'Manage partitions in Athena tables' command :partition do |c| c.desc "List partitions in a table" c.arg_name "database.table" c.command :list do |add| add.action do |global_options,options,args| cmd = AmazonAthena::Commands::ShowPartitions.new(args.first) render athena.run(cmd, global_options[:p]) end end c.desc "TODO: Add partition(s) to a table" c.arg_name "database.table" c.arg_name "key=value,key=value:key=value,key=value" c.command :add do |add| add.action do |global_options,options,args| render "NOT IMPLEMENTED" end end c.desc "TODO: Drop partition(s) from a table" c.arg_name "database.table" c.arg_name "key=value,key=value:key=value,key=value" c.command :drop do |add| add.action do |global_options,options,args| render "NOT IMPLEMENTED" end end c.default_command :list end desc 'Run a query from a file or STDIN' arg_name '[file]' command :query do |c| c.action do |global_options,options,args| input = args.first sql = case input when "", NilClass $stdin.read if !STDIN.tty? when String if File.exists?(input) File.read(input) else input end else exit_now! "SQL string or file path required" end if global_options[:p] render sql else begin render athena.query(sql).map(&:to_h) rescue Exception => e render e.message end end end end pre do |global,command,options,args| @access_key = global[:key] @access_secret = global[:secret] @staging_folder = global[:"staging-dir"] self.check_aws_settings self.check_class_path true end post do |global,command,options,args| # Post logic here # Use skips_post before a command to skip this # block on that command only end on_error do |exception| # Error logic here # return false to skip default error handling true end def self.access_key @access_key end def self.access_secret @access_secret end def self.staging_folder @staging_folder end def self.class_path ENV['CLASSPATH'] end def self.check_aws_settings if self.access_key.nil? || self.access_secret.nil? msg = <<~MSG athena-cli needs your AWS credentials. You can pass them via the --key and --secret flags or set AWS_ACCESS_KEY and AWS_SECRET_KEY environment variables. MSG exit_now! msg end if self.staging_folder.nil? msg = <<~MSG athena-cli requires an S3 location to use as a scratch folder. Please provide one via the --staging-dir flag or set the ATHENA_S3_STAGING_DIR environment variable. MSG exit_now! msg, 1 end end def self.check_class_path unless self.class_path =~ /AthenaJDBC/ jar_path = File.expand_path(File.join([File.dirname(__FILE__),'..', 'jdbc','AthenaJDBC41-1.0.0.jar'])) msg = <<~MSG JRuby requires JDBC driver to be in your Java class path. For download instructions, see: http://docs.aws.amazon.com/athena/latest/ug/connect-with-jdbc.html After downloading, add /path/to/driver.jar to your CLASSPATH environment variable. Example: export CLASSPATH="$CLASSPATH:~/src/AthenaJDB41-1.0.0.jar" MSG exit_now! msg, 1 end end def self.details(data) longest_key = data.keys.max_by(&:length) data.each do |key, value| printf "%##{longest_key.length}s %s\n", key, value end end def self.render(output) case output when Hash details(output) when String puts output when Array case output.first when Hash tp output else puts output end else return end end end end