kiba-plus 0.1.0

data/lib/kiba/plus/destination/pg.rb

@@ -0,0 +1,51 @@
+require 'pg'
+
+module Kiba::Plus::Destination
+  class Pg
+    attr_reader :options
+
+    def initialize(options = {})
+      @options = options
+      @options.assert_valid_keys(
+        :connect_url,
+        :prepare_name,
+        :prepare_sql,
+        :columns
+      )
+      @conn = PG.connect(connect_url)
+      @conn.prepare(prepare_name, prepare_sql)
+    end
+
+    def write(row)
+      @conn.exec_prepared(prepare_name,
+        row.values_at(*columns))
+    rescue PG::Error => ex
+      Kiba::Plus.logger.error "ERROR for #{row}"
+      Kiba::Plus.logger.error ex.message
+      # Maybe, write to db table or file
+    end
+
+    def close
+      @conn.close
+      @conn = nil
+    end
+
+    private
+
+    def connect_url
+      options.fetch(:connect_url)
+    end
+
+    def prepare_name
+      options.fetch(:prepare_name, self.class.to_s.downcase + "stmt")
+    end
+
+    def prepare_sql
+      options.fetch(:prepare_sql)
+    end
+
+    def columns
+      options.fetch(:columns)
+    end
+  end
+end

data/lib/kiba/plus/destination/pg_bulk.rb

@@ -0,0 +1,118 @@
+module Kiba::Plus::Destination
+  class PgBulk
+    attr_reader :options
+
+    def initialize(options = {})
+      @options = options
+      @options.assert_valid_keys(
+       :connect_url,
+       :input_file,
+       :table_name,
+       :columns,
+       :truncate,
+       :incremental,
+       :unique_by
+       )
+      @conn = PG.connect(connect_url)
+      if truncate
+        truncate_staging_table
+        truncate_target_table
+      end
+    end
+
+    def connect_url
+      options.fetch(:connect_url)
+    end
+
+    def table_name
+      options.fetch(:table_name)
+    end
+
+    def input_file
+      options.fetch(:input_file)
+    end
+
+    def columns
+      options.fetch(:columns)
+    end
+
+    def truncate
+      options.fetch(:truncate, false)
+    end
+
+    def incremental
+      options.fetch(:incremental, true)
+    end
+
+    def unique_by
+      options.fetch(:unique_by, :id)
+    end
+
+    def write(row)
+      # blank!
+    end
+
+    def staging_table_name
+      table_name + "_staging"
+    end
+
+    def create_staging_table
+      sql = "CREATE TABLE IF NOT EXISTS #{staging_table_name} (LIKE #{table_name} INCLUDING DEFAULTS INCLUDING CONSTRAINTS INCLUDING INDEXES)"
+      Kiba::Plus.logger.info sql
+      @conn.exec(sql)
+    end
+
+    def truncate_staging_table
+      truncate_sql = "TRUNCATE TABLE #{staging_table_name}"
+      Kiba::Plus.logger.info truncate_sql
+      @conn.exec(truncate_sql) rescue nil
+    end
+
+    def truncate_target_table
+      truncate_sql = "TRUNCATE TABLE #{table_name};"
+      Kiba::Plus.logger.info truncate_sql
+      @conn.exec(truncate_sql)
+    end
+
+    def copy_to_target_table
+      sql = "COPY #{table_name} (#{columns.join(', ')}) FROM '#{File.expand_path(input_file)}' WITH DELIMITER ',' NULL '\\N' CSV"
+      Kiba::Plus.logger.info sql
+      @conn.exec(sql)
+    end
+
+    def copy_to_staging_table
+      sql = "COPY #{staging_table_name} (#{columns.join(', ')}) FROM '#{File.expand_path(input_file)}' WITH DELIMITER ',' NULL '\\N' CSV"
+      Kiba::Plus.logger.info sql
+      @conn.exec(sql)
+    end
+
+    # TODO add where condition to speed up deleting.
+    def delete_before_insert
+      where = Array(unique_by).map{|x| ["#{staging_table_name}.#{x}", "#{table_name}.#{x}"].join(" = ") }.join(" AND ")
+      sql = "DELETE FROM #{table_name} USING #{staging_table_name} WHERE #{where}"
+      Kiba::Plus.logger.info sql
+      @conn.exec(sql)
+    end
+
+    def merge_to_target_table
+      sql = "INSERT INTO #{table_name} (SELECT * FROM #{staging_table_name})"
+      Kiba::Plus.logger.info sql
+      @conn.exec(sql)
+    end
+
+    def close
+      if incremental
+        truncate_staging_table
+        create_staging_table
+        copy_to_staging_table
+        delete_before_insert
+        merge_to_target_table
+        truncate_staging_table
+      else
+        copy_to_target_table
+      end
+      @conn.close
+      @conn = nil
+    end
+  end
+end

data/lib/kiba/plus/destination/pg_bulk2.rb

@@ -0,0 +1,68 @@
+require 'pg'
+require 'csv'
+module Kiba::Plus::Destination
+  class PgBulk2
+    attr_reader :options
+
+    def initialize(options = {})
+      @options = options
+      @options.assert_valid_keys(:table_name,
+                                 :columns,
+                                 :connect_url,
+                                 :truncate,
+                                 :incremental
+                                 )
+
+      @conn = PG.connect(connect_url)
+      truncate_sql = "TRUNCATE TABLE #{table_name};"
+      if truncate
+        Kiba::Plus.logger.info truncate_sql
+        @conn.exec(truncate_sql)
+      end
+      sql = "COPY #{table_name} (#{columns.join(', ')}) FROM STDIN WITH DELIMITER ',' NULL '\\N' CSV"
+      Kiba::Plus.logger.info sql
+      @res  = @conn.exec(sql)
+    end
+
+    def connect_url
+      options.fetch(:connect_url)
+    end
+
+    def table_name
+      options.fetch(:table_name)
+    end
+
+    def write(row)
+      begin
+        @conn.put_copy_data CSV.generate_line(row.values_at(*columns))
+      rescue Exception => err
+        errmsg = "%s while copy data: %s" % [ err.class.name, err.message ]
+        @conn.put_copy_end( errmsg )
+        Kiba::Plus.logger.error @conn.get_result
+        raise
+      end
+    end
+
+    def columns
+      options.fetch(:columns)
+    end
+
+    def truncate
+      options.fetch(:truncate, false)
+    end
+
+    def incremental
+      options.fetch(:incremental, true)
+    end
+
+    def close
+      @conn.put_copy_end
+      @conn.get_last_result
+    rescue
+      raise
+    ensure
+      @conn.close
+      @conn = nil
+    end
+  end
+end

data/lib/kiba/plus/helper.rb

@@ -0,0 +1,22 @@
+require 'uri'
+module Kiba
+  module Plus
+    module Helper
+      def connect_hash(url)
+        u = URI.parse(url)
+        {
+          host: u.host,
+          username: u.user,
+          password: u.password,
+          port: u.port,
+          database: u.path[1..-1]
+        }
+      end
+
+      def scheme(url)
+        u = URI.parse(url)
+        u.scheme
+      end
+    end
+  end
+end

data/lib/kiba/plus/job.rb

@@ -0,0 +1,149 @@
+require 'uri'
+
+module Kiba
+  require 'pg'
+  require 'mysql2'
+  require 'uri'
+
+  module Plus
+    class Job
+      include Kiba::Plus::Helper
+
+      attr_reader :options, :client
+      def initialize(options)
+        @options = options
+        @options.assert_valid_keys(:connect_url, :job_id, :job_name, :start_at, :completed_at)
+        url = URI.parse(connect_url)
+        if url.scheme =~ /mysql/i
+          @client = Mysql2::Client.new(connect_hash(connect_url))
+        elsif url.scheme =~ /postgres/i
+          @client = PG.connect(connect_url)
+        else
+          raise 'No Imp!'
+        end
+      end
+
+      def job_id
+        options.fetch(:job_id, nil)
+      end
+
+      def connect_url
+        options.fetch(:connect_url)
+      end
+
+      def job_name
+        options.fetch(:job_name)
+      end
+
+      def start_at
+        options.fetch(:start_at, Time.now)
+      end
+
+      def completed_at
+        options.fetch(:completed_at, Time.now)
+      end
+
+      def start
+        create_table
+        result = create_job
+        result.first["id"].to_i
+      end
+
+      def last_pull_at
+        sql = "SELECT MAX(created_at) AS last_pull_at FROM jobs WHERE status = 'completed' AND job_name = '#{job_name}'"
+        Kiba::Plus.logger.info sql
+        client.query(sql).first["last_pull_at"]
+      end
+
+      def complete
+        complete_job
+      end
+
+      private
+
+      def create_table
+        url = URI.parse(connect_url)
+        if url.scheme =~ /mysql/i
+          create_table_mysql
+        elsif url.scheme =~ /postgres/i
+          create_table_pg
+        else
+          raise 'No Imp!'
+        end
+      end
+
+      def create_job
+        if @client.is_a?(Mysql2::Client)
+          create_job_mysql
+        else
+          create_job_pg
+        end
+      end
+
+      def create_job_mysql
+        sql = <<-SQL
+          INSERT INTO jobs (
+            completed_at,
+            job_name,
+            created_at,
+            status) VALUES
+            (NULL, '#{job_name}', '#{start_at.to_s}', 'executing')
+        SQL
+        Kiba::Plus.logger.info sql
+        @client.query(sql)
+        returning_id_sql = "SELECT LAST_INSERT_ID() AS id"
+        Kiba::Plus.logger.info returning_id_sql
+        @client.query(returning_id_sql)
+      end
+
+      def create_job_pg
+        sql = <<-SQL
+          INSERT INTO jobs (
+            completed_at,
+            job_name,
+            created_at,
+            status) VALUES
+            (NULL, '#{job_name}', '#{start_at.to_s}', 'executing') RETURNING id
+        SQL
+        Kiba::Plus.logger.info sql
+        @client.query(sql)
+      end
+
+      def create_table_pg
+        sql = <<-SQL
+        CREATE TABLE IF NOT EXISTS jobs (
+          id SERIAL,
+          job_name varchar(255) NOT NULL,
+          created_at TIMESTAMP without time zone,
+          completed_at TIMESTAMP without time zone,
+          status varchar(255) DEFAULT NULL,
+          PRIMARY KEY (id)
+          )
+        SQL
+        Kiba::Plus.logger.info sql
+        @client.query(sql)
+      end
+
+      def create_table_mysql
+        sql = <<-SQL
+        CREATE TABLE IF NOT EXISTS jobs (
+          id integer(11) NOT NULL AUTO_INCREMENT,
+          job_name varchar(255) NOT NULL,
+          created_at datetime NOT NULL,
+          completed_at datetime DEFAULT NULL,
+          status varchar(255) DEFAULT NULL,
+          PRIMARY KEY (id)
+          ) AUTO_INCREMENT=1
+        SQL
+        Kiba::Plus.logger.info sql
+        @client.query(sql)
+      end
+
+      def complete_job
+        sql = "UPDATE jobs SET status = 'completed', completed_at = '#{completed_at.to_s}' WHERE id = #{job_id} AND job_name = '#{job_name}'"
+        Kiba::Plus.logger.info sql
+        @client.query(sql)
+      end
+    end
+  end
+end

data/lib/kiba/plus/logger.rb

@@ -0,0 +1,12 @@
+require 'logger'
+module Kiba
+  module Plus
+    def self.logger
+      @logger ||= Logger.new($stdout)
+    end
+
+    def self.logger=(logger)
+      @logger = logger
+    end
+  end
+end

data/lib/kiba/plus/source/mysql.rb

@@ -0,0 +1,50 @@
+require 'mysql2'
+require 'uri'
+
+module Kiba
+  module Plus::Source
+    class Mysql
+      include Kiba::Plus::Helper
+      attr_reader :options, :client
+
+      def initialize(options = {})
+        @options = options
+        @options.assert_valid_keys(
+          :query,
+          :output,
+          :last_pull_at,
+          :incremental,
+          :connect_url
+          )
+        @client = Mysql2::Client.new(connect_hash(connect_url))
+      end
+
+      def each
+        results = client.query(query, as: :hash, symbolize_keys: true, stream: true)
+        results.each do |row|
+          yield(row)
+        end
+      end
+
+      def query
+        options.fetch(:query)
+      end
+
+      def output
+        options.fetch(:output)
+      end
+
+      def last_pull_at
+        options[:last_pull_at]
+      end
+
+      def incremental
+        options.fetch(:incremental, true)
+      end
+
+      def connect_url
+        options.fetch(:connect_url)
+      end
+    end
+  end
+end