postgres_upsert 5.0.0 → 5.1.0

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

data/lib/postgres_upsert/writer.rb

@@ -1,36 +1,31 @@
 module PostgresUpsert
   class Writer
 
-    def initialize(klass, source, options = {})
+    def initialize(klass, destination, source, options = {})
       @klass = klass
+      @destination = destination
+      @source = source
       @options = options.reverse_merge({
-        :delimiter => ",",
-        :header => true,
-        :unique_key => [primary_key],
-        :update_only => false})
+        delimiter: ',',
+        header: true,
+        unique_key: [primary_key],
+        update_only: false
+      })
+      @source = source
       @options[:unique_key] = Array.wrap(@options[:unique_key])
-      @source = source.instance_of?(String) ? File.open(source, 'r') : source
-      @columns_list = get_columns
-      generate_temp_table_name
+
     end
 
     def write
-      if @columns_list.empty?
-        raise "Either the :columns option or :header => true are required"
-      end
-
-      csv_options = "DELIMITER '#{@options[:delimiter]}' CSV"
+      validate_options
 
-      copy_table = @temp_table_name
-      columns_string = columns_string_for_copy
+      
       create_temp_table
 
-      @copy_result = database_connection.raw_connection.copy_data %{COPY #{copy_table} #{columns_string} FROM STDIN #{csv_options}} do
-
-        while line = @source.gets do
-          next if line.strip.size == 0
-          database_connection.raw_connection.put_copy_data line
-        end
+      if @source.continuous_write_enabled
+        write_continuous
+      else
+        write_batched
       end
 
       upsert_from_temp_table
@@ -41,44 +36,54 @@ module PostgresUpsert
 
   private
 
+    def write_continuous
+      csv_options = "DELIMITER '#{@options[:delimiter]}' CSV"
+      @copy_result = database_connection.raw_connection.copy_data %{COPY #{@temp_table_name} #{columns_string_for_copy} FROM STDIN #{csv_options}} do
+        while (line = @source.gets)
+          next if line.strip.empty?
+
+          database_connection.raw_connection.put_copy_data line
+        end
+      end
+    end
+
+    def write_batched
+      @source.gets do |line|
+        @copy_result = database_connection.raw_connection.copy_data %{COPY #{@temp_table_name} #{columns_string_for_copy} FROM STDIN} do
+          database_connection.raw_connection.put_copy_data line
+        end
+      end
+    end
+
     def database_connection
-      @klass.connection
+      @destination.database_connection
     end
 
     def summarize_results
       result = PostgresUpsert::Result.new(@insert_result, @update_result, @copy_result)
       expected_rows = @options[:update_only] ? result.updated_rows : result.copied_rows
-      
+
       if result.changed_rows != expected_rows
         raise "#{expected_rows} rows were copied, but #{result.changed_rows} were upserted to destination table.  Check to make sure your key is unique."
       end
 
-      return result
+      result
     end
 
     def primary_key
-      @klass.primary_key
+      @destination.primary_key
     end
 
-    def column_names
-      @klass.column_names
+    def destination_columns
+      @destination.column_names
     end
 
     def quoted_table_name
-      @klass.quoted_table_name
+      @destination.quoted_table_name
     end
 
-    def get_columns
-      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 = @source.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]
-      return columns_list
+    def source_columns
+      @source.columns
     end
 
     def columns_string_for_copy
@@ -87,29 +92,37 @@ module PostgresUpsert
     end
 
     def columns_string_for_select
-      columns = @columns_list.clone
-      columns << "created_at" if column_names.include?("created_at")
-      columns << "updated_at" if column_names.include?("updated_at")
-      str = get_columns_string(columns)
+      columns = source_columns.clone
+      columns << 'created_at' if inject_create_timestamp?
+      columns << 'updated_at' if inject_update_timestamp?
+      get_columns_string(columns)
     end
 
     def columns_string_for_insert
-      columns = @columns_list.clone
-      columns << "created_at" if column_names.include?("created_at")
-      columns << "updated_at" if column_names.include?("updated_at")
-      str = get_columns_string(columns)
+      columns = source_columns.clone
+      columns << 'created_at' if inject_create_timestamp?
+      columns << 'updated_at' if inject_update_timestamp?
+      get_columns_string(columns)
     end
 
     def select_string_for_insert
-      columns = @columns_list.clone
+      columns = source_columns.clone
       str = get_columns_string(columns)
-      str << ",'#{DateTime.now.utc}'" if column_names.include?("created_at")
-      str << ",'#{DateTime.now.utc}'" if column_names.include?("updated_at")
+      str << ",'#{DateTime.now.utc}'" if inject_create_timestamp?
+      str << ",'#{DateTime.now.utc}'" if inject_update_timestamp?
       str
     end
 
+    def inject_create_timestamp?
+      destination_columns.include?('created_at') &&  !source_columns.include?('created_at')
+    end
+
+    def inject_update_timestamp?
+      destination_columns.include?('updated_at') &&  !source_columns.include?('updated_at')
+    end
+
     def select_string_for_create
-      columns = @columns_list.map(&:to_sym)
+      columns = source_columns.map(&:to_sym)
       @options[:unique_key].each do |key_component|
         columns << key_component.to_sym unless columns.include?(key_component.to_sym)
       end
@@ -117,12 +130,12 @@ module PostgresUpsert
     end
 
     def get_columns_string(columns = nil)
-      columns ||= @columns_list
-      columns.size > 0 ? "\"#{columns.join('","')}\"" : ""
+      columns ||= source_columns
+      !columns.empty? ? "\"#{columns.join('","')}\"" : ''
     end
 
     def generate_temp_table_name
-      @temp_table_name = "#{@table_name}_temp_#{rand(1000)}"
+      @temp_table_name ||= "#{@table_name}_temp_#{rand(1000)}"
     end
 
     def upsert_from_temp_table
@@ -135,17 +148,19 @@ module PostgresUpsert
         UPDATE #{quoted_table_name} AS d
           #{update_set_clause}
           FROM #{@temp_table_name} as t
-          WHERE #{unique_key_select("t", "d")}
-          AND #{unique_key_present("d")}
+          WHERE #{unique_key_select('t', 'd')}
+          AND #{unique_key_present('d')}
       SQL
     end
 
     def update_set_clause
-      command = @columns_list.map do |col|
+      command = source_columns.map do |col|
         "\"#{col}\" = t.\"#{col}\""
       end
-      command << "\"updated_at\" = '#{DateTime.now.utc}'" if column_names.include?("updated_at")
-      "SET #{command.join(',')}"
+      unless source_columns.include?('updated_at')
+        command << "\"updated_at\" = '#{DateTime.now.utc}'" if destination_columns.include?('updated_at')
+      end
+        "SET #{command.join(',')}"
     end
 
     def insert_from_temp_table
@@ -158,35 +173,38 @@ module PostgresUpsert
           WHERE NOT EXISTS
             (SELECT 1
                   FROM #{quoted_table_name} as d
-                  WHERE #{unique_key_select("t", "d")});
+                  WHERE #{unique_key_select('t', 'd')});
       SQL
     end
 
     def unique_key_select(source, dest)
-      @options[:unique_key].map {|field| "#{source}.#{field} = #{dest}.#{field}"}.join(' AND ')
+      @options[:unique_key].map { |field| "#{source}.#{field} = #{dest}.#{field}" }.join(' AND ')
     end
 
     def unique_key_present(source)
-      @options[:unique_key].map {|field| "#{source}.#{field} IS NOT NULL"}.join(' AND ')
+      @options[:unique_key].map { |field| "#{source}.#{field} IS NOT NULL" }.join(' AND ')
     end
 
     def create_temp_table
-      columns_string = select_string_for_create
-      verify_temp_has_key
+      generate_temp_table_name
       database_connection.execute <<-SQL
         SET client_min_messages=WARNING;
         DROP TABLE IF EXISTS #{@temp_table_name};
 
         CREATE TEMP TABLE #{@temp_table_name}
-          AS SELECT #{columns_string} FROM #{quoted_table_name} WHERE 0 = 1;
+          AS SELECT #{select_string_for_create} FROM #{quoted_table_name} WHERE 0 = 1;
       SQL
     end
 
-    def verify_temp_has_key
+    def validate_options
+      if source_columns.empty?
+        raise 'Either the :columns option or :header => true are required'
+      end
+
       @options[:unique_key].each do |key_component|
-        unless @columns_list.include?(key_component.to_s)
-          raise "Expected a unique column '#{key_component}' but the source data does not include this column.  Update the :columns list or explicitly set the unique_key option.}"
-       end
+        unless source_columns.include?(key_component.to_s)
+          raise "Expected column '#{key_component}' was not found in source"
+        end
       end
     end
 

data/postgres_upsert.gemspec

@@ -5,7 +5,7 @@ $:.unshift lib unless $:.include?(lib)
 
 Gem::Specification.new do |s|
   s.name = "postgres_upsert"
-  s.version = "5.0.0"
+  s.version = "5.1.0"
 
   s.platform    = Gem::Platform::RUBY
   s.required_ruby_version     = ">= 1.8.7"
@@ -13,6 +13,7 @@ Gem::Specification.new do |s|
   s.date = "2014-09-12"
   s.description = "Uses Postgres's powerful COPY command to upsert large sets of data into ActiveRecord tables"
   s.email = "thestevemitchell@gmail.com"
+  s.license = "MIT"
   git_files            = `git ls-files`.split("\n") rescue ''
   s.files              = git_files
   s.test_files         = `git ls-files -- {test,spec,features}/*`.split("\n")
@@ -27,7 +28,9 @@ Gem::Specification.new do |s|
   s.add_dependency "rails", '>= 3.0.0'
   s.add_development_dependency "bundler"
   s.add_development_dependency "pry-rails"
-  s.add_development_dependency "rspec", "~> 2.12"
-  s.add_development_dependency "rspec-rails", "~> 2.0"
+  s.add_development_dependency "pry-nav"
+  s.add_development_dependency "rspec-rails", ">= 3.9"
+  s.add_development_dependency "database_cleaner-active_record"
+  s.add_development_dependency "rubocop"
 end
 

data/spec/composite_key_spec.rb

@@ -1,12 +1,6 @@
 require File.expand_path(File.dirname(__FILE__) + '/spec_helper')
 
 describe "pg_upsert from file with CSV format" do
-  before(:each) do
-    ActiveRecord::Base.connection.execute %{
-      TRUNCATE TABLE composite_key_models;
-      SELECT setval('composite_key_models_id_seq', 1, false);
-    }
-  end
 
   before do
     DateTime.stub_chain(:now, :utc).and_return (DateTime.parse("2012-01-01").utc)

data/spec/fixtures/comma_with_header_duplicate.csv

@@ -0,0 +1,3 @@
+id,data
+1,test data 1
+1,test data dupe

data/spec/fixtures/test_model_copy.rb

@@ -0,0 +1,4 @@
+require 'postgres_upsert'
+
+class TestModelCopy < ActiveRecord::Base
+end

data/spec/from_table_spec.rb

@@ -0,0 +1,40 @@
+require "rails_helper"
+
+describe PostgresUpsert do
+  context "when passing ActiveRecord class as destination" do
+    context "when passing ActiveRecord clas as Source" do
+      let(:original_created_at) {5.days.ago.utc}
+
+      before(:each) do
+        TestModel.create(data: "From the before time, in the long long ago", :created_at => original_created_at)
+      end
+
+      it "copies the source to destination" do
+          PostgresUpsert.write TestModelCopy, TestModel
+          expect(
+            TestModelCopy.first.attributes
+          ).to eq(TestModelCopy.first.attributes)
+      end
+
+      context "with a large table" do
+        before do
+          csv_string = CSV.generate do |csv|
+            csv << %w(id data)    # CSV header row
+            (1..100_000).each do |n|
+              csv << ["#{n}", "data about #{n}"]
+            end
+          end
+          io = StringIO.new(csv_string)
+          PostgresUpsert.write TestModel, io
+        end
+
+        it "moves like the poop through a goose" do
+          expect{
+            PostgresUpsert.write TestModelCopy, TestModel
+          }.to change{TestModelCopy.count}.by(100_000)
+
+        end
+      end
+    end
+  end
+end