naginegi 0.3.0

data/Rakefile

@@ -0,0 +1,6 @@
+require 'bundler/gem_tasks'
+require 'rspec/core/rake_task'
+
+RSpec::Core::RakeTask.new(:spec)
+
+task :default => :spec

data/bin/console

@@ -0,0 +1,14 @@
+#!/usr/bin/env ruby
+
+require "bundler/setup"
+require "naginegi"
+
+# You can add fixtures and/or initialization code here to make experimenting
+# with your gem easier. You can also use a different console, if you like.
+
+# (If you use this, don't forget to add pry to your Gemfile!)
+# require "pry"
+# Pry.start
+
+require "irb"
+IRB.start(__FILE__)

data/bin/setup

@@ -0,0 +1,8 @@
+#!/usr/bin/env bash
+set -euo pipefail
+IFS=$'\n\t'
+set -vx
+
+bundle install
+
+# Do any other automated setup that you need to do here

data/lib/naginegi.rb

@@ -0,0 +1,57 @@
+require 'naginegi/version'
+require 'naginegi/embulk_config'
+require 'naginegi/embulk'
+require 'naginegi/mysql'
+require 'naginegi/postgresql'
+require 'logger'
+
+module Naginegi
+  class EmbulkRunner
+    def initialize(db_configs: nil, log_level: 'warn', embulk_run_option: '')
+      @logger = Logger.new(STDOUT)
+      @logger.datetime_format = '%Y-%m-%d %H:%M:%S'
+
+      @db_configs = db_configs || YAML.load_file('database.yml')
+      @log_level = log_level
+      @embulk_run_option = embulk_run_option
+    end
+
+    def generate_config(bq_config)
+      Naginegi::EmbulkConfig.new.generate_config(@db_configs, bq_config)
+    end
+
+    def run(bq_config, target_table_names = [], retry_max = 0)
+      cmd = 'embulk --version'
+      unless system(cmd)
+        @logger.error('Cannot execute Embulk!!')
+        @logger.error('Cofirm Embulk install and environment')
+        return
+      end
+
+      error_tables = run_and_retry(bq_config, target_table_names, retry_max, 0)
+      error_tables.empty?
+    end
+
+    private
+
+    def run_and_retry(bq_config, target_table_names, retry_max, retry_count)
+      error_tables = Naginegi::Embulk.new(@log_level, @embulk_run_option).run(
+        @db_configs,
+        table_configs,
+        bq_config,
+        target_table_names
+      )
+      if !error_tables.empty? && retry_count < retry_max
+        @logger.warn('------------------------------------')
+        @logger.warn("retry start -> #{retry_count + 1} time")
+        @logger.warn('------------------------------------')
+        error_tables = run_and_retry(bq_config, error_tables, retry_max, retry_count + 1)
+      end
+      error_tables
+    end
+
+    def table_configs
+      @table_configs ||= Naginegi::TableConfig.generate_table_configs
+    end
+  end
+end

data/lib/naginegi/bigquery.rb

@@ -0,0 +1,107 @@
+require 'json'
+require 'erb'
+require 'google/cloud/bigquery'
+require 'unindent'
+require 'date'
+
+module Naginegi
+  class BigQuery
+    CONTENTS = <<-EOS.unindent
+    in:
+      type: <%= db_type %>
+      host: <%= host %>
+      user: <%= user %>
+      password: <%= password %>
+      database: <%= database %>
+      ssl: <%= ssl %>
+      query: |
+        <%= query %>
+      <%= options %>
+    out:
+      type: bigquery
+      auth_method: <%= auth_method %>
+      json_keyfile: <%= json_keyfile %>
+        <%= json_key_content %>
+      project: <%= project %>
+      service_account_email: <%= service_account_email %>
+      dataset: <%= dataset %>
+      table: <%= table_name %>
+      schema_file: <%= schema_file %>
+      auto_create_table: true
+      path_prefix: <%= path_prefix %>
+      source_format: NEWLINE_DELIMITED_JSON
+      file_ext: .json.gz
+      delete_from_local_when_job_end: 1
+      formatter:
+        type: jsonl
+      encoders:
+      - {type: gzip}
+    EOS
+
+    def initialize(config)
+      @config = config.dup
+      @current_date = Date.today
+    end
+
+    def self.generate_schema(columns)
+      json_body = columns.map(&:to_json).join(",\n")
+      "[\n" + json_body + "\n]\n"
+    end
+
+    def self.generate_sql(table_config, columns)
+      columns = columns.map(&:converted_value)
+      sql = "SELECT #{columns.join(',')}"
+      sql << " FROM #{table_config.name}"
+      sql << " WHERE #{table_config.condition}" if table_config.condition
+      sql << "\n"
+      sql
+    end
+
+    def generate_embulk_config(db_name, db_config, table_config, columns)
+      db_type = db_config['db_type']
+      host = db_config['host']
+      user = db_config['username']
+      password = db_config['password']
+      database = db_config['database']
+      ssl = db_config['embulk_ssl_enable'] || false
+      options = if db_type == 'mysql'
+                  "options: {useLegacyDatetimeCode: false, serverTimezone: #{db_config['timezone']}}"
+                else
+                  ''
+                end
+      query = Naginegi::BigQuery.generate_sql(table_config, columns)
+
+      auth_method = @config['auth_method']
+      if @config['json_key']
+        values = @config['json_key'].map do |k, v|
+          value = v.gsub("\n", '\\n')
+          "\"#{k}\": \"#{value}\""
+        end
+        json_key_content = "content: |\n      {#{values.join(',')}}"
+      else
+        json_keyfile = @config['json_keyfile']
+      end
+      project = @config['project_id']
+      service_account_email = @config['service_email']
+      dataset = db_config['bq_dataset']
+      table_name = actual_table_name(table_config.name, db_config['daily_snapshot'] || table_config.daily_snapshot)
+      schema_file = "#{@config['schema_dir']}/#{db_name}/#{table_config.name}.json"
+      path_prefix = "/var/tmp/embulk_#{db_name}_#{table_config.name}"
+
+      ERB.new(CONTENTS).result(binding)
+    end
+
+    def delete_table(dataset, table_name)
+      bq = Google::Cloud::Bigquery.new(
+        project: @config['project_id'],
+        keyfile: @config['json_keyfile']
+      )
+      bq.service.delete_table(dataset, table_name)
+    end
+
+    def actual_table_name(table_name, daily_snapshot)
+      return table_name unless daily_snapshot
+      table_name + @current_date.strftime('%Y%m%d')
+    end
+  end
+end

data/lib/naginegi/embulk.rb

@@ -0,0 +1,75 @@
+require 'logger'
+
+module Naginegi
+  class Embulk
+    def initialize(log_level, embulk_run_option)
+      @logger = Logger.new(STDOUT)
+      @logger.datetime_format = '%Y-%m-%d %H:%M:%S'
+
+      @log_level = log_level
+      @embulk_run_option = embulk_run_option
+    end
+
+    def run(db_configs, all_table_configs, bq_config, target_table_names = [])
+      error_tables = []
+      db_configs.keys.each do |db_name|
+        table_configs = select_table_configs(all_table_configs[db_name], target_table_names)
+        error_tables += run_by_database(
+          db_name,
+          table_configs,
+          bq_config,
+          db_configs[db_name]['bq_dataset']
+        )
+      end
+      error_tables
+    end
+
+    def select_table_configs(table_configs, target_table_names)
+      return table_configs if target_table_names.empty?
+      table_configs.select { |table_config| target_table_names.include?(table_config.name) }
+    end
+
+    private
+
+    def run_by_database(db_name, table_configs, bq_config, bq_dataset)
+      process_times = []
+      error_tables = []
+
+      bq_utility = Naginegi::BigQuery.new(bq_config)
+
+      table_configs.each do |table_config|
+        start_time = Time.now
+        @logger.info("table: #{table_config.name} - start")
+
+        begin
+          bq_utility.delete_table(bq_dataset, table_config.name)
+          @logger.info("#{table_config.name} is deleted")
+        rescue => e
+          @logger.warn(e.message)
+        end
+
+        cmd = "embulk run #{@embulk_run_option} #{bq_config['config_dir']}/#{db_name}/#{table_config.name}.yml --log-level #{@log_level}"
+        @logger.info("cmd: #{cmd}")
+
+        if system(cmd)
+          result = 'success'
+        else
+          result = 'error'
+          error_tables << table_config.name
+        end
+
+        process_time = "table: #{table_config.name} - result: #{result}  #{format('%10.1f', Time.now - start_time)}sec"
+        @logger.info(process_time)
+
+        process_times << process_time
+      end
+
+      @logger.info('------------------------------------')
+      @logger.info("db_name: #{db_name}")
+
+      process_times.each { |process_time| @logger.info(process_time) }
+
+      error_tables
+    end
+  end
+end

data/lib/naginegi/embulk_config.rb

@@ -0,0 +1,73 @@
+module Naginegi
+  class EmbulkConfig
+    def generate_config(db_configs, bq_config)
+      bq_utility = BigQuery.new(bq_config)
+
+      db_configs.keys.each do |db_name|
+        db_config = db_configs[db_name]
+        table_configs = all_table_configs[db_name]
+        db_type = db_config['db_type']
+
+        case db_type
+        when 'mysql'
+          sql_client = MySQL::MySQLClient.new(db_config)
+        when 'postgresql'
+          sql_client = PostgreSQL::PgClient.new(db_config)
+        end
+
+        table_configs.each do |table_config|
+          write(
+            "#{bq_config['schema_dir']}/#{db_name}",
+            "#{table_config.name}.json",
+            sql_client.generate_bq_schema(table_config.name)
+          )
+          write(
+            "#{bq_config['config_dir']}/#{db_name}",
+            "#{table_config.name}.yml",
+            bq_utility.generate_embulk_config(
+              db_name,
+              db_config,
+              table_config,
+              sql_client.columns(table_config.name)
+            )
+          )
+        end
+      end
+    end
+
+    private
+
+    def write(directory, file_name, content)
+      FileUtils.mkdir_p(directory) unless FileTest.exist?(directory)
+      File.write("#{directory}/#{file_name}", content)
+    end
+
+    def all_table_configs
+      @all_table_configs ||= Naginegi::TableConfig.generate_table_configs
+    end
+  end
+
+  class TableConfig
+    attr_reader :name, :daily_snapshot, :condition
+
+    def initialize(config)
+      @name = config['name']
+      @daily_snapshot = config['daily_snapshot'] || false
+      @condition = config['condition']
+    end
+
+    def self.generate_table_configs(file_path = 'table.yml')
+      configs = YAML.load_file(file_path)
+      configs.each_with_object({}) do |(db, database_config), table_configs|
+        table_configs[db] = database_config['tables'].map { |config| TableConfig.new(config) }
+        table_configs
+      end
+    end
+
+    def ==(other)
+      instance_variables.all? do |v|
+        instance_variable_get(v) == other.instance_variable_get(v)
+      end
+    end
+  end
+end

data/lib/naginegi/mysql.rb

@@ -0,0 +1,95 @@
+require 'mysql2-cs-bind'
+require 'json'
+require 'yaml'
+require 'fileutils'
+require 'naginegi/bigquery'
+
+module Naginegi
+  module MySQL
+    class MySQLClient
+      COLUMN_SQL = <<-SQL.freeze
+        SELECT column_name, data_type
+        FROM INFORMATION_SCHEMA.COLUMNS
+        WHERE table_schema = ?
+        AND table_name = ?
+        ORDER BY ordinal_position
+      SQL
+
+      def initialize(database_config)
+        @database_config = database_config
+      end
+
+      def client
+        @client ||= Mysql2::Client.new(
+          host: @database_config['host'],
+          username: @database_config['username'],
+          password: @database_config['password'],
+          database: @database_config['database']
+        )
+      end
+
+      def generate_bq_schema(table_name)
+        infos = columns(table_name)
+        BigQuery.generate_schema(infos)
+      end
+
+      def columns(table_name)
+        rows = client.xquery(COLUMN_SQL, @database_config['database'], table_name)
+        rows.map { |row| Column.new(row['column_name'], row['data_type']) }
+      end
+    end
+
+    class Column
+      attr_reader :column_name, :data_type
+
+      TYPE_MAPPINGS = {
+        'int' => 'INT64',
+        'tinyint' => 'INT64',
+        'smallint' => 'INT64',
+        'mediumint' => 'INT64',
+        'bigint' => 'INT64',
+        'float' => 'FLOAT64',
+        'double' => 'FLOAT64',
+        'decimal' => 'FLOAT64',
+        'char' => 'STRING',
+        'varchar' => 'STRING',
+        'tinytext' => 'STRING',
+        'text' => 'STRING',
+        'date' => 'TIMESTAMP',
+        'datetime' => 'TIMESTAMP',
+        'timestamp' => 'TIMESTAMP'
+      }.freeze
+
+      def initialize(column_name, data_type)
+        @column_name = column_name
+        @data_type = data_type
+      end
+
+      def bigquery_data_type
+        TYPE_MAPPINGS[@data_type]
+      end
+
+      def converted_value
+        if bigquery_data_type == 'TIMESTAMP'
+          # time zone translate to UTC
+          "UNIX_TIMESTAMP(#{escaped_column_name}) AS #{escaped_column_name}"
+        elsif data_type == 'tinyint'
+          # for MySQL tinyint(1) problem
+          "CAST(#{escaped_column_name} AS signed) AS #{escaped_column_name}"
+        else
+          escaped_column_name
+        end
+      end
+
+      def to_json(*a)
+        { 'name' => @column_name, 'type' => bigquery_data_type }.to_json(*a)
+      end
+
+      private
+
+      def escaped_column_name
+        "`#{@column_name}`"
+      end
+    end
+  end
+end

data/lib/naginegi/postgresql.rb

@@ -0,0 +1,93 @@
+require 'pg'
+require 'json'
+require 'yaml'
+require 'fileutils'
+require 'naginegi/bigquery'
+
+module Naginegi
+  module PostgreSQL
+    class PgClient
+      COLUMN_SQL = <<-SQL.freeze
+        SELECT column_name, data_type
+        FROM INFORMATION_SCHEMA.COLUMNS
+        WHERE table_name = $1
+        ORDER BY ordinal_position
+      SQL
+
+      def initialize(db_config)
+        @db_config = db_config
+      end
+
+      def client
+        @client ||= PG::Connection.new(
+          host: @db_config['host'],
+          user: @db_config['username'],
+          password: @db_config['password'],
+          dbname: @db_config['database']
+        )
+      end
+
+      def generate_bq_schema(table_name)
+        infos = columns(table_name)
+        BigQuery.generate_schema(infos)
+      end
+
+      def columns(table_name)
+        rows = client.exec_params(COLUMN_SQL, [table_name])
+        rows.map { |row| Column.new(row['column_name'], row['data_type']) }
+      end
+    end
+
+    class Column
+      attr_reader :column_name, :data_type
+
+      TYPE_MAPPINGS = {
+        'smallint' => 'INT64',
+        'integer' => 'INT64',
+        'bigint' => 'INT64',
+        'smallserial' => 'INT64',
+        'serial' => 'INT64',
+        'bigserial' => 'INT64',
+        'decimal' => 'FLOAT64',
+        'numeric' => 'FLOAT64',
+        'real' => 'FLOAT64',
+        'double precision' => 'FLOAT64',
+        'character' => 'STRING',
+        'character varying' => 'STRING',
+        'text' => 'STRING',
+        'date' => 'TIMESTAMP',
+        'timestamp' => 'TIMESTAMP',
+        'timestamp with time zone' => 'TIMESTAMP',
+        'boolean' => 'BOOL'
+      }.freeze
+
+      def initialize(column_name, data_type)
+        @column_name = column_name
+        @data_type = data_type
+      end
+
+      def bigquery_data_type
+        TYPE_MAPPINGS[@data_type]
+      end
+
+      def converted_value
+        if bigquery_data_type == 'TIMESTAMP'
+          # time zone translate to UTC
+          "EXTRACT(EPOCH FROM #{escaped_column_name}) AS #{escaped_column_name}"
+        else
+          escaped_column_name
+        end
+      end
+
+      def to_json(*a)
+        { 'name' => @column_name, 'type' => bigquery_data_type }.to_json(*a)
+      end
+
+      private
+
+      def escaped_column_name
+        "\"#{@column_name}\""
+      end
+    end
+  end
+end