<commit_before>require 'uri'
require 'java'
require 'google/apis/storage_v1'
require_relative 'google_client'
require_relative 'helper'

# ToDo: Use https://cloud.google.com/storage/docs/streaming if google-api-ruby-client supports streaming transfers
# ToDo: Tests are not written because this implementation will probably entirely changed on supporting streaming transfers
module Embulk
  module Output
    class Bigquery < OutputPlugin
      class GcsClient < GoogleClient
        def initialize(task)
          scope = "https://www.googleapis.com/auth/cloud-platform"
          client_class = Google::Apis::StorageV1::StorageService
          super(task, scope, client_class)

          @project = @task['project']
          @bucket = @task['gcs_bucket']
        end

        def insert_bucket(bucket = nil)
          bucket ||= @bucket
          begin
            Embulk.logger.info { "embulk-output-bigquery: Insert bucket... #{@project}:#{bucket}" }
            body  = {
              name: bucket,
            }
            opts = {}

            Embulk.logger.debug { "embulk-output-bigquery: insert_bucket(#{@project}, #{body}, #{opts})" }
            client.insert_bucket(@project, body, opts)
          rescue Google::Apis::ServerError, Google::Apis::ClientError, Google::Apis::AuthorizationError => e
            if e.status_code == 409 && /conflict:/ =~ e.message
              # ignore 'Already Exists' error
              return nil
            end
            response = {status_code: e.status_code, message: e.message, error_class: e.class}
            Embulk.logger.error {
              "embulk-output-bigquery: insert_bucket(#{@project}, #{body}, #{opts}), response:#{response}"
            }
            raise Error, "failed to insert bucket #{@project}:#{bucket}, response:#{response}"
          end
        end

        def insert_object(path, object: nil, bucket: nil)
          bucket ||= @bucket
          object ||= path
          object = object.start_with?('/') ? object[1..-1] : object
          object_uri = URI.join("gs://#{bucket}", object).to_s

          started = Time.now
          retries = 0
          begin
            Embulk.logger.info { "embulk-output-bigquery: Insert object... #{path} => #{@project}:#{object_uri}" }
            body = {
              name: object,
            }
            opts = {
              upload_source: path,
              content_type: 'application/octet-stream'
            }

            Embulk.logger.debug { "embulk-output-bigquery: insert_object(#{bucket}, #{body}, #{opts})" }
            # memo: gcs is strongly consistent for insert (read-after-write). ref: https://cloud.google.com/storage/docs/consistency
            client.insert_object(bucket, body, opts)
          rescue Google::Apis::ServerError, Google::Apis::ClientError, Google::Apis::AuthorizationError => e
            response = {status_code: e.status_code, message: e.message, error_class: e.class}
            Embulk.logger.error {
              "embulk-output-bigquery: insert_object(#{bucket}, #{body}, #{opts}), response:#{response}"
            }
            raise Error, "failed to insert object #{@project}:#{object_uri}, response:#{response}"
          rescue ::Java::Java.net.SocketException => e
            # I encountered `java.net.SocketException: Broken pipe` and `Connection reset` serveral times
            # google-api-ruby-client itself has a retry feature, but it does not retry with SocketException
            if e.message == 'Broken pipe' || e.message == 'Connection reset'
              if retries < @task['retries']
                response = {message: e.message, error_class: e.class}
                Embulk.logger.warn {
                  "embulk-output-bigquery: RETRY: insert_object(#{bucket}, #{body}, #{opts}), response:#{response}"
                }
                retries += 1 # want to share with google-api-ruby-client, but it is difficult
                retry
              end
            end
            raise e
          end
        end

        def insert_objects(paths, objects: nil, bucket: nil)
          return [] if paths.empty?
          bucket ||= @bucket
          objects ||= paths
          raise "number of paths and objects are different" if paths.size != objects.size

          responses = []
          paths.each_with_index do |path, idx|
            object = objects[idx]
            responses << insert_object(path, object: object, bucket: bucket)
          end
          responses
        end
<commit_msg>add a comment<commit_after>require 'uri'
require 'java'
require 'google/apis/storage_v1'
require_relative 'google_client'
require_relative 'helper'

# ToDo: Use https://cloud.google.com/storage/docs/streaming if google-api-ruby-client supports streaming transfers
# ToDo: Tests are not written because this implementation will probably entirely changed on supporting streaming transfers
module Embulk
  module Output
    class Bigquery < OutputPlugin
      class GcsClient < GoogleClient
        def initialize(task)
          scope = "https://www.googleapis.com/auth/cloud-platform"
          client_class = Google::Apis::StorageV1::StorageService
          super(task, scope, client_class)

          @project = @task['project']
          @bucket = @task['gcs_bucket']
        end

        def insert_bucket(bucket = nil)
          bucket ||= @bucket
          begin
            Embulk.logger.info { "embulk-output-bigquery: Insert bucket... #{@project}:#{bucket}" }
            body  = {
              name: bucket,
            }
            opts = {}

            Embulk.logger.debug { "embulk-output-bigquery: insert_bucket(#{@project}, #{body}, #{opts})" }
            client.insert_bucket(@project, body, opts)
          rescue Google::Apis::ServerError, Google::Apis::ClientError, Google::Apis::AuthorizationError => e
            if e.status_code == 409 && /conflict:/ =~ e.message
              # ignore 'Already Exists' error
              return nil
            end
            response = {status_code: e.status_code, message: e.message, error_class: e.class}
            Embulk.logger.error {
              "embulk-output-bigquery: insert_bucket(#{@project}, #{body}, #{opts}), response:#{response}"
            }
            raise Error, "failed to insert bucket #{@project}:#{bucket}, response:#{response}"
          end
        end

        def insert_object(path, object: nil, bucket: nil)
          bucket ||= @bucket
          object ||= path
          object = object.start_with?('/') ? object[1..-1] : object
          object_uri = URI.join("gs://#{bucket}", object).to_s

          started = Time.now
          retries = 0
          begin
            Embulk.logger.info { "embulk-output-bigquery: Insert object... #{path} => #{@project}:#{object_uri}" }
            body = {
              name: object,
            }
            opts = {
              upload_source: path,
              content_type: 'application/octet-stream'
            }

            Embulk.logger.debug { "embulk-output-bigquery: insert_object(#{bucket}, #{body}, #{opts})" }
            # memo: gcs is strongly consistent for insert (read-after-write). ref: https://cloud.google.com/storage/docs/consistency
            client.insert_object(bucket, body, opts)
          rescue Google::Apis::ServerError, Google::Apis::ClientError, Google::Apis::AuthorizationError => e
            response = {status_code: e.status_code, message: e.message, error_class: e.class}
            Embulk.logger.error {
              "embulk-output-bigquery: insert_object(#{bucket}, #{body}, #{opts}), response:#{response}"
            }
            raise Error, "failed to insert object #{@project}:#{object_uri}, response:#{response}"
          rescue ::Java::Java.net.SocketException => e
            # I encountered `java.net.SocketException: Broken pipe` and `Connection reset` serveral times
            # I am doubting as this is caused by Google's unstable network
            # google-api-ruby-client itself has a retry feature, but it does not retry with SocketException
            if e.message == 'Broken pipe' || e.message == 'Connection reset'
              if retries < @task['retries']
                response = {message: e.message, error_class: e.class}
                Embulk.logger.warn {
                  "embulk-output-bigquery: RETRY: insert_object(#{bucket}, #{body}, #{opts}), response:#{response}"
                }
                retries += 1 # want to share with google-api-ruby-client, but it is difficult
                retry
              end
            end
            raise e
          end
        end

        def insert_objects(paths, objects: nil, bucket: nil)
          return [] if paths.empty?
          bucket ||= @bucket
          objects ||= paths
          raise "number of paths and objects are different" if paths.size != objects.size

          responses = []
          paths.each_with_index do |path, idx|
            object = objects[idx]
            responses << insert_object(path, object: object, bucket: bucket)
          end
          responses
        end

{'fraction_non_alphanumeric': 0.11888803428194704, 'fraction_numerical': 0.001977804636853093, 'mean_word_length': 2.1117948717948716, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 11, 'https://': 6, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4810544', 'n_tokens_mistral': 2453, 'n_tokens_neox': 2406, 'n_words': 627}