postgres_upsert 3.0.0 → 4.0.0

data/lib/postgres_upsert.rb

@@ -1,14 +1,17 @@
 require 'rubygems'
 require 'active_record'
-require 'postgres_upsert/active_record'
 require 'postgres_upsert/writer'
+require 'postgres_upsert/table_writer'
+require 'postgres_upsert/result'
 require 'rails'
 
-class PostgresCopy < Rails::Railtie
+module PostgresUpsert
 
-  initializer 'postgres_upsert' do
-    ActiveSupport.on_load :active_record do
-      require "postgres_upsert/active_record"
+  class << self
+    def write class_or_table, path_or_io, options = {}
+      writer = class_or_table.is_a?(String) ?  
+         TableWriter : Writer
+      writer.new(class_or_table, path_or_io, options).write
     end
   end
 end

data/lib/postgres_upsert/result.rb

@@ -0,0 +1,11 @@
+module PostgresUpsert
+  class Result
+    attr_reader :inserted, :updated
+
+    def initialize(insert_result, update_result)
+      @inserted = insert_result ? insert_result.cmd_tuples : 0
+      @updated = update_result ? update_result.cmd_tuples : 0  
+    end
+  end
+end
+

data/lib/postgres_upsert/table_writer.rb

@@ -0,0 +1,46 @@
+module PostgresUpsert
+  # alternate version of PostgresUpsert::Writer which does not rely on AR table information.  We 
+  # 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 primary_key
+      @primary_key ||= begin
+        query = <<-sql
+          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
+        sql
+
+        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 = '#{@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/writer.rb

@@ -1,12 +1,10 @@
 module PostgresUpsert
-
   class Writer
 
-    def initialize(table_name, source, options = {})
-      @table_name = table_name
+    def initialize(klass, source, options = {})
+      @klass = klass
       @options = options.reverse_merge({
         :delimiter => ",", 
-        :format => :csv, 
         :header => true, 
         :key_column => primary_key,
         :update_only => false})
@@ -20,16 +18,15 @@ module PostgresUpsert
         raise "Either the :columns option or :header => true are required"
       end
 
-      csv_options = @options[:format] == :binary ? "BINARY" : "DELIMITER '#{@options[:delimiter]}' CSV"
+      csv_options = "DELIMITER '#{@options[:delimiter]}' CSV"
 
       copy_table = @temp_table_name
-
       columns_string = columns_string_for_copy
       create_temp_table
 
       ActiveRecord::Base.connection.raw_connection.copy_data %{COPY #{copy_table} #{columns_string} FROM STDIN #{csv_options}} do
 
-        while line = read_input_line do
+        while line = @source.gets do
           next if line.strip.size == 0
           ActiveRecord::Base.connection.raw_connection.put_copy_data line
         end
@@ -37,46 +34,31 @@ module PostgresUpsert
 
       upsert_from_temp_table
       drop_temp_table
+      PostgresUpsert::Result.new(@insert_result, @update_result)
     end
 
   private
 
     def primary_key
-      @primary_key ||= begin
-        query = <<-sql
-          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
-        sql
-
-        pg_result = ActiveRecord::Base.connection.execute query
-        pg_result.each{ |row| return row['attname'] }
-      end
+      @klass.primary_key
     end
 
     def column_names
-      @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
+      @klass.column_names
+    end
+
+    def quoted_table_name
+      @klass.quoted_table_name
     end
 
     def get_columns
-      columns_list = @options[:columns] || []
-      if @options[:format] != :binary && @options[:header]
+      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
+        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
@@ -120,33 +102,17 @@ module PostgresUpsert
       columns.size > 0 ? "\"#{columns.join('","')}\"" : ""
     end
 
-    def quoted_table_name
-      @quoted_table_name ||= ActiveRecord::Base.connection.quote_table_name(@table_name)
-    end
-
     def generate_temp_table_name
       @temp_table_name = "#{@table_name}_temp_#{rand(1000)}"
     end
 
-    def read_input_line
-      if @options[:format] == :binary
-        begin
-          return @source.readpartial(10240)
-        rescue EOFError
-        end
-      else
-        line = @source.gets
-        return line
-      end
-    end
-
     def upsert_from_temp_table
       update_from_temp_table
       insert_from_temp_table unless @options[:update_only]
     end
 
     def update_from_temp_table
-      ActiveRecord::Base.connection.execute <<-SQL
+      @update_result = ActiveRecord::Base.connection.execute <<-SQL
         UPDATE #{quoted_table_name} AS d
           #{update_set_clause}
           FROM #{@temp_table_name} as t
@@ -166,7 +132,7 @@ module PostgresUpsert
     def insert_from_temp_table
       columns_string = columns_string_for_insert
       select_string = select_string_for_insert
-      ActiveRecord::Base.connection.execute <<-SQL
+      @insert_result = ActiveRecord::Base.connection.execute <<-SQL
         INSERT INTO #{quoted_table_name} (#{columns_string})
           SELECT #{select_string}
           FROM #{@temp_table_name} as t
@@ -180,6 +146,7 @@ module PostgresUpsert
 
     def create_temp_table
       columns_string = select_string_for_create
+      verify_temp_has_key
       ActiveRecord::Base.connection.execute <<-SQL
         SET client_min_messages=WARNING;
         DROP TABLE IF EXISTS #{@temp_table_name};
@@ -189,12 +156,16 @@ module PostgresUpsert
       SQL
     end
 
+    def verify_temp_has_key
+      unless @columns_list.include?(@options[:key_column].to_s)
+        raise "Expected a unique column '#{@options[:key_column]}' but the source data does not include this column.  Update the :columns list or explicitly set the :key_column option.}"
+      end
+    end
+
     def drop_temp_table
       ActiveRecord::Base.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 = "3.0.0"
+  s.version = "4.0.0"
 
   s.platform    = Gem::Platform::RUBY
   s.required_ruby_version     = ">= 1.8.7"
@@ -26,8 +26,8 @@ Gem::Specification.new do |s|
   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 "rspec-rails", "~> 2.0"
 end
 

data/spec/pg_upsert_csv_spec.rb

@@ -18,14 +18,14 @@ describe "pg_upsert from file with CSV format" do
   end
 
   it "should import from file if path is passed without field_map" do
-    TestModel.pg_upsert File.expand_path('spec/fixtures/comma_with_header.csv')
+    PostgresUpsert.write TestModel, File.expand_path('spec/fixtures/comma_with_header.csv')
     expect(
       TestModel.first.attributes
     ).to include('data' => 'test data 1', 'created_at' => timestamp, 'updated_at' => timestamp)
   end
 
   it "correctly handles delimiters in content" do
-    TestModel.pg_upsert File.expand_path('spec/fixtures/comma_with_header_and_comma_values.csv')
+    PostgresUpsert.write TestModel, File.expand_path('spec/fixtures/comma_with_header_and_comma_values.csv')
     expect(
       TestModel.first.attributes
     ).to include('data' => 'test, the data 1', 'created_at' => timestamp, 'updated_at' => timestamp)
@@ -33,39 +33,33 @@ describe "pg_upsert from file with CSV format" do
 
   it "throws error if csv is malformed" do
     expect{
-      TestModel.pg_upsert File.expand_path('spec/fixtures/comma_with_header_and_unquoted_comma.csv')
-    }.to raise_error
-  end
-
-  it "throws error if the csv has mixed delimiters" do
-    expect{
-      TestModel.pg_upsert File.expand_path('spec/fixtures/tab_with_error.csv'), :delimiter => "\t"
+      PostgresUpsert.write TestModel, File.expand_path('spec/fixtures/comma_with_header_and_unquoted_comma.csv')
     }.to raise_error
   end
 
   it "should import from IO without field_map" do
-    TestModel.pg_upsert File.open(File.expand_path('spec/fixtures/comma_with_header.csv'), 'r')
+    PostgresUpsert.write TestModel, File.open(File.expand_path('spec/fixtures/comma_with_header.csv'), 'r')
     expect(
       TestModel.first.attributes
     ).to include('data' => 'test data 1', 'created_at' => timestamp, 'updated_at' => timestamp)
   end
 
   it "should import with custom delimiter from path" do
-    TestModel.pg_upsert File.expand_path('spec/fixtures/semicolon_with_header.csv'), :delimiter => ';'
+    PostgresUpsert.write TestModel, File.expand_path('spec/fixtures/semicolon_with_header.csv'), :delimiter => ';'
     expect(
       TestModel.first.attributes
     ).to include('data' => 'test data 1', 'created_at' => timestamp, 'updated_at' => timestamp)
   end
 
   it "should import with custom delimiter from IO" do
-    TestModel.pg_upsert File.open(File.expand_path('spec/fixtures/semicolon_with_header.csv'), 'r'), :delimiter => ';'
+    PostgresUpsert.write TestModel, File.open(File.expand_path('spec/fixtures/semicolon_with_header.csv'), 'r'), :delimiter => ';'
     expect(
       TestModel.first.attributes
     ).to include('data' => 'test data 1', 'created_at' => timestamp, 'updated_at' => timestamp)
   end
 
   it "should not expect a header when :header is false" do
-    TestModel.pg_upsert(File.open(File.expand_path('spec/fixtures/comma_without_header.csv'), 'r'), :header => false, :columns => [:id,:data])
+    PostgresUpsert.write(TestModel, File.open(File.expand_path('spec/fixtures/comma_without_header.csv'), 'r'), :header => false, :columns => [:id,:data])
 
     expect(
       TestModel.first.attributes
@@ -73,7 +67,7 @@ describe "pg_upsert from file with CSV format" do
   end
 
   it "should be able to map the header in the file to diferent column names" do
-    TestModel.pg_upsert(File.open(File.expand_path('spec/fixtures/tab_with_different_header.csv'), 'r'), :delimiter => "\t", :map => {'cod' => 'id', 'info' => 'data'})
+    PostgresUpsert.write(TestModel, File.open(File.expand_path('spec/fixtures/tab_with_different_header.csv'), 'r'), :delimiter => "\t", :map => {'cod' => 'id', 'info' => 'data'})
 
     expect(
       TestModel.first.attributes
@@ -81,7 +75,7 @@ describe "pg_upsert from file with CSV format" do
   end
 
   it "should be able to map the header in the file to diferent column names with custom delimiter" do
-    TestModel.pg_upsert(File.open(File.expand_path('spec/fixtures/semicolon_with_different_header.csv'), 'r'), :delimiter => ';', :map => {'cod' => 'id', 'info' => 'data'})
+    PostgresUpsert.write(TestModel, File.open(File.expand_path('spec/fixtures/semicolon_with_different_header.csv'), 'r'), :delimiter => ';', :map => {'cod' => 'id', 'info' => 'data'})
 
     expect(
       TestModel.first.attributes
@@ -89,7 +83,7 @@ describe "pg_upsert from file with CSV format" do
   end
 
   it "should ignore empty lines" do
-    TestModel.pg_upsert(File.open(File.expand_path('spec/fixtures/tab_with_extra_line.csv'), 'r'), :delimiter => "\t")
+    PostgresUpsert.write(TestModel, File.open(File.expand_path('spec/fixtures/tab_with_extra_line.csv'), 'r'), :delimiter => "\t")
 
     expect(
       TestModel.first.attributes
@@ -97,7 +91,7 @@ describe "pg_upsert from file with CSV format" do
   end
 
   it "should not create timestamps when the model does not include them" do
-    ReservedWordModel.pg_upsert File.expand_path('spec/fixtures/reserved_words.csv'), :delimiter => "\t"
+    PostgresUpsert.write ReservedWordModel, File.expand_path('spec/fixtures/reserved_words.csv'), :delimiter => "\t"
 
     expect(
       ReservedWordModel.first.attributes
@@ -113,12 +107,12 @@ describe "pg_upsert from file with CSV format" do
 
     it "should not violate primary key constraint" do
       expect{
-        TestModel.pg_upsert File.expand_path('spec/fixtures/comma_with_header.csv')
+        PostgresUpsert.write TestModel, File.expand_path('spec/fixtures/comma_with_header.csv')
       }.to_not raise_error
     end
 
     it "should upsert (update existing records and insert new records)" do
-      TestModel.pg_upsert File.expand_path('spec/fixtures/tab_with_two_lines.csv'), :delimiter => "\t"
+      PostgresUpsert.write TestModel, File.expand_path('spec/fixtures/tab_with_two_lines.csv'), :delimiter => "\t"
 
       expect(
         TestModel.find(1).attributes
@@ -128,21 +122,33 @@ describe "pg_upsert from file with CSV format" do
       ).to eq("id"=>2, "data"=>"test data 2", "created_at" => timestamp, "updated_at" => timestamp)
     end
 
+    it "should return updated and inserted results" do
+      result = PostgresUpsert.write TestModel, File.expand_path('spec/fixtures/tab_with_two_lines.csv'), :delimiter => "\t"
+
+     expect(
+        result.updated
+      ).to eq(1)
+
+      expect(
+        result.inserted
+      ).to eq(1)
+    end
+
     it "should require columns option if no header" do
       expect{
-        TestModel.pg_upsert File.expand_path('spec/fixtures/2_col_binary_data.dat'), :format => :binary
+        PostgresUpsert.write TestModel, File.expand_path('spec/fixtures/comma_without_header.csv'), :header => false
       }.to raise_error("Either the :columns option or :header => true are required")
     end
 
     it "should clean up the temp table after completion" do
-      TestModel.pg_upsert File.expand_path('spec/fixtures/tab_with_two_lines.csv'), :delimiter => "\t"
+      PostgresUpsert.write TestModel, File.expand_path('spec/fixtures/tab_with_two_lines.csv'), :delimiter => "\t"
       
       ActiveRecord::Base.connection.tables.should_not include("test_models_temp")
     end
 
     it "should gracefully drop the temp table if it already exists" do
       ActiveRecord::Base.connection.execute "CREATE TEMP TABLE test_models_temp (LIKE test_models);"
-      TestModel.pg_upsert File.expand_path('spec/fixtures/tab_with_two_lines.csv'), :delimiter => "\t"
+      PostgresUpsert.write TestModel, File.expand_path('spec/fixtures/tab_with_two_lines.csv'), :delimiter => "\t"
 
       expect(
         TestModel.find(1).attributes
@@ -154,7 +160,7 @@ describe "pg_upsert from file with CSV format" do
 
     it "should be able to copy using custom set of columns" do
       ThreeColumn.create(id: 1, data: "old stuff", extra: "neva change!", created_at: original_created_at)
-      ThreeColumn.pg_upsert(File.open(File.expand_path('spec/fixtures/tab_only_data.csv'), 'r'), :delimiter => "\t", :columns => ["id", "data"])
+      PostgresUpsert.write(ThreeColumn, File.open(File.expand_path('spec/fixtures/tab_only_data.csv'), 'r'), :delimiter => "\t", :columns => ["id", "data"])
 
       expect(
         ThreeColumn.first.attributes
@@ -167,8 +173,7 @@ describe "pg_upsert from file with CSV format" do
       three_col = ThreeColumn.create(id: 1, data: "old stuff", extra: "neva change!")
       file = File.open(File.expand_path('spec/fixtures/no_id.csv'), 'r')
 
-
-      ThreeColumn.pg_upsert(file, :key_column => "data")
+      PostgresUpsert.write(ThreeColumn, file, :key_column => "data")
       expect(
         three_col.reload.extra
       ).to eq("ABC: Always Be Changing.")
@@ -177,11 +182,29 @@ describe "pg_upsert from file with CSV format" do
     it 'inserts records if the passed match column doesnt exist' do
       file = File.open(File.expand_path('spec/fixtures/no_id.csv'), 'r')
 
-      ThreeColumn.pg_upsert(file, :key_column => "data")
+      PostgresUpsert.write(ThreeColumn, file, :key_column => "data")
       expect(
         ThreeColumn.last.attributes
-      ).to include("id" => 1, "data" => "old stuff", "extra" => "ABC: Always Be Changing.")
+      ).to include("data" => "old stuff", "extra" => "ABC: Always Be Changing.")
+    end
+
+    it 'allows key column to be a string or symbol' do
+      file = File.open(File.expand_path('spec/fixtures/no_id.csv'), 'r')
+
+      PostgresUpsert.write(ThreeColumn, file, :header => true, :key_column => :data)
+      expect(
+        ThreeColumn.last.attributes
+      ).to include("data" => "old stuff", "extra" => "ABC: Always Be Changing.")
+    end
+
+    it 'raises an error if the expected key column is not in data' do
+      file = File.open(File.expand_path('spec/fixtures/no_id.csv'), 'r')
+
+      expect{
+      PostgresUpsert.write(ThreeColumn, file, :header => true)
+      }.to raise_error (/Expected a unique column 'id'/)
     end
+
   end
 
   context 'update only' do
@@ -189,9 +212,9 @@ describe "pg_upsert from file with CSV format" do
     before(:each) do
       TestModel.create(id: 1, data: "From the before time, in the long long ago", :created_at => original_created_at)
     end
-    it 'will only update and not insert if insert_only flag is passed.' do
-      TestModel.pg_upsert File.expand_path('spec/fixtures/tab_with_two_lines.csv'), :delimiter => "\t", :update_only => true
 
+    it 'will only update and not insert if insert_only flag is passed.' do
+      PostgresUpsert.write TestModel, File.expand_path('spec/fixtures/tab_with_two_lines.csv'), :delimiter => "\t", :update_only => true
       expect(
         TestModel.find(1).attributes
       ).to eq("id"=>1, "data"=>"test data 1", "created_at" => original_created_at  , "updated_at" => timestamp)
@@ -201,6 +224,41 @@ describe "pg_upsert from file with CSV format" do
 
     end
 
+    it 'will return the number of updated rows' do
+      a = PostgresUpsert.write TestModel, File.expand_path('spec/fixtures/tab_with_two_lines.csv'), :delimiter => "\t", :update_only => true
+      expect(
+        a.updated
+      ).to eq(1)
+
+      expect(
+        a.inserted
+      ).to eq(0)
+    end
+
+  end
+
+  context 'using table_name' do
+    it "should import from file if path is passed without field_map" do
+      PostgresUpsert.write TestModel.table_name, File.expand_path('spec/fixtures/comma_with_header.csv')
+      expect(
+        TestModel.first.attributes
+      ).to include('data' => 'test data 1', 'created_at' => timestamp, 'updated_at' => timestamp)
+    end
+
+    it "should still report results" do
+      TestModel.create(data: "test data 1")
+      result = result = PostgresUpsert.write TestModel.table_name, File.expand_path('spec/fixtures/tab_with_two_lines.csv'), :delimiter => "\t"
+
+      expect(
+        result.updated
+      ).to eq(1)
+
+      expect(
+        result.inserted
+      ).to eq(1)
+    end
+
+
   end
 end