embulk-output-bigquery 0.4.14 → 0.6.3

data/lib/embulk/output/bigquery/auth.rb

@@ -0,0 +1,35 @@
+require 'googleauth'
+
+module Embulk
+  module Output
+    class Bigquery < OutputPlugin
+      class Auth
+
+        attr_reader :auth_method, :json_key, :scope
+
+        def initialize(task, scope)
+          @auth_method = task['auth_method']
+          @json_key = task['json_keyfile']
+          @scope = scope
+        end
+
+        def authenticate
+          case auth_method
+          when 'authorized_user'
+            key = StringIO.new(json_key)
+            return Google::Auth::UserRefreshCredentials.make_creds(json_key_io: key, scope: scope)
+          when 'compute_engine'
+            return Google::Auth::GCECredentials.new
+          when 'service_account', 'json_key' # json_key is for backward compatibility
+            key = StringIO.new(json_key)
+            return Google::Auth::ServiceAccountCredentials.make_creds(json_key_io: key, scope: scope)
+          when 'application_default'
+            return Google::Auth.get_application_default([scope])
+          else
+            raise ConfigError.new("Unknown auth method: #{auth_method}")
+          end
+        end
+      end
+    end
+  end
+end

data/lib/embulk/output/bigquery/bigquery_client.rb

@@ -79,11 +79,7 @@ module Embulk
             begin
               # As https://cloud.google.com/bigquery/docs/managing_jobs_datasets_projects#managingjobs says,
               # we should generate job_id in client code, otherwise, retrying would cause duplication
-              if @task['prevent_duplicate_insert'] and (@task['mode'] == 'append' or @task['mode'] == 'append_direct')
-                job_id = Helper.create_load_job_id(@task, path, fields)
-              else
-                job_id = "embulk_load_job_#{SecureRandom.uuid}"
-              end
+              job_id = "embulk_load_job_#{SecureRandom.uuid}"
               Embulk.logger.info { "embulk-output-bigquery: Load job starting... job_id:[#{job_id}] #{object_uris} => #{@project}:#{@dataset}.#{table} in #{@location_for_log}" }
 
               body = {
@@ -174,11 +170,7 @@ module Embulk
               if File.exist?(path)
                 # As https://cloud.google.com/bigquery/docs/managing_jobs_datasets_projects#managingjobs says,
                 # we should generate job_id in client code, otherwise, retrying would cause duplication
-                if @task['prevent_duplicate_insert'] and (@task['mode'] == 'append' or @task['mode'] == 'append_direct')
-                  job_id = Helper.create_load_job_id(@task, path, fields)
-                else
-                  job_id = "embulk_load_job_#{SecureRandom.uuid}"
-                end
+                job_id = "embulk_load_job_#{SecureRandom.uuid}"
                 Embulk.logger.info { "embulk-output-bigquery: Load job starting... job_id:[#{job_id}] #{path} => #{@project}:#{@dataset}.#{table} in #{@location_for_log}" }
               else
                 Embulk.logger.info { "embulk-output-bigquery: Load job starting... #{path} does not exist, skipped" }
@@ -437,7 +429,6 @@ module Embulk
                 type: options['time_partitioning']['type'],
                 expiration_ms: options['time_partitioning']['expiration_ms'],
                 field: options['time_partitioning']['field'],
-                require_partition_filter: options['time_partitioning']['require_partition_filter'],
               }
             end
 

data/lib/embulk/output/bigquery/google_client.rb

@@ -1,4 +1,4 @@
-require 'google/api_client/auth/key_utils'
+require_relative 'auth'
 
 module Embulk
   module Output
@@ -14,6 +14,7 @@ module Embulk
         def initialize(task, scope, client_class)
           @task = task
           @scope = scope
+          @auth = Auth.new(task, scope)
           @client_class = client_class
         end
 
@@ -37,39 +38,7 @@ module Embulk
           Embulk.logger.debug { "embulk-output-bigquery: client_options: #{client.client_options.to_h}" }
           Embulk.logger.debug { "embulk-output-bigquery: request_options: #{client.request_options.to_h}" }
 
-          case @task['auth_method']
-          when 'private_key'
-            private_key_passphrase = 'notasecret'
-            key = Google::APIClient::KeyUtils.load_from_pkcs12(@task['p12_keyfile'], private_key_passphrase)
-            auth = Signet::OAuth2::Client.new(
-              token_credential_uri: "https://accounts.google.com/o/oauth2/token",
-              audience: "https://accounts.google.com/o/oauth2/token",
-              scope: @scope,
-              issuer: @task['service_account_email'],
-              signing_key: key)
-
-          when 'compute_engine'
-            auth = Google::Auth::GCECredentials.new
-
-          when 'json_key'
-            json_key = @task['json_keyfile']
-            if File.exist?(json_key)
-              auth = File.open(json_key) do |f|
-                Google::Auth::ServiceAccountCredentials.make_creds(json_key_io: f, scope: @scope)
-              end
-            else
-              key = StringIO.new(json_key)
-              auth = Google::Auth::ServiceAccountCredentials.make_creds(json_key_io: key, scope: @scope)
-            end
-
-          when 'application_default'
-            auth = Google::Auth.get_application_default([@scope])
-
-          else
-            raise ConfigError, "Unknown auth method: #{@task['auth_method']}"
-          end
-
-          client.authorization = auth
+          client.authorization = @auth.authenticate
 
           @cached_client_expiration = Time.now + 1800
           @cached_client = client

data/lib/embulk/output/bigquery/value_converter_factory.rb

@@ -203,6 +203,13 @@ module Embulk
                 val # Users must care of BQ timestamp format
               }
             end
+          when 'DATE'
+            Proc.new {|val|
+              next nil if val.nil?
+              with_typecast_error(val) do |val|
+                TimeWithZone.set_zone_offset(Time.parse(val), zone_offset).strftime("%Y-%m-%d")
+              end
+            }
           when 'RECORD'
             Proc.new {|val|
               next nil if val.nil?
@@ -240,6 +247,11 @@ module Embulk
               next nil if val.nil?
               val.strftime("%Y-%m-%d %H:%M:%S.%6N %:z")
             }
+          when 'DATE'
+            Proc.new {|val|
+              next nil if val.nil?
+              val.localtime(zone_offset).strftime("%Y-%m-%d")
+            }
           else
             raise NotSupportedType, "cannot take column type #{type} for timestamp column"
           end

data/test/test_bigquery_client.rb

@@ -32,7 +32,7 @@ else
             'dataset'          => 'your_dataset_name',
             'table'            => 'your_table_name',
             'auth_method'      => 'json_key',
-            'json_keyfile'     => JSON_KEYFILE,
+            'json_keyfile'     => File.read(JSON_KEYFILE),
             'retries'          => 3,
             'timeout_sec'      => 300,
             'open_timeout_sec' => 300,
@@ -61,10 +61,6 @@ else
           def test_json_keyfile
             assert_nothing_raised { BigqueryClient.new(least_task, schema).client }
           end
-
-          def test_p12_keyfile
-            # pending
-          end
         end
 
         sub_test_case "create_dataset" do

data/test/test_configure.rb

@@ -18,10 +18,9 @@ module Embulk
 
       def least_config
         DataSource.new({
-          'project' => 'your_project_name',
-          'dataset' => 'your_dataset_name',
-          'table'   => 'your_table_name',
-          'p12_keyfile' => __FILE__, # fake
+          'project'      => 'your_project_name',
+          'dataset'      => 'your_dataset_name',
+          'table'        => 'your_table_name',
         })
       end
 
@@ -43,9 +42,7 @@ module Embulk
       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 "application_default", task['auth_method']
         assert_equal nil, task['json_keyfile']
         assert_equal "your_project_name", task['project']
         assert_equal "your_dataset_name", task['dataset']
@@ -55,14 +52,13 @@ module Embulk
         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 true, 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']
@@ -133,11 +129,6 @@ module Embulk
         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")
@@ -162,22 +153,22 @@ module Embulk
       end
 
       def test_payload_column
-        config = least_config.merge('payload_column' => schema.first.name)
+        config = least_config.merge('payload_column' => schema.first.name, 'auto_create_table' => false, 'mode' => 'append_direct')
         task = Bigquery.configure(config, schema, processor_count)
         assert_equal task['payload_column_index'], 0
 
-        config = least_config.merge('payload_column' => 'not_exist')
+        config = least_config.merge('payload_column' => 'not_exist', 'auto_create_table' => false, 'mode' => 'append_direct')
         assert_raise { Bigquery.configure(config, schema, processor_count) }
       end
 
       def test_payload_column_index
-        config = least_config.merge('payload_column_index' => 0)
+        config = least_config.merge('payload_column_index' => 0, 'auto_create_table' => false, 'mode' => 'append_direct')
         assert_nothing_raised { Bigquery.configure(config, schema, processor_count) }
 
-        config = least_config.merge('payload_column_index' => -1)
+        config = least_config.merge('payload_column_index' => -1, 'auto_create_table' => false, 'mode' => 'append_direct')
         assert_raise { Bigquery.configure(config, schema, processor_count) }
 
-        config = least_config.merge('payload_column_index' => schema.size)
+        config = least_config.merge('payload_column_index' => schema.size, 'auto_create_table' => false, 'mode' => 'append_direct')
         assert_raise { Bigquery.configure(config, schema, processor_count) }
       end
 

data/test/test_example.rb

@@ -33,7 +33,6 @@ else
     files.each do |config_path|
       if %w[
         config_expose_errors.yml
-        config_prevent_duplicate_insert.yml
         ].include?(File.basename(config_path))
         define_method(:"test_#{File.basename(config_path, ".yml")}") do
           assert_false embulk_run(config_path)

data/test/test_helper.rb

@@ -62,7 +62,8 @@ module Embulk
             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}),
+            Column.new({index: 5, name: 'date', type: :timestamp}),
+            Column.new({index: 6, name: 'json', type: :json}),
           ])
           task = {
             'column_options' => [
@@ -71,6 +72,7 @@ module Embulk
               {'name' => 'double',    'type' => 'STRING'},
               {'name' => 'string',    'type' => 'INTEGER'},
               {'name' => 'timestamp', 'type' => 'INTEGER'},
+              {'name' => 'date',      'type' => 'DATE'},
               {'name' => 'json',      'type' => 'RECORD', 'fields' => [
                 { 'name' => 'key1',   'type' => 'STRING' },
               ]},
@@ -82,6 +84,7 @@ module Embulk
             {name: 'double',    type: 'STRING'},
             {name: 'string',    type: 'INTEGER'},
             {name: 'timestamp', type: 'INTEGER'},
+            {name: 'date',      type: 'DATE'},
             {name: 'json',      type: 'RECORD', fields: [
               {name: 'key1',    type: 'STRING'},
             ]},

data/test/test_transaction.rb

@@ -8,12 +8,11 @@ module Embulk
     class TestTransaction < Test::Unit::TestCase
       def least_config
         DataSource.new({
-          'project'     => 'your_project_name',
-          'dataset'     => 'your_dataset_name',
-          'table'       => 'your_table_name',
-          'p12_keyfile' => __FILE__, # fake
-          'temp_table'  => 'temp_table', # randomly created is not good for our test
-          'path_prefix' => 'tmp/', # randomly created is not good for our test
+          'project'      => 'your_project_name',
+          'dataset'      => 'your_dataset_name',
+          'table'        => 'your_table_name',
+          'temp_table'   => 'temp_table', # randomly created is not good for our test
+          'path_prefix'  => 'tmp/', # randomly created is not good for our test
         })
       end
 
@@ -41,8 +40,8 @@ module Embulk
       end
 
       sub_test_case "append_direct" do
-        def test_append_direct
-          config = least_config.merge('mode' => 'append_direct')
+        def test_append_direc_without_auto_create
+          config = least_config.merge('mode' => 'append_direct', 'auto_create_dataset' => false, 'auto_create_table' => false)
           any_instance_of(BigqueryClient) do |obj|
             mock(obj).get_dataset(config['dataset'])
             mock(obj).get_table(config['table'])
@@ -60,8 +59,8 @@ module Embulk
           Bigquery.transaction(config, schema, processor_count, &control)
         end
 
-        def test_append_direct_with_partition
-          config = least_config.merge('mode' => 'append_direct', 'table' => 'table$20160929')
+        def test_append_direct_with_partition_without_auto_create
+          config = least_config.merge('mode' => 'append_direct', 'table' => 'table$20160929', 'auto_create_dataset' => false, 'auto_create_table' => false)
           any_instance_of(BigqueryClient) do |obj|
             mock(obj).get_dataset(config['dataset'])
             mock(obj).get_table(config['table'])
@@ -86,7 +85,7 @@ module Embulk
           task = Bigquery.configure(config, schema, processor_count)
           any_instance_of(BigqueryClient) do |obj|
             mock(obj).get_dataset(config['dataset'])
-            mock(obj).delete_partition(config['table'])
+            mock(obj).delete_table_or_partition(config['table'])
             mock(obj).create_table_if_not_exists(config['table'])
           end
           Bigquery.transaction(config, schema, processor_count, &control)
@@ -97,7 +96,7 @@ module Embulk
           task = Bigquery.configure(config, schema, processor_count)
           any_instance_of(BigqueryClient) do |obj|
             mock(obj).get_dataset(config['dataset'])
-            mock(obj).delete_partition(config['table'])
+            mock(obj).delete_table_or_partition(config['table'])
             mock(obj).create_table_if_not_exists(config['table'])
           end
           Bigquery.transaction(config, schema, processor_count, &control)
@@ -111,6 +110,7 @@ module Embulk
           any_instance_of(BigqueryClient) do |obj|
             mock(obj).get_dataset(config['dataset'])
             mock(obj).create_table_if_not_exists(config['temp_table'])
+            mock(obj).create_table_if_not_exists(config['table'])
             mock(obj).copy(config['temp_table'], config['table'], write_disposition: 'WRITE_TRUNCATE')
             mock(obj).delete_table(config['temp_table'])
           end
@@ -120,19 +120,6 @@ module Embulk
         def test_replace_with_partitioning
           config = least_config.merge('mode' => 'replace', 'table' => 'table$20160929')
           task = Bigquery.configure(config, schema, processor_count)
-          any_instance_of(BigqueryClient) do |obj|
-            mock(obj).get_dataset(config['dataset'])
-            mock(obj).create_table_if_not_exists(config['temp_table'])
-            mock(obj).get_table(config['table'])
-            mock(obj).copy(config['temp_table'], config['table'], write_disposition: 'WRITE_TRUNCATE')
-            mock(obj).delete_table(config['temp_table'])
-          end
-          Bigquery.transaction(config, schema, processor_count, &control)
-        end
-
-        def test_replace_with_partitioning_with_auto_create_table
-          config = least_config.merge('mode' => 'replace', 'table' => 'table$20160929', 'auto_create_table' => true)
-          task = Bigquery.configure(config, schema, processor_count)
           any_instance_of(BigqueryClient) do |obj|
             mock(obj).get_dataset(config['dataset'])
             mock(obj).create_table_if_not_exists(config['temp_table'])
@@ -152,8 +139,10 @@ module Embulk
             mock(obj).get_dataset(config['dataset'])
             mock(obj).get_dataset(config['dataset_old'])
             mock(obj).create_table_if_not_exists(config['temp_table'])
+            mock(obj).create_table_if_not_exists(config['table'])
+            mock(obj).create_table_if_not_exists(config['table_old'], dataset: config['dataset_old'])
 
-            mock(obj).get_table_or_partition(task['table'])
+            mock(obj).get_table_or_partition(config['table'])
             mock(obj).copy(config['table'], config['table_old'], config['dataset_old'])
 
             mock(obj).copy(config['temp_table'], config['table'], write_disposition: 'WRITE_TRUNCATE')
@@ -168,9 +157,11 @@ module Embulk
           any_instance_of(BigqueryClient) do |obj|
             mock(obj).create_dataset(config['dataset'])
             mock(obj).create_dataset(config['dataset_old'], reference: config['dataset'])
+            mock(obj).create_table_if_not_exists(config['table'])
             mock(obj).create_table_if_not_exists(config['temp_table'])
+            mock(obj).create_table_if_not_exists(config['table_old'], dataset: config['dataset_old'])
 
-            mock(obj).get_table_or_partition(task['table'])
+            mock(obj).get_table_or_partition(config['table'])
             mock(obj).copy(config['table'], config['table_old'], config['dataset_old'])
 
             mock(obj).copy(config['temp_table'], config['table'], write_disposition: 'WRITE_TRUNCATE')
@@ -180,35 +171,16 @@ module Embulk
         end
 
         def test_replace_backup_with_partitioning
-          config = least_config.merge('mode' => 'replace_backup', 'table' => 'table$20160929', 'dataset_old' => 'dataset_old', 'table_old' => 'table_old$20190929', 'temp_table' => 'temp_table')
-          task = Bigquery.configure(config, schema, processor_count)
-          any_instance_of(BigqueryClient) do |obj|
-            mock(obj).get_dataset(config['dataset'])
-            mock(obj).get_dataset(config['dataset_old'])
-            mock(obj).create_table_if_not_exists(config['temp_table'])
-            mock(obj).get_table(task['table'])
-            mock(obj).get_table(task['table_old'], dataset: config['dataset_old'])
-
-            mock(obj).get_table_or_partition(task['table'])
-            mock(obj).copy(config['table'], config['table_old'], config['dataset_old'])
-
-            mock(obj).copy(config['temp_table'], config['table'], write_disposition: 'WRITE_TRUNCATE')
-            mock(obj).delete_table(config['temp_table'])
-          end
-          Bigquery.transaction(config, schema, processor_count, &control)
-        end
-
-        def test_replace_backup_with_partitioning_auto_create_table
           config = least_config.merge('mode' => 'replace_backup', 'table' => 'table$20160929', 'dataset_old' => 'dataset_old', 'table_old' => 'table_old$20160929', 'temp_table' => 'temp_table', 'auto_create_table' => true)
           task = Bigquery.configure(config, schema, processor_count)
           any_instance_of(BigqueryClient) do |obj|
             mock(obj).get_dataset(config['dataset'])
             mock(obj).get_dataset(config['dataset_old'])
             mock(obj).create_table_if_not_exists(config['temp_table'])
-            mock(obj).create_table_if_not_exists(task['table'])
-            mock(obj).create_table_if_not_exists(task['table_old'], dataset: config['dataset_old'])
+            mock(obj).create_table_if_not_exists(config['table'])
+            mock(obj).create_table_if_not_exists(config['table_old'], dataset: config['dataset_old'])
 
-            mock(obj).get_table_or_partition(task['table'])
+            mock(obj).get_table_or_partition(config['table'])
             mock(obj).copy(config['table'], config['table_old'], config['dataset_old'])
 
             mock(obj).copy(config['temp_table'], config['table'], write_disposition: 'WRITE_TRUNCATE')
@@ -225,6 +197,7 @@ module Embulk
           any_instance_of(BigqueryClient) do |obj|
             mock(obj).get_dataset(config['dataset'])
             mock(obj).create_table_if_not_exists(config['temp_table'])
+            mock(obj).create_table_if_not_exists(config['table'])
             mock(obj).copy(config['temp_table'], config['table'], write_disposition: 'WRITE_APPEND')
             mock(obj).delete_table(config['temp_table'])
           end
@@ -232,19 +205,6 @@ module Embulk
         end
 
         def test_append_with_partitioning
-          config = least_config.merge('mode' => 'append', 'table' => 'table$20160929')
-          task = Bigquery.configure(config, schema, processor_count)
-          any_instance_of(BigqueryClient) do |obj|
-            mock(obj).get_dataset(config['dataset'])
-            mock(obj).create_table_if_not_exists(config['temp_table'])
-            mock(obj).get_table(config['table'])
-            mock(obj).copy(config['temp_table'], config['table'], write_disposition: 'WRITE_APPEND')
-            mock(obj).delete_table(config['temp_table'])
-          end
-          Bigquery.transaction(config, schema, processor_count, &control)
-        end
-
-        def test_append_with_partitioning_with_auto_create_table
           config = least_config.merge('mode' => 'append', 'table' => 'table$20160929', 'auto_create_table' => true)
           task = Bigquery.configure(config, schema, processor_count)
           any_instance_of(BigqueryClient) do |obj|