embulk-output-bigquery 0.2.3 → 0.3.0.pre1

data/test/test_configure.rb

@@ -0,0 +1,254 @@
+require_relative './helper'
+require 'embulk/output/bigquery'
+
+Bigquery = Embulk::Output::Bigquery unless defined?(Bigquery)
+
+module Embulk
+  class Output::Bigquery
+    class TestConfigure < Test::Unit::TestCase
+      class << self
+        def startup
+          FileUtils.mkdir_p('tmp')
+        end
+
+        def shutdown
+          FileUtils.rm_rf('tmp')
+        end
+      end
+
+      def least_config
+        DataSource.new({
+          'project' => 'your_project_name',
+          'dataset' => 'your_dataset_name',
+          'table'   => 'your_table_name',
+          'p12_keyfile' => __FILE__, # fake
+        })
+      end
+
+      def schema
+        Schema.new([
+          Column.new({index: 0, name: 'boolean', type: :boolean}),
+          Column.new({index: 1, name: 'long', type: :long}),
+          Column.new({index: 2, name: 'double', type: :double}),
+          Column.new({index: 3, name: 'string', type: :string}),
+          Column.new({index: 4, name: 'timestamp', type: :timestamp}),
+          Column.new({index: 5, name: 'json', type: :json}),
+        ])
+      end
+
+      def processor_count
+        1
+      end
+
+      def test_configure_default
+        task = Bigquery.configure(least_config, schema, processor_count)
+        assert_equal "append", task['mode']
+        assert_equal "private_key", task['auth_method']
+        assert_equal nil, task['service_account_email']
+        assert_equal __FILE__, task['p12_keyfile']
+        assert_equal nil, task['json_keyfile']
+        assert_equal "your_project_name", task['project']
+        assert_equal "your_dataset_name", task['dataset']
+        assert_equal "your_table_name", task['table']
+        assert_equal nil, task['dataset_old']
+        assert_equal nil, task['table_old']
+        assert_equal nil, task['table_name_old']
+        assert_equal false, task['auto_create_dataset']
+        assert_equal false, task['auto_create_table']
+        assert_equal nil, task['schema_file']
+        assert_equal nil, task['template_table']
+        assert_equal true, task['delete_from_local_when_job_end']
+        assert_equal 3600, task['job_status_max_polling_time']
+        assert_equal 10, task['job_status_polling_interval']
+        assert_equal false, task['is_skip_job_result_check']
+        assert_equal false, task['prevent_duplicate_insert']
+        assert_equal false, task['with_rehearsal']
+        assert_equal 1000, task['rehearsal_counts']
+        assert_equal [], task['column_options']
+        assert_equal "UTC", task['default_timezone']
+        assert_equal "%Y-%m-%d %H:%M:%S.%6N", task['default_timestamp_format']
+        assert_equal nil, task['payload_column']
+        assert_equal nil, task['payload_column_index']
+        assert_equal 300, task['timeout_sec']
+        assert_equal 300, task['open_timeout_sec']
+        assert_equal 5, task['retries']
+        assert_equal "Embulk BigQuery plugin", task['application_name']
+        # assert_equal "/tmp/embulk_output_bigquery_20160228-27184-pubcn0", task['path_prefix']
+        assert_equal ".%d.%03d", task['sequence_format']
+        assert_equal ".csv", task['file_ext']
+        assert_equal false, task['skip_file_generation']
+        assert_equal "NONE", task['compression']
+        assert_equal "CSV", task['source_format']
+        assert_equal 0, task['max_bad_records']
+        assert_equal ",", task['field_delimiter']
+        assert_equal "UTF-8", task['encoding']
+        assert_equal false, task['ignore_unknown_values']
+        assert_equal false, task['allow_quoted_newlines']
+        assert_equal false, task['skip_load']
+      end
+
+      def test_mode
+        config = least_config.merge('mode' => 'foobar')
+        assert_raise { Bigquery.configure(config, schema, processor_count) }
+
+        config = least_config.merge('mode' => 'append')
+        assert_nothing_raised { Bigquery.configure(config, schema, processor_count) }
+
+        config = least_config.merge('mode' => 'replace')
+        assert_nothing_raised { Bigquery.configure(config, schema, processor_count) }
+
+        config = least_config.merge('mode' => 'delete_in_advance')
+        assert_nothing_raised { Bigquery.configure(config, schema, processor_count) }
+
+        config = least_config.merge('mode' => 'replace_backup')
+        assert_raise { Bigquery.configure(config, schema, processor_count) }
+      end
+
+      def test_dataset_table_old
+        task = nil
+        config = least_config.merge('mode' => 'replace_backup', 'table_old' => 'backup')
+        assert_nothing_raised { task = Bigquery.configure(config, schema, processor_count) }
+        assert_equal task['dataset_old'], task['dataset']
+        assert_equal task['table_old'],   'backup'
+
+        config = least_config.merge('mode' => 'replace_backup', 'dataset_old' => 'backup')
+        assert_nothing_raised { task = Bigquery.configure(config, schema, processor_count) }
+        assert_equal task['dataset_old'], 'backup'
+        assert_equal task['table_old'],   task['table']
+      end
+
+      def test_auth_method
+        config = least_config.merge('auth_method' => 'foobar')
+        assert_raise { Bigquery.configure(config, schema, processor_count) }
+
+        config = least_config.merge('auth_method' => 'private_key').tap {|h| h.delete('p12_keyfile') }
+        assert_raise { Bigquery.configure(config, schema, processor_count) }
+        config = least_config.merge('auth_method' => 'private_key', 'p12_keyfile' => 'dummy')
+        assert_nothing_raised { Bigquery.configure(config, schema, processor_count) }
+
+        config = least_config.merge('auth_method' => 'json_key').tap {|h| h.delete('json_keyfile') }
+        assert_raise { Bigquery.configure(config, schema, processor_count) }
+        config = least_config.merge('auth_method' => 'json_key', 'json_keyfile' => "#{EXAMPLE_ROOT}/json_key.json")
+        assert_nothing_raised { Bigquery.configure(config, schema, processor_count) }
+
+        config = least_config.merge('auth_method' => 'compute_engine')
+        assert_nothing_raised { Bigquery.configure(config, schema, processor_count) }
+      end
+
+      def test_json_keyfile
+        json_keyfile = "#{EXAMPLE_ROOT}/json_key.json"
+        config = least_config.merge('auth_method' => 'json_key', 'json_keyfile' => json_keyfile).tap {|h| h.delete('project') }
+        task = Bigquery.configure(config, schema, processor_count)
+        assert_not_equal nil, task['project'] # project is obtained from json_keyfile if available
+
+        config = least_config.merge('auth_method' => 'json_key', 'json_keyfile' => { 'content' => File.read(json_keyfile) }).tap {|h| h.delete('project') }
+        task = Bigquery.configure(config, schema, processor_count)
+        assert_not_equal nil, task['project'] # project is obtained from json_keyfile if available
+
+        config = least_config.merge('auth_method' => 'json_key', 'json_keyfile' => { 'content' => 'not a json' })
+        assert_raise { Bigquery.configure(config, schema, processor_count) }
+      end
+
+      def test_payload_column
+        config = least_config.merge('payload_column' => schema.first.name)
+        task = Bigquery.configure(config, schema, processor_count)
+        assert_equal task['payload_column_index'], 0
+
+        config = least_config.merge('payload_column' => 'not_exist')
+        assert_raise { Bigquery.configure(config, schema, processor_count) }
+      end
+
+      def test_payload_column_index
+        config = least_config.merge('payload_column_index' => 0)
+        assert_nothing_raised { Bigquery.configure(config, schema, processor_count) }
+
+        config = least_config.merge('payload_column_index' => -1)
+        assert_raise { Bigquery.configure(config, schema, processor_count) }
+
+        config = least_config.merge('payload_column_index' => schema.size)
+        assert_raise { Bigquery.configure(config, schema, processor_count) }
+      end
+
+      def test_auto_create_table_with_payload_column
+        config = least_config.merge('auto_create_table' => true, 'payload_column' => 'json')
+        assert_raise { Bigquery.configure(config, schema, processor_count) }
+
+        config = least_config.merge('auto_create_table' => true, 'payload_column' => 'json', 'schema_file' => "#{EXAMPLE_ROOT}/schema.json")
+        assert_nothing_raised { Bigquery.configure(config, schema, processor_count) }
+
+        config = least_config.merge('auto_create_table' => true, 'payload_column' => 'json', 'template_table' => 'foo')
+        assert_nothing_raised { Bigquery.configure(config, schema, processor_count) }
+      end
+
+      def test_auto_create_table_with_payload_column_index
+        config = least_config.merge('auto_create_table' => true, 'payload_column_index' => 0)
+        assert_raise { Bigquery.configure(config, schema, processor_count) }
+
+        config = least_config.merge('auto_create_table' => true, 'payload_column_index' => 0, 'schema_file' => "#{EXAMPLE_ROOT}/schema.json")
+        assert_nothing_raised { Bigquery.configure(config, schema, processor_count) }
+
+        config = least_config.merge('auto_create_table' => true, 'payload_column_index' => 0, 'template_table' => 'foo')
+        assert_nothing_raised { Bigquery.configure(config, schema, processor_count) }
+      end
+
+      def test_schema_file
+        config = least_config.merge('schema_file' => "#{EXAMPLE_ROOT}/schema.json")
+        assert_nothing_raised { Bigquery.configure(config, schema, processor_count) }
+
+        config = least_config.merge('schema_file' => "not_found.json")
+        assert_raise { Bigquery.configure(config, schema, processor_count) }
+
+        File.write("tmp/bad_schema.json", "not_a_json")
+        config = least_config.merge('schema_file' => "tmp/bad_schema.json")
+        assert_raise { Bigquery.configure(config, schema, processor_count) }
+      end
+
+      def test_source_format
+        config = least_config.merge('source_format' => 'csv')
+        assert_nothing_raised { Bigquery.configure(config, schema, processor_count) }
+
+        config = least_config.merge('source_format' => 'jsonl')
+        assert_nothing_raised { Bigquery.configure(config, schema, processor_count) }
+
+        config = least_config.merge('source_format' => 'newline_delimited_json')
+        assert_nothing_raised { Bigquery.configure(config, schema, processor_count) }
+
+        config = least_config.merge('source_format' => 'foobar')
+        assert_raise { Bigquery.configure(config, schema, processor_count) }
+      end
+
+      def test_compression
+        config = least_config.merge('compression' => 'gzip')
+        assert_nothing_raised { Bigquery.configure(config, schema, processor_count) }
+
+        config = least_config.merge('compression' => 'none')
+        assert_nothing_raised { Bigquery.configure(config, schema, processor_count) }
+
+        config = least_config.merge('compression' => 'foobar')
+        assert_raise { Bigquery.configure(config, schema, processor_count) }
+      end
+
+      def test_file_ext
+        config = least_config.merge('source_format' => 'csv', 'compression' => 'gzip')
+        task = Bigquery.configure(config, schema, processor_count)
+        assert_equal '.csv.gz', task['file_ext']
+
+        config = least_config.merge('source_format' => 'NEWLINE_DELIMITED_JSON', 'compression' => 'gzip')
+        task = Bigquery.configure(config, schema, processor_count)
+        assert_equal '.jsonl.gz', task['file_ext']
+
+        config = least_config.merge('source_format' => 'csv', 'compression' => 'none')
+        task = Bigquery.configure(config, schema, processor_count)
+        assert_equal '.csv', task['file_ext']
+
+        config = least_config.merge('source_format' => 'NEWLINE_DELIMITED_JSON', 'compression' => 'none')
+        task = Bigquery.configure(config, schema, processor_count)
+        assert_equal '.jsonl', task['file_ext']
+
+        config = least_config.merge('file_ext' => '.foo')
+        task = Bigquery.configure(config, schema, processor_count)
+        assert_equal '.foo', task['file_ext']
+      end
+    end
+  end
+end

data/test/test_example.rb

@@ -0,0 +1,34 @@
+require_relative './helper'
+
+# 1. Prepare /tmp/your-project-000.json
+# 2. embulk bundle
+# 3. CONNECT=1 bundle exec ruby test/test_example.rb
+
+if ENV['CONNECT']
+  class TestExample < Test::Unit::TestCase
+    def embulk_path
+      if File.exist?("#{ENV['PATH']}/.embulk/bin/embulk")
+        "#{ENV['PATH']}/.embulk/bin/embulk"
+      elsif File.exist?("/usr/local/bin/embulk")
+        "/usr/local/bin/embulk"
+      else
+        "embulk"
+      end
+    end
+
+    files = Dir.glob("#{APP_ROOT}/example/config_*.yml").sort
+    files = files.reject {|file| File.symlink?(file) }
+    # files.shift
+    files.each do |config_path|
+      define_method(:"test_#{File.basename(config_path, ".yml")}") do
+        success = Bundler.with_clean_env do
+          cmd = "#{embulk_path} run -X page_size=1 -b . -l trace #{config_path}"
+          puts "=" * 64
+          puts cmd
+          system(cmd)
+        end
+        assert_true success
+      end
+    end
+  end
+end

data/test/test_file_writer.rb

@@ -0,0 +1,129 @@
+require_relative './helper'
+require 'embulk/output/bigquery/file_writer'
+require 'fileutils'
+require 'zlib'
+
+module Embulk
+  class Output::Bigquery
+    class TestFileWriter < Test::Unit::TestCase
+      class << self
+        def startup
+          FileUtils.mkdir_p('tmp')
+        end
+
+        def shutdown
+          FileUtils.rm_rf('tmp')
+        end
+      end
+
+      def default_task
+        {
+          'compression' => 'GZIP',
+          'payload_column' => nil,
+          'source_format' => 'CSV',
+          'path_prefix' => 'tmp/path_prefix',
+          'sequence_format' => '.%d.%03d',
+          'file_ext' => nil,
+        }
+      end
+
+      def schema
+        Schema.new([
+          Column.new({index: 0, name: 'boolean', type: :boolean}),
+          Column.new({index: 1, name: 'long', type: :long}),
+          Column.new({index: 2, name: 'double', type: :double}),
+          Column.new({index: 3, name: 'string', type: :string}),
+          Column.new({index: 4, name: 'timestamp', type: :timestamp}),
+          Column.new({index: 5, name: 'json', type: :json}),
+        ])
+      end
+
+      def converters
+        @converters ||= ValueConverterFactory.create_converters(default_task, schema)
+      end
+
+      sub_test_case "path" do
+        def test_path
+          task = default_task.merge('path_prefix' => 'tmp/foo', 'sequence_format' => '', 'file_ext' => '.1')
+          file_writer = FileWriter.new(task, schema, 0, converters)
+          assert_equal 'tmp/foo.1', file_writer.instance_variable_get(:@path)
+        end
+      end
+
+      sub_test_case "formatter" do
+        def record
+          [true, 1, 1.1, 'foo', Time.parse("2016-02-26 00:00:00 +09:00"), {"foo"=>"foo"}]
+        end
+
+        def test_payload_column_index
+          task = default_task.merge('payload_column_index' => 0)
+          file_writer = FileWriter.new(task, schema, 0, converters)
+          formatter_proc = file_writer.instance_variable_get(:@formatter_proc)
+          assert_equal :to_payload, formatter_proc.name
+
+          assert_equal %Q[true\n], formatter_proc.call(record)
+        end
+
+        def test_csv
+          task = default_task.merge('source_format' => 'CSV')
+          file_writer = FileWriter.new(task, schema, 0, converters)
+          formatter_proc = file_writer.instance_variable_get(:@formatter_proc)
+          assert_equal :to_csv, formatter_proc.name
+
+          expected = %Q[true,1,1.1,foo,1456412400.0,"{""foo"":""foo""}"\n]
+          assert_equal expected, formatter_proc.call(record)
+        end
+
+        def test_jsonl
+          task = default_task.merge('source_format' => 'NEWLINE_DELIMITED_JSON')
+          file_writer = FileWriter.new(task, schema, 0, converters)
+          formatter_proc = file_writer.instance_variable_get(:@formatter_proc)
+          assert_equal :to_jsonl, formatter_proc.name
+
+          expected = %Q[{"boolean":true,"long":1,"double":1.1,"string":"foo","timestamp":1456412400.0,"json":"{\\"foo\\":\\"foo\\"}"}\n]
+          assert_equal expected, formatter_proc.call(record)
+        end
+      end
+
+      sub_test_case "compression" do
+        def record
+          [true, 1, 1.1, 'foo', Time.parse("2016-02-26 00:00:00 +09:00"), {"foo"=>"foo"}]
+        end
+
+        def page
+          [record]
+        end
+
+        def test_gzip
+          task = default_task.merge('compression' => 'GZIP')
+          file_writer = FileWriter.new(task, schema, 0, converters)
+          io = file_writer.instance_variable_get(:@io)
+          assert_equal Zlib::GzipWriter, io.class
+
+          begin
+            file_writer.add(page)
+          ensure
+            file_writer.commit
+          end
+          assert_true File.exist?(file_writer.path)
+          assert_nothing_raised { Zlib::GzipReader.open(file_writer.path) {|gz| } }
+        end
+
+        def test_uncompressed
+          task = default_task.merge('compression' => 'NONE')
+          file_writer = FileWriter.new(task, schema, 0, converters)
+          io = file_writer.instance_variable_get(:@io)
+          assert_equal File, io.class
+
+          begin
+            file_writer.add(page)
+          ensure
+            file_writer.commit
+          end
+          assert_true File.exist?(file_writer.path)
+          assert_raise { Zlib::GzipReader.open(file_writer.path) {|gz| } }
+        end
+      end
+    end
+  end
+end

data/test/test_helper.rb

@@ -0,0 +1,103 @@
+require_relative './helper'
+require 'embulk/output/bigquery/helper'
+
+module Embulk
+  class Output::Bigquery
+    class TestHelper < Test::Unit::TestCase
+      class << self
+        def startup
+          FileUtils.mkdir_p('tmp')
+        end
+
+        def shutdown
+          FileUtils.rm_rf('tmp')
+        end
+      end
+
+      def bq_type_from_embulk_type
+        assert_equal 'BOOLEAN',   Helper.bq_type_from_embulk_type(:boolean)
+        assert_equal 'STRING',    Helper.bq_type_from_embulk_type(:string)
+        assert_equal 'FLOAT',     Helper.bq_type_from_embulk_type(:double)
+        assert_equal 'STRING',    Helper.bq_type_from_embulk_type(:string)
+        assert_equal 'TIMESTAMP', Helper.bq_type_from_embulk_type(:timestamp)
+        assert_equal 'STRING',    Helper.bq_type_from_embulk_type(:json)
+      end
+
+      sub_test_case "fields_from_embulk_schema" do
+        def test_fields_from_embulk_schema_without_column_options
+          schema = Schema.new([
+            Column.new({index: 0, name: 'boolean', type: :boolean}),
+            Column.new({index: 1, name: 'long', type: :long}),
+            Column.new({index: 2, name: 'double', type: :double}),
+            Column.new({index: 3, name: 'string', type: :string}),
+            Column.new({index: 4, name: 'timestamp', type: :timestamp}),
+            Column.new({index: 5, name: 'json', type: :json}),
+          ])
+          expected = [
+            {name: 'boolean',   type: 'BOOLEAN'},
+            {name: 'long',      type: 'INTEGER'},
+            {name: 'double',    type: 'FLOAT'},
+            {name: 'string',    type: 'STRING'},
+            {name: 'timestamp', type: 'TIMESTAMP'},
+            {name: 'json',      type: 'STRING'},
+          ]
+          fields = Helper.fields_from_embulk_schema({}, schema)
+          assert_equal expected, fields
+        end
+
+        def test_fields_from_embulk_schema_with_column_options
+          schema = Schema.new([
+            Column.new({index: 0, name: 'boolean', type: :boolean}),
+            Column.new({index: 1, name: 'long', type: :long}),
+            Column.new({index: 2, name: 'double', type: :double}),
+            Column.new({index: 3, name: 'string', type: :string}),
+            Column.new({index: 4, name: 'timestamp', type: :timestamp}),
+            Column.new({index: 5, name: 'json', type: :json}),
+          ])
+          task = {
+            'column_options' => [
+              {'name' => 'boolean',   'type' => 'STRING', 'mode' => 'REQUIRED'},
+              {'name' => 'long',      'type' => 'STRING'},
+              {'name' => 'double',    'type' => 'STRING'},
+              {'name' => 'string',    'type' => 'INTEGER'},
+              {'name' => 'timestamp', 'type' => 'INTEGER'},
+              {'name' => 'json',      'type' => 'RECORD', 'fields' => [
+                { 'name' => 'key1',   'type' => 'STRING' },
+              ]},
+            ],
+          }
+          expected = [
+            {name: 'boolean',   type: 'STRING', mode: 'REQUIRED'},
+            {name: 'long',      type: 'STRING'},
+            {name: 'double',    type: 'STRING'},
+            {name: 'string',    type: 'INTEGER'},
+            {name: 'timestamp', type: 'INTEGER'},
+            {name: 'json',      type: 'RECORD', fields: [
+              {name: 'key1',    type: 'STRING'},
+            ]},
+          ]
+          fields = Helper.fields_from_embulk_schema(task, schema)
+          assert_equal expected, fields
+        end
+      end
+
+      def test_create_job_id
+        task = {
+          'dataset' => 'your_dataset_name',
+          'source_format' => 'CSV',
+          'max_bad_records' => nil,
+          'field_delimiter' => ',',
+          'encoding' => 'UTF-8',
+          'ignore_unknown_values' => nil,
+          'allow_quoted_newlines' => nil,
+        }
+        fields = {
+          name: 'a', type: 'STRING',
+        }
+        File.write("tmp/your_file_name", "foobarbaz")
+        job_id = Helper.create_job_id(task, 'tmp/your_file_name', 'your_table_name', fields)
+        assert job_id.is_a?(String)
+      end
+    end
+  end
+end