embulk-output-bigquery 0.3.0.pre5 → 0.3.0.pre6

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA1:
-  metadata.gz: 278ad1e9acdd3c7ae22ef63b9304e4c5fec6871e
-  data.tar.gz: a310f2abd73d6cd07808fedc21a11b649a4339ac
+  metadata.gz: ce1897b1eda7e46719b031b5f7e960cdce0ac014
+  data.tar.gz: 1bab2f7575b470df798c6aad4ecaa26d00e43b8c
 SHA512:
-  metadata.gz: 008bc08aaa1b48718b4e0cd6d79b24def61eaa350274560c42212571fe99c7f67dfbd1644186a8d5c3f87de367ccb0d6346b091a994e87bfb31f8d03b42919f3
-  data.tar.gz: f0f87a847877b5c70e95088424e657796112d35b044f6b069733446ade00206308ca39d106751750a4a47229c552c3ac91e487dc25b2c5415127c2aeec7ee6a5
+  metadata.gz: 772f2b087ede71e6edb9d3ccae7b870b389db66f2b21ef0e559a34c8923eaf2900fe1dc8532bea2bde8505c97434c81289e5f0dff1e1141c9eb31bf5df9773ce
+  data.tar.gz: f6037ccfbfbd3a0c440cf3d65b3bb5b1e62976267ed3299e7581542deb04f79308c6e377b4e53f61a8461a8195b938e417d7b69ac227654d41be8d13ae4a28df

data/embulk-output-bigquery.gemspec

@@ -1,6 +1,6 @@
 Gem::Specification.new do |spec|
   spec.name          = "embulk-output-bigquery"
-  spec.version       = "0.3.0.pre5"
+  spec.version       = "0.3.0.pre6"
   spec.authors       = ["Satoshi Akama", "Naotoshi Seo"]
   spec.summary       = "Google BigQuery output plugin for Embulk"
   spec.description   = "Embulk plugin that insert records to Google BigQuery."

data/example/config_with_rehearsal.yml

@@ -30,3 +30,4 @@ out:
   with_rehearsal: true
   rehearsal_counts: 1
   skip_load: true # for debug
+  compression: GZIP

data/lib/embulk/output/bigquery.rb

@@ -1,5 +1,6 @@
 require 'json'
 require 'tempfile'
+require 'fileutils'
 require_relative 'bigquery/bigquery_client'
 require_relative 'bigquery/file_writer'
 require_relative 'bigquery/value_converter_factory'
@@ -211,23 +212,27 @@ module Embulk
         @converters
       end
 
+      def self.rehearsal_thread
+        @rehearsal_thread
+      end
+
+      def self.rehearsal_thread=(rehearsal_thread)
+        @rehearsal_thread = rehearsal_thread
+      end
+
       def self.transaction_report(task_reports, responses)
+        num_input_rows = task_reports.inject(0) do |sum, task_report|
+          sum + task_report['num_input_rows']
+        end
+        num_output_rows = responses.inject(0) do |sum, response|
+          sum + (response ? response.statistics.load.output_rows.to_i : 0)
+        end
+        num_rejected_rows = num_input_rows - num_output_rows
         transaction_report = {
-          'num_input_rows' => 0,
-          'num_output_rows' => 0,
-          'num_rejected_rows' => 0,
+          'num_input_rows' => num_input_rows,
+          'num_output_rows' => num_output_rows,
+          'num_rejected_rows' => num_rejected_rows,
         }
-        (0...task_reports.size).each do |idx|
-          task_report       = task_reports[idx]
-          response          = responses[idx]
-          num_input_rows    = task_report['num_input_rows']
-          num_output_rows   = response ? response.statistics.load.output_rows.to_i : 0
-          num_rejected_rows = num_input_rows - num_output_rows
-          transaction_report['num_input_rows']    += num_input_rows
-          transaction_report['num_output_rows']   += num_output_rows
-          transaction_report['num_rejected_rows'] += num_rejected_rows
-        end
-        transaction_report
       end
 
       def self.transaction(config, schema, task_count, &control)
@@ -278,7 +283,14 @@ module Embulk
             task_reports = yield(task) # generates local files
             Embulk.logger.info { "embulk-output-bigquery: task_reports: #{task_reports.to_json}" }
             paths = FileWriter.paths
-            FileWriter.ios.each {|io| io.close rescue nil }
+            FileWriter.ios.values.each do |io|
+              Embulk.logger.debug { "close #{io.path}" }
+              io.close rescue nil
+            end
+          end
+
+          if rehearsal_thread
+            rehearsal_thread.join
           end
 
           if task['skip_load'] # only for debug
@@ -332,7 +344,6 @@ module Embulk
         super
 
         if task['with_rehearsal'] and @index == 0
-          @bigquery = self.class.bigquery
           @rehearsaled = false
           @num_rows = 0
         end
@@ -351,13 +362,7 @@ module Embulk
         if task['with_rehearsal'] and @index == 0 and !@rehearsaled
           page = page.to_a # to avoid https://github.com/embulk/embulk/issues/403
           if @num_rows >= task['rehearsal_counts']
-            Embulk.logger.info { "embulk-output-bigquery: Rehearsal started" }
-            begin
-              @bigquery.create_table(task['rehearsal_table'])
-              @bigquery.load(FileWriter.paths.first, task['rehearsal_table'])
-            ensure
-              @bigquery.delete_table(task['rehearsal_table'])
-            end
+            load_rehearsal
             @rehearsaled = true
           end
           @num_rows += page.to_a.size
@@ -368,6 +373,30 @@ module Embulk
         end
       end
 
+      def load_rehearsal
+        bigquery = self.class.bigquery
+        Embulk.logger.info { "embulk-output-bigquery: Rehearsal started" }
+
+        io = @file_writer.close # need to close once for gzip
+        rehearsal_path = "#{io.path}.rehearsal"
+        Embulk.logger.debug { "embulk_output_bigquery: cp #{io.path} #{rehearsal_path}" }
+        FileUtils.cp(io.path, rehearsal_path)
+        @file_writer.reopen
+
+        self.class.rehearsal_thread = Thread.new do
+          begin
+            bigquery.create_table(task['rehearsal_table'])
+            response = bigquery.load(rehearsal_path, task['rehearsal_table'])
+            num_output_rows = response ? response.statistics.load.output_rows.to_i : 0
+            Embulk.logger.info { "embulk-output-bigquery: Loaded rehearsal #{num_output_rows}" }
+          ensure
+            Embulk.logger.debug { "embulk_output_bigquery: delete #{rehearsal_path}" }
+            File.unlink(rehearsal_path) rescue nil
+            bigquery.delete_table(task['rehearsal_table'])
+          end
+        end
+      end
+
       def finish
       end
 

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

@@ -31,14 +31,14 @@ module Embulk
         end
 
         @mutex = Mutex.new
-        @ios = Set.new
+        @ios = Hash.new
 
         def self.mutex
           @mutex
         end
 
         def self.reset_ios
-          @ios = Set.new
+          @ios = Hash.new
         end
 
         def self.ios
@@ -46,7 +46,7 @@ module Embulk
         end
 
         def self.paths
-          ios.map {|io| io.path }
+          @ios.keys
         end
 
         THREAD_LOCAL_IO_KEY = :embulk_output_bigquery_file_writer_io
@@ -69,22 +69,35 @@ module Embulk
             File.unlink(path) rescue nil
           end
           Embulk.logger.info { "embulk-output-bigquery: create #{path}" }
-          file_io = File.open(path, 'w')
 
+          open(path, 'w')
+        end
+
+        def open(path, mode = 'w')
+          file_io = File.open(path, mode)
           case @task['compression'].downcase
           when 'gzip'
             io = Zlib::GzipWriter.new(file_io)
           else
             io = file_io
           end
-
           self.class.mutex.synchronize do
-            self.class.ios.add(io)
+            self.class.ios[path] = io
           end
-
           Thread.current[THREAD_LOCAL_IO_KEY] = io
         end
 
+        def close
+          io = thread_io
+          io.close rescue nil
+          io
+        end
+
+        def reopen
+          io = thread_io
+          open(io.path, 'a')
+        end
+
         def to_payload(record)
           "#{record[@payload_column_index]}\n"
         end

data/test/test_file_writer.rb

@@ -108,7 +108,7 @@ module Embulk
 
           begin
             file_writer.add(page)
-            io = FileWriter.ios.first
+            io = FileWriter.ios.values.first
             assert_equal Zlib::GzipWriter, io.class
           ensure
             io.close rescue nil
@@ -124,7 +124,7 @@ module Embulk
 
           begin
             file_writer.add(page)
-            io = FileWriter.ios.first
+            io = FileWriter.ios.values.first
             assert_equal File, io.class
           ensure
             io.close rescue nil

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: embulk-output-bigquery
 version: !ruby/object:Gem::Version
-  version: 0.3.0.pre5
+  version: 0.3.0.pre6
 platform: ruby
 authors:
 - Satoshi Akama
@@ -9,7 +9,7 @@ authors:
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2016-03-24 00:00:00.000000000 Z
+date: 2016-03-30 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: google-api-client