Merge pull request #226 from treasure-data/send-table-import-requests…

…-to-api-import-directly

Send table:import requests to api-import endpoints directly

lib/td/command/common.rb

@@ -47,7 +47,7 @@ def get_client(opts={})
 
  # optional, if not provided a default is used from the ruby client library
  begin
- if Config.endpoint
+ if !opts[:endpoint] && Config.endpoint
  opts[:endpoint] = Config.endpoint
  end
  rescue ConfigNotFoundError => e
@@ -67,6 +67,17 @@ def get_client(opts={})
  Client.new(apikey, opts)
  end
 
+ DEFAULT_IMPORT_ENDPOINT = "https://" + TreasureData::API::DEFAULT_IMPORT_ENDPOINT
+
+ def get_import_client
+ import_endpoint = begin
+ Config.import_endpoint || DEFAULT_IMPORT_ENDPOINT
+ rescue TreasureData::ConfigNotFoundError
+ DEFAULT_IMPORT_ENDPOINT
+ end
+ get_client(endpoint: import_endpoint)
+ end
+
  def get_ssl_client(opts={})
  opts[:ssl] = true
  get_client(opts)

lib/td/command/runner.rb

@@ -6,11 +6,12 @@ def initialize
  @config_path = nil
  @apikey = nil
  @endpoint = nil
+ @import_endpoint = nil
  @prog_name = nil
  @insecure = false
  end
 
- attr_accessor :apikey, :endpoint, :config_path, :prog_name, :insecure
+ attr_accessor :apikey, :endpoint, :import_endpoint, :config_path, :prog_name, :insecure
 
  def run(argv=ARGV)
  require 'td/version'
@@ -73,6 +74,7 @@ def run(argv=ARGV)
  config_path = @config_path
  apikey = @apikey
  endpoint = @endpoint
+ import_endpoint = @import_endpoint || @endpoint
  insecure = nil
  $verbose = false
  #$debug = false
@@ -94,6 +96,12 @@ def run(argv=ARGV)
  endpoint = e
  }
 
+ op.on('--import-endpoint API_IMPORT_SERVER', "specify the URL for API Import server to use (default: https://api-import.treasuredata.com).") { |e|
+ require 'td/command/common'
+ Command.validate_api_endpoint(e)
+ import_endpoint = e
+ }
+
  op.on('--insecure', "Insecure access: disable SSL (enabled by default)") {|b|
  insecure = true
  }
@@ -140,6 +148,10 @@ def run(argv=ARGV)
  Config.endpoint = endpoint
  Config.cl_endpoint = true
  end
+ if import_endpoint
+ Config.import_endpoint = endpoint
+ Config.cl_import_endpoint = true
+ end
  if insecure
  Config.secure = false
  end

lib/td/command/table.rb

@@ -1,5 +1,6 @@
 require 'td/helpers'
 require 'td/command/job'
+require 'td/client/api'
 
 module TreasureData
 module Command
@@ -453,17 +454,18 @@ def table_import(op)
  import_params[:table_name] = table_name
  import_params[:paths] = paths
 
- client = get_client
+ api_client = get_client
+ import_client = get_import_client
 
  if auto_create
- create_database_and_table_if_not_exist(client, db_name, table_name)
+ create_database_and_table_if_not_exist(api_client, db_name, table_name)
  end
 
- do_table_import(client, import_params)
+ do_table_import(api_client, import_client, import_params)
  end
 
  private
- def do_table_import(client, import_params)
+ def do_table_import(api_client, import_client, import_params)
  case import_params[:format]
  when 'json', 'msgpack'
  #unless time_key
@@ -488,7 +490,7 @@ def do_table_import(client, import_params)
  end
 
  begin
- db = client.database(import_params[:db_name])
+ db = api_client.database(import_params[:db_name])
  rescue ForbiddenError => e
  $stdout.puts "Warning: database and table validation skipped - #{e.message}"
  else
@@ -521,7 +523,7 @@ def do_table_import(client, import_params)
  #require 'thread'
 
  files.zip(import_params[:paths]).each {|file, path|
- import_log_file(file, path, client, import_params[:db_name], import_params[:table_name], parser)
+ import_log_file(file, path, import_client, import_params[:db_name], import_params[:table_name], parser)
  }
 
  $stdout.puts "done."

lib/td/config.rb

@@ -21,6 +21,9 @@ class Config
  @@endpoint = ENV['TREASURE_DATA_API_SERVER'] || ENV['TD_API_SERVER']
  @@endpoint = nil if @@endpoint == ""
  @@cl_endpoint = false # flag to indicate whether an endpoint has been provided through the command-line
+ @@import_endpoint = ENV['TREASURE_DATA_API_IMPORT_SERVER'] || ENV['TD_API_IMPORT_SERVER']
+ @@import_endpoint = nil if @@endpoint == ""
+ @@cl_import_endpoint = false # flag to indicate whether an endpoint has been provided through the command-line option
  @@secure = true
  @@retry_post_requests = false
 
@@ -164,6 +167,22 @@ def self.cl_endpoint=(flag)
  @@cl_endpoint = flag
  end
 
+ def self.import_endpoint
+ @@import_endpoint || Config.read['account.import_endpoint']
+ end
+
+ def self.import_endpoint=(endpoint)
+ @@import_endpoint = endpoint
+ end
+
+ def self.cl_import_endpoint
+ @@cl_import_endpoint
+ end
+
+ def self.cl_import_endpoint=(flag)
+ @@cl_import_endpoint = flag
+ end
+
  def self.workflow_endpoint
  case self.endpoint_domain
  when /\Aapi(-(?:staging|development))?(-[a-z0-9]+)?\.(connect\.)?(eu01\.)?treasuredata\.(com|co\.jp)\z/i
@@ -176,10 +195,9 @@ def self.workflow_endpoint
  # renders the apikey and endpoint options as a string for the helper commands
  def self.cl_options_string
  string = ""
- string += "-k #{@@apikey}" if @@cl_apikey
- string += " " unless string.empty?
- string += "-e #{@@endpoint}" if @@cl_endpoint
- string += " " unless string.empty?
+ string += "-k #{@@apikey} " if @@cl_apikey
+ string += "-e #{@@endpoint} " if @@cl_endpoint
+ string += "--import-endpoint #{@@import_endpoint} " if @@cl_import_endpoint
  string
  end
 

spec/td/command/table_spec.rb

@@ -302,6 +302,7 @@ module TreasureData::Command
  let(:db_name) { 'database' }
  let(:table_name) { 'table' }
  let(:client) { double('client') }
+ let(:import_client) { double('import-client') }
  let(:command) { Class.new { include TreasureData::Command }.new }
 
  describe 'auto create table' do
@@ -336,6 +337,7 @@ module TreasureData::Command
  describe 'time key' do
  before do
  allow(command).to receive(:get_client) { client }
+ allow(command).to receive(:get_import_client) { import_client }
  allow(command).to receive(:do_table_import)
  end
  let(:input_params) {{
@@ -355,7 +357,7 @@ module TreasureData::Command
  end
 
  it "with '#{tk_option}' option" do
- expect(command).to receive(:do_table_import).with(client, input_params)
+ expect(command).to receive(:do_table_import).with(client, import_client, input_params)
  command.table_import(option)
  end
  end
@@ -371,7 +373,7 @@ module TreasureData::Command
  end
 
  it 'without \'-t / --time-key\' option' do
- expect(command).to receive(:do_table_import).with(client, input_params)
+ expect(command).to receive(:do_table_import).with(client, import_client, input_params)
  command.table_import(option)
  end
  end