postgres_upsert 2.0.0 → 5.1.0

data/db/migrate/20150214192135_create_test_tables.rb

@@ -0,0 +1,24 @@
+class CreateTestTables < ActiveRecord::Migration[6.0]
+  def change
+    create_table :test_models do |t|
+      t.string :data
+      t.timestamps
+    end
+
+    create_table :test_model_copies do |t|
+      t.string :data
+      t.timestamps
+    end
+
+    create_table :three_columns do |t|
+      t.string :data
+      t.string :extra
+      t.timestamps
+    end
+
+    create_table :reserved_word_models do |t|
+      t.string :select
+      t.string :group
+    end
+  end
+end

data/db/migrate/20150710162236_create_composite_models_table.rb

@@ -0,0 +1,9 @@
+class CreateCompositeModelsTable < ActiveRecord::Migration[6.0]
+  def change
+    create_table :composite_key_models do |t|
+      t.integer :comp_key_1
+      t.integer :comp_key_2
+      t.string :data
+    end
+  end
+end

data/db/schema.rb

@@ -0,0 +1,48 @@
+# This file is auto-generated from the current state of the database. Instead
+# of editing this file, please use the migrations feature of Active Record to
+# incrementally modify your database, and then regenerate this schema definition.
+#
+# This file is the source Rails uses to define your schema when running `rails
+# db:schema:load`. When creating a new database, `rails db:schema:load` tends to
+# be faster and is potentially less error prone than running all of your
+# migrations from scratch. Old migrations may fail to apply correctly if those
+# migrations use external dependencies or application code.
+#
+# It's strongly recommended that you check this file into your version control system.
+
+ActiveRecord::Schema.define(version: 2015_07_10_162236) do
+
+  # These are extensions that must be enabled in order to support this database
+  enable_extension "plpgsql"
+
+  create_table "composite_key_models", force: :cascade do |t|
+    t.integer "comp_key_1"
+    t.integer "comp_key_2"
+    t.string "data"
+  end
+
+  create_table "reserved_word_models", force: :cascade do |t|
+    t.string "select"
+    t.string "group"
+  end
+
+  create_table "test_model_copies", force: :cascade do |t|
+    t.string "data"
+    t.datetime "created_at", precision: 6, null: false
+    t.datetime "updated_at", precision: 6, null: false
+  end
+
+  create_table "test_models", force: :cascade do |t|
+    t.string "data"
+    t.datetime "created_at", precision: 6, null: false
+    t.datetime "updated_at", precision: 6, null: false
+  end
+
+  create_table "three_columns", force: :cascade do |t|
+    t.string "data"
+    t.string "extra"
+    t.datetime "created_at", precision: 6, null: false
+    t.datetime "updated_at", precision: 6, null: false
+  end
+
+end

data/db/seeds.rb

@@ -0,0 +1,7 @@
+# This file should contain all the record creation needed to seed the database with its default values.
+# The data can then be loaded with the rake db:seed (or created alongside the db with db:setup).
+#
+# Examples:
+#
+#   cities = City.create([{ name: 'Chicago' }, { name: 'Copenhagen' }])
+#   Mayor.create(name: 'Emanuel', city: cities.first)

data/lib/postgres_upsert.rb

@@ -1,14 +1,46 @@
 require 'rubygems'
 require 'active_record'
-require 'postgres_upsert/active_record'
+require 'postgres_upsert/read_adapters/active_record_adapter'
+require 'postgres_upsert/read_adapters/file_adapter'
+require 'postgres_upsert/read_adapters/io_adapter'
+require 'postgres_upsert/write_adapters/active_record_adapter'
+require 'postgres_upsert/write_adapters/table_adapter'
 require 'postgres_upsert/writer'
+require 'postgres_upsert/table_writer'
+require 'postgres_upsert/model_to_model_adapter'
+require 'postgres_upsert/result'
 require 'rails'
 
-class PostgresCopy < Rails::Railtie
+module PostgresUpsert
+  class << self
+    def write(destination, source, options = {})
+      read_adapter = read_adapter(source).new(source, options)
+      write_adapter = write_adapter(destination).new(destination, options)
+      Writer.new(destination, write_adapter, read_adapter, options).write
+    end
+
+    def read_adapter(source)
+      if [StringIO, File].include?(source.class)
+        ReadAdapters::IOAdapter
+      elsif [String].include?(source.class)
+        ReadAdapters::FileAdapter
+      elsif source < ActiveRecord::Base
+        ReadAdapters::ActiveRecordAdapter
+      else
+        raise "Source must be a Filename string, StringIO of data, or a ActiveRecord Class."
+      end
+    end
 
-  initializer 'postgres_upsert' do
-    ActiveSupport.on_load :active_record do
-      require "postgres_upsert/active_record"
+    def write_adapter(destination)
+      if [String].include?(destination.class)
+        WriteAdapters::TableAdapter
+      elsif destination <= ActiveRecord::Base
+        WriteAdapters::ActiveRecordAdapter
+      # elsif source < ActiveRecord::Base && destination < ActiveRecord::Base
+        #ModelToModelAdapter
+      else
+        raise "Destination must be an ActiveRecord class or a table name string"
+      end
     end
   end
-end
+end

data/lib/postgres_upsert/model_to_model_adapter.rb

@@ -0,0 +1,37 @@
+require 'csv'
+
+module PostgresUpsert
+  class ModelToModelAdapter
+    def initialize(destination_model, source_model, options = {})
+      @destination_model = destination_model
+      @source_model = source_model
+      @options = options
+    end
+
+    def write
+        source_table = @source_model.table_name
+        source_conn = @source_model.connection.raw_connection
+
+        to_stdout_sql = "COPY #{source_table} TO STDOUT"
+      
+        csv_string = CSV.generate do |csv|
+          csv << @source_model.column_names # CSV header row
+          source_conn.copy_data(to_stdout_sql) do
+            while (line = source_conn.get_copy_data) do
+              csv << line.split("\t")
+            end
+          end
+        end
+        io = StringIO.new(csv_string)
+        Writer.new(@destination_model, io, @options).write
+    end
+  
+  private
+
+    def get_columns
+      # columns_list = @options[:columns] 
+      # columns_list ||= 
+      @source.column_names
+    end
+  end
+end

data/lib/postgres_upsert/read_adapters/active_record_adapter.rb

@@ -0,0 +1,37 @@
+module PostgresUpsert
+  module ReadAdapters
+    class ActiveRecordAdapter
+      def initialize(source, options)
+        @options = sanitize_options(options)
+        @source = source
+      end
+
+      def sanitize_options(options)
+        options.slice(
+          :columns, :map, :unique_key
+        )
+      end
+
+      def continuous_write_enabled
+        false
+      end
+
+      def gets(&block)
+        batch_size = 1_000
+        line = ""
+        conn = @source.connection.raw_connection
+
+        conn.copy_data("COPY #{@source.table_name} TO STDOUT") do
+          while (line_read = conn.get_copy_data) do
+            line << line_read
+          end
+        end
+        yield line
+      end
+
+      def columns
+        @source.column_names
+      end
+    end
+  end
+end

data/lib/postgres_upsert/read_adapters/file_adapter.rb

@@ -0,0 +1,42 @@
+module PostgresUpsert
+  module ReadAdapters
+    class FileAdapter
+      def initialize(source, options)
+        @options = sanitize_options(options)
+        @source = File.open(source, 'r')
+      end
+
+      def sanitize_options(options)
+        options.slice(
+          :delimiter, :header, :columns, :map, :unique_key
+        ).reverse_merge(
+          header: true,
+          delimiter: ',',
+        )
+      end
+
+      def continuous_write_enabled
+        true
+      end
+
+      def gets
+        @source.gets
+      end
+
+      def columns
+        @columns ||= begin
+          columns_list = @options[:columns] ? @options[:columns].map(&:to_s) : []
+          if @options[:header]
+            # if header is present, we need to strip it from io, whether we use it for the columns list or not.
+            line = gets
+            if columns_list.empty?
+              columns_list = line.strip.split(@options[:delimiter])
+            end
+          end
+          columns_list = columns_list.map { |c| @options[:map][c.to_s] } if @options[:map]
+          columns_list
+        end
+      end
+    end
+  end
+end

data/lib/postgres_upsert/read_adapters/io_adapter.rb

@@ -0,0 +1,42 @@
+module PostgresUpsert
+  module ReadAdapters
+    class IOAdapter
+      def initialize(source, options)
+        @options = sanitize_options(options)
+        @source = source
+      end
+
+      def sanitize_options(options)
+        options.slice(
+          :delimiter, :header, :columns, :map, :unique_key
+        ).reverse_merge(
+          header: true,
+          delimiter: ',',
+        )
+      end
+      
+      def continuous_write_enabled
+        true
+      end
+
+      def gets
+        @source.gets
+      end
+
+      def columns
+        @columns ||= begin
+          columns_list = @options[:columns] ? @options[:columns].map(&:to_s) : []
+          if @options[:header]
+            # if header is present, we need to strip it from io, whether we use it for the columns list or not.
+            line = gets
+            if columns_list.empty?
+              columns_list = line.strip.split(@options[:delimiter])
+            end
+          end
+          columns_list = columns_list.map { |c| @options[:map][c.to_s] } if @options[:map]
+          columns_list
+        end
+      end
+    end
+  end
+end

data/lib/postgres_upsert/result.rb

@@ -0,0 +1,23 @@
+module PostgresUpsert
+  class Result
+    attr_reader :inserted, :updated
+
+    def initialize(insert_result, update_result, copy_result)
+      @inserted = insert_result ? insert_result.cmd_tuples : 0
+      @updated = update_result ? update_result.cmd_tuples : 0 
+      @copied = copy_result ? copy_result.cmd_tuples : 0
+    end
+
+    def changed_rows
+      @inserted + @updated
+    end
+
+    def copied_rows
+      @copied
+    end
+
+    def updated_rows
+      @updated
+    end
+  end
+end

data/lib/postgres_upsert/table_writer.rb

@@ -0,0 +1,48 @@
+module PostgresUpsert
+  # alternate version of PostgresUpsert::Writer which does not rely on AR table information.  
+  # We can use this model to upsert data into views, or tables not associated to rails models
+  class TableWriter < Writer
+    def initialize(table_name, source, options = {})
+      @table_name = table_name
+      super(nil, source, options)
+    end
+
+  private
+
+    def database_connection
+      ActiveRecord::Base.connection
+    end
+
+    def primary_key
+      @primary_key ||= begin
+        query = <<-SELECT_KEY
+          SELECT
+            pg_attribute.attname,
+            format_type(pg_attribute.atttypid, pg_attribute.atttypmod)
+          FROM pg_index, pg_class, pg_attribute
+          WHERE
+            pg_class.oid = '#{@table_name}'::regclass AND
+            indrelid = pg_class.oid AND
+            pg_attribute.attrelid = pg_class.oid AND
+            pg_attribute.attnum = any(pg_index.indkey)
+          AND indisprimary
+        SELECT_KEY
+
+        pg_result = ActiveRecord::Base.connection.execute query
+        pg_result.each { |row| return row['attname'] }
+      end
+    end
+
+    def destination_columns
+      @column_names ||= begin
+        query = "SELECT * FROM information_schema.columns WHERE TABLE_NAME = '#{@table_name}'"
+        pg_result = ActiveRecord::Base.connection.execute query
+        pg_result.map { |row| row['column_name'] }
+      end
+    end
+
+    def quoted_table_name
+      @quoted_table_name ||= ActiveRecord::Base.connection.quote_table_name(@table_name)
+    end
+  end
+end

data/lib/postgres_upsert/write_adapters/active_record_adapter.rb

@@ -0,0 +1,36 @@
+module PostgresUpsert
+  module WriteAdapters
+    class ActiveRecordAdapter
+      def initialize(destination, options)
+        @destination = destination
+        @options = sanitize_options(options)
+        
+      end
+
+      def sanitize_options(options)
+        options.slice(
+          :delimiter, :unique_key
+        ).reverse_merge(
+          delimiter: ',',
+          unique_key: [primary_key],
+        )
+      end
+
+      def database_connection
+        @destination.connection
+      end
+      
+      def primary_key
+        @destination.primary_key
+      end
+
+      def column_names
+        @destination.column_names
+      end
+
+      def quoted_table_name
+        @destination.quoted_table_name
+      end
+    end
+  end
+end

data/lib/postgres_upsert/write_adapters/table_adapter.rb

@@ -0,0 +1,56 @@
+module PostgresUpsert
+  module WriteAdapters
+    class TableAdapter
+      def initialize(destination, options)
+        @destination = destination
+        @options = sanitize_options(options)
+      end
+
+      def sanitize_options(options)
+        options.slice(
+          :delimiter, :unique_key
+        ).reverse_merge(
+          delimiter: ',',
+          unique_key: [primary_key],
+        )
+      end
+
+      def database_connection
+        ActiveRecord::Base.connection
+      end
+
+      def primary_key
+        @primary_key ||= begin
+          query = <<-SELECT_KEY
+            SELECT
+              pg_attribute.attname,
+              format_type(pg_attribute.atttypid, pg_attribute.atttypmod)
+            FROM pg_index, pg_class, pg_attribute
+            WHERE
+              pg_class.oid = '#{@destination}'::regclass AND
+              indrelid = pg_class.oid AND
+              pg_attribute.attrelid = pg_class.oid AND
+              pg_attribute.attnum = any(pg_index.indkey)
+            AND indisprimary
+          SELECT_KEY
+  
+          pg_result = ActiveRecord::Base.connection.execute query
+          pg_result.each { |row| return row['attname'] }
+        end
+      end
+
+      def column_names
+        @column_names ||= begin
+          query = "SELECT * FROM information_schema.columns WHERE TABLE_NAME = '#{@destination}'"
+          pg_result = ActiveRecord::Base.connection.execute query
+          pg_result.map { |row| row['column_name'] }
+        end
+      end
+  
+      def quoted_table_name
+        @quoted_table_name ||= database_connection.quote_table_name(@destination)
+      end
+
+    end
+  end
+end