postgres_upsert 2.0.0 → 5.1.0

data/lib/postgres_upsert/writer.rb

@@ -1,60 +1,89 @@
 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 => ",", 
-        :format => :csv, 
-        :header => true, 
-        :key_column => @klass.primary_key,
-        :update_only => false})
-      @source = source.instance_of?(String) ? File.open(source, 'r') : source
-      @columns_list = get_columns
-      generate_temp_table_name
+        delimiter: ',',
+        header: true,
+        unique_key: [primary_key],
+        update_only: false
+      })
+      @source = source
+      @options[:unique_key] = Array.wrap(@options[:unique_key])
+
     end
 
     def write
-      if @columns_list.empty? 
-        raise "Either the :columns option or :header => true are required"
+      validate_options
+
+      
+      create_temp_table
+
+      if @source.continuous_write_enabled
+        write_continuous
+      else
+        write_batched
       end
 
-      csv_options = @options[:format] == :binary ? "BINARY" : "DELIMITER '#{@options[:delimiter]}' CSV"
+      upsert_from_temp_table
+      drop_temp_table
 
-      copy_table = @temp_table_name
-      destination_table = get_table_name
+      summarize_results
+    end
 
-      columns_string = columns_string_for_copy
-      create_temp_table
+  private
 
-      ActiveRecord::Base.connection.raw_connection.copy_data %{COPY #{copy_table} #{columns_string} FROM STDIN #{csv_options}} do
+    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?
 
-        while line = read_input_line do
-          next if line.strip.size == 0
-          ActiveRecord::Base.connection.raw_connection.put_copy_data line
+          database_connection.raw_connection.put_copy_data line
         end
       end
+    end
 
-      if destination_table
-        upsert_from_temp_table
-        drop_temp_table
+    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
 
-  private
+    def database_connection
+      @destination.database_connection
+    end
 
-    def get_columns
-      columns_list = @options[:columns] || []
-      if @options[:format] != :binary && @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
+    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
-      columns_list = columns_list.map{|c| @options[:map][c.to_s] } if @options[:map]
-      return columns_list
+
+      result
+    end
+
+    def primary_key
+      @destination.primary_key
+    end
+
+    def destination_columns
+      @destination.column_names
+    end
+
+    def quoted_table_name
+      @destination.quoted_table_name
+    end
+
+    def source_columns
+      @source.columns
     end
 
     def columns_string_for_copy
@@ -63,60 +92,50 @@ module PostgresUpsert
     end
 
     def columns_string_for_select
-      columns = @columns_list.clone
-      columns << "created_at" if @klass.column_names.include?("created_at")
-      columns << "updated_at" if @klass.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 @klass.column_names.include?("created_at")
-      columns << "updated_at" if @klass.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 @klass.column_names.include?("created_at")
-      str << ",'#{DateTime.now.utc}'" if @klass.column_names.include?("updated_at")
+      str << ",'#{DateTime.now.utc}'" if inject_create_timestamp?
+      str << ",'#{DateTime.now.utc}'" if inject_update_timestamp?
       str
     end
 
-    def select_string_for_create
-      columns = @columns_list.map(&:to_sym)
-      columns << @options[:key_column].to_sym unless columns.include?(@options[:key_column].to_sym)
-      get_columns_string(columns)
+    def inject_create_timestamp?
+      destination_columns.include?('created_at') &&  !source_columns.include?('created_at')
     end
 
-    def get_columns_string(columns = nil)
-      columns ||= @columns_list
-      columns.size > 0 ? "\"#{columns.join('","')}\"" : ""
+    def inject_update_timestamp?
+      destination_columns.include?('updated_at') &&  !source_columns.include?('updated_at')
     end
 
-    def get_table_name
-      if @options[:table]
-        connection.quote_table_name(@options[:table])
-      else
-        @klass.quoted_table_name
+    def select_string_for_create
+      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
+      get_columns_string(columns)
     end
 
-    def generate_temp_table_name
-      @temp_table_name = "#{@klass.table_name}_temp_#{rand(1000)}"
+    def get_columns_string(columns = nil)
+      columns ||= source_columns
+      !columns.empty? ? "\"#{columns.join('","')}\"" : ''
     end
 
-    def read_input_line
-      if @options[:format] == :binary
-        begin
-          return @source.readpartial(10240)
-        rescue EOFError
-        end
-      else
-        line = @source.gets
-        return line
-      end
+    def generate_temp_table_name
+      @temp_table_name ||= "#{@table_name}_temp_#{rand(1000)}"
     end
 
     def upsert_from_temp_table
@@ -125,55 +144,74 @@ module PostgresUpsert
     end
 
     def update_from_temp_table
-      ActiveRecord::Base.connection.execute <<-SQL
-        UPDATE #{get_table_name} AS d
+      @update_result = database_connection.execute <<-SQL
+        UPDATE #{quoted_table_name} AS d
           #{update_set_clause}
           FROM #{@temp_table_name} as t
-          WHERE t.#{@options[:key_column]} = d.#{@options[:key_column]}
-          AND d.#{@options[:key_column]} IS NOT NULL;
+          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 @klass.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
       columns_string = columns_string_for_insert
       select_string = select_string_for_insert
-      ActiveRecord::Base.connection.execute <<-SQL
-        INSERT INTO #{get_table_name} (#{columns_string})
+      @insert_result = database_connection.execute <<-SQL
+        INSERT INTO #{quoted_table_name} (#{columns_string})
           SELECT #{select_string}
           FROM #{@temp_table_name} as t
-          WHERE NOT EXISTS 
-            (SELECT 1 
-                  FROM #{get_table_name} as d 
-                  WHERE d.#{@options[:key_column]} = t.#{@options[:key_column]})
-          AND t.#{@options[:key_column]} IS NOT NULL;
+          WHERE NOT EXISTS
+            (SELECT 1
+                  FROM #{quoted_table_name} as 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 ')
+    end
+
+    def unique_key_present(source)
+      @options[:unique_key].map { |field| "#{source}.#{field} IS NOT NULL" }.join(' AND ')
+    end
+
     def create_temp_table
-      columns_string = select_string_for_create
-      ActiveRecord::Base.connection.execute <<-SQL
+      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 #{get_table_name} WHERE 0 = 1;
+        CREATE TEMP TABLE #{@temp_table_name}
+          AS SELECT #{select_string_for_create} FROM #{quoted_table_name} WHERE 0 = 1;
       SQL
     end
 
+    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 source_columns.include?(key_component.to_s)
+          raise "Expected column '#{key_component}' was not found in source"
+        end
+      end
+    end
+
     def drop_temp_table
-      ActiveRecord::Base.connection.execute <<-SQL
-        DROP TABLE #{@temp_table_name} 
+      database_connection.execute <<-SQL
+        DROP TABLE #{@temp_table_name}
       SQL
     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 = "2.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")
@@ -22,12 +23,14 @@ Gem::Specification.new do |s|
   s.require_paths = ["lib"]
   s.summary = "A rubygem that integrates with ActiveRecord to insert/update large data sets into the database efficiently"
 
-  s.add_dependency "pg", '~> 0.17.0'
+  s.add_dependency "pg", '>= 0.17.0'
   s.add_dependency "activerecord", '>= 3.0.0'
   s.add_dependency "rails", '>= 3.0.0'
   s.add_development_dependency "bundler"
-  s.add_development_dependency "rdoc"
   s.add_development_dependency "pry-rails"
-  s.add_development_dependency "rspec", "~> 2.12"
+  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

@@ -0,0 +1,50 @@
+require File.expand_path(File.dirname(__FILE__) + '/spec_helper')
+
+describe "pg_upsert from file with CSV format" do
+
+  before do
+    DateTime.stub_chain(:now, :utc).and_return (DateTime.parse("2012-01-01").utc)
+  end
+
+  def timestamp
+    DateTime.now.utc
+  end
+
+  context 'composite_key_support' do
+    it 'inserts records if the passed match composite key doesnt exist' do
+      file = File.open(File.expand_path('spec/fixtures/composite_key_with_header.csv'), 'r')
+
+      PostgresUpsert.write(CompositeKeyModel, file, :unique_key => ["comp_key_1", "comp_key_2"])
+      expect(
+        CompositeKeyModel.last.attributes
+      ).to include("data" => "test data 2")
+    end
+
+    it 'updates records if the passed composite key exists' do
+      file = File.open(File.expand_path('spec/fixtures/composite_key_with_header.csv'), 'r')
+      existing = CompositeKeyModel.create(comp_key_1: 2, comp_key_2:3, data: "old stuff")
+
+      PostgresUpsert.write(CompositeKeyModel, file, :unique_key => ["comp_key_1", "comp_key_2"])
+
+      expect(
+        CompositeKeyModel.find_by({comp_key_1: 2, comp_key_2:3}).attributes
+      ).to include("data" => "test data 2")
+
+      expect(
+        CompositeKeyModel.find_by({comp_key_1: 1, comp_key_2:2}).attributes
+      ).to include("data" => "test data 1")
+    end
+
+    it 'fails if composite keys are not unique.' do
+      file = File.open(File.expand_path('spec/fixtures/composite_nonkey_with_header.csv'), 'r')
+      existing = CompositeKeyModel.create(comp_key_1: 1, comp_key_2:2, data: "old stuff")
+
+       expect{
+        PostgresUpsert.write(CompositeKeyModel, file, :unique_key => ["comp_key_1", "comp_key_2"])
+      }.to raise_error(/Check to make sure your key is unique/)
+    end
+    
+  end
+
+
+end

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/composite_key_model.rb

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

data/spec/fixtures/composite_key_with_header.csv

@@ -0,0 +1,3 @@
+comp_key_1,comp_key_2,data
+1,2,test data 1
+2,3,test data 2

data/spec/fixtures/composite_nonkey_with_header.csv

@@ -0,0 +1,3 @@
+comp_key_1,comp_key_2,data
+1,2,test data 1
+1,2,test data 4

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