upsert 0.1.2 → 0.2.0

data/.gitignore

@@ -1,3 +1,4 @@
+.DS_Store
 *.gem
 *.rbc
 .bundle

data/README.md

@@ -18,7 +18,9 @@ Let's say you have...
     document = {:breed => 'beagle'}
     upsert.row selector, document
 
-### Multiple upserts bundled together for speed
+### Streaming upserts (fastest)
+
+Rows are buffered in memory until it's efficient to send them to the database.
 
     Upsert.stream(Pet.connection, Pet.table_name) do |upsert|
       # [...]
@@ -28,7 +30,14 @@ Let's say you have...
       # [...]
     end
 
-Rows are buffered in memory until it's efficient to send them to the database.
+### With a helper method
+
+For bulk upserts, you probably still want to use `Upsert.stream`.
+
+    # be sure to require 'upsert/active_record_upsert' - it's not required by default
+    selector = {:name => 'Jerry'}
+    document = {:breed => 'beagle'}
+    Pet.upsert selector, document
 
 ## Real-world usage
 
@@ -54,9 +63,9 @@ Using the [mysql2](https://rubygems.org/gems/mysql2) driver.
 From the tests:
 
     Upsert was 77% faster than find + new/set/save
-    Upsert was 84% faster than create + rescue/find/update
-    Upsert was 82% faster than find_or_create + update_attributes
-    Upsert was 47% faster than faking upserts with activerecord-import
+    Upsert was 58% faster than create + rescue/find/update
+    Upsert was 80% faster than find_or_create + update_attributes
+    Upsert was 39% faster than faking upserts with activerecord-import
 
 #### SQL MERGE trick
 
@@ -199,11 +208,6 @@ This, however, only works on MySQL and requires ActiveRecord—and if all yo
 
 The `selector` and `document` arguments are inspired by the upsert functionality of the [mongo-ruby-driver's update method](http://api.mongodb.org/ruby/1.6.4/Mongo/Collection.html#update-instance_method).
 
-## Wishlist
-
-1. `Pet.upsert`... duh
-2. Don't need a separate buffer class... just extend an instance of Upsert with the appropriate database driver module.
-
 ## Copyright
 
 Copyright 2012 Brighter Planet, Inc.

data/Rakefile

@@ -10,7 +10,7 @@ Rake::TestTask.new(:_test) do |test|
 end
 
 task :test_each_db_adapter do
-  %w{ mysql2 sqlite pg active_record_connection_adapter }.each do |database|
+  %w{ active_record_upsert mysql2 sqlite pg active_record_connection_adapter }.each do |database|
     puts
     puts "#{'*'*10} Running #{database} tests"
     puts

data/lib/upsert.rb

@@ -2,12 +2,10 @@ require 'bigdecimal'
 
 require 'upsert/version'
 require 'upsert/binary'
-require 'upsert/buffer'
-require 'upsert/quoter'
 require 'upsert/row'
-require 'upsert/buffer/mysql2_client'
-require 'upsert/buffer/pg_connection'
-require 'upsert/buffer/sqlite3_database'
+require 'upsert/mysql2_client'
+require 'upsert/pg_connection'
+require 'upsert/sqlite3_database'
 
 class Upsert
   class << self
@@ -33,9 +31,9 @@ class Upsert
     #   end
     def stream(connection, table_name)
       upsert = new connection, table_name
-      upsert.buffer.async!
+      upsert.async!
       yield upsert
-      upsert.buffer.sync!
+      upsert.sync!
     end
   end
 
@@ -43,13 +41,38 @@ class Upsert
   class TooBig < RuntimeError
   end
 
+  SINGLE_QUOTE = %{'}
+  DOUBLE_QUOTE = %{"}
+  BACKTICK = %{`}
+  E_AND_SINGLE_QUOTE = %{E'}
+  X_AND_SINGLE_QUOTE = %{x'}
+  USEC_SPRINTF = '%06d'
+  ISO8601_DATETIME = '%Y-%m-%d %H:%M:%S'
+  ISO8601_DATE = '%F'
+
+  # @return [Mysql2::Client,Sqlite3::Database,PG::Connection,#raw_connection]
+  attr_reader :connection
+
+  # @return [String,Symbol]
+  attr_reader :table_name
+
   # @private
-  attr_reader :buffer
+  attr_reader :rows
 
   # @param [Mysql2::Client,Sqlite3::Database,PG::Connection,#raw_connection] connection A supported database connection.
   # @param [String,Symbol] table_name The name of the table into which you will be upserting.
   def initialize(connection, table_name)
-    @buffer = Buffer.for connection, table_name
+    @table_name = table_name
+    @rows = []
+
+    @connection = if connection.respond_to?(:raw_connection)
+      # deal with ActiveRecord::Base.connection or ActiveRecord::Base.connection_pool.checkout
+       connection.raw_connection
+    else
+      connection
+    end
+
+    extend Upsert.const_get(@connection.class.name.gsub(/\W+/, '_'))
   end
 
   # Upsert a row given a selector and a document.
@@ -68,7 +91,69 @@ class Upsert
   #   upsert.row({:name => 'Jerry'}, :breed => 'beagle')
   #   upsert.row({:name => 'Pierre'}, :breed => 'tabby')
   def row(selector, document)
-    buffer.add selector, document
+    rows << Row.new(self, selector, document)
+    if sql = chunk
+      execute sql
+    end
     nil
   end
+
+  # @private
+  def async?
+    !!@async
+  end
+
+  # @private
+  def async!
+    @async = true
+  end
+
+  # @private
+  def sync!
+    @async = false
+    while sql = chunk
+      execute sql
+    end
+  end
+
+  # @private
+  def quote_value(v)
+    case v
+    when NilClass
+      'NULL'
+    when Upsert::Binary
+      quote_binary v # must be defined by base
+    when String
+      quote_string v # must be defined by base
+    when TrueClass, FalseClass
+      quote_boolean v
+    when BigDecimal
+      quote_big_decimal v
+    when Numeric
+      v
+    when Symbol
+      quote_string v.to_s
+    when Time, DateTime
+      quote_time v # must be defined by base
+    when Date
+      quote_string v.strftime(ISO8601_DATE)
+    else
+      raise "not sure how to quote #{v.class}: #{v.inspect}"
+    end
+  end
+
+  # @private
+  def quote_idents(idents)
+    idents.map { |k| quote_ident(k) }.join(',') # must be defined by base
+  end
+
+  # @private
+  def quote_values(values)
+    values.map { |v| quote_value(v) }.join(',')
+  end
+  
+  # @private
+  def quote_pairs(pairs)
+    pairs.map { |k, v| [quote_ident(k),quote_value(v)].join('=') }.join(',')
+  end
 end

data/lib/upsert/active_record_upsert.rb

@@ -0,0 +1,12 @@
+class Upsert
+  module ActiveRecordUpsert
+    def upsert(selector, document)
+      ActiveRecord::Base.connection_pool.with_connection do |c|
+        upsert = Upsert.new c, table_name
+        upsert.row selector, document
+      end
+    end
+  end
+end
+
+ActiveRecord::Base.extend Upsert::ActiveRecordUpsert

data/lib/upsert/mysql2_client.rb

@@ -0,0 +1,160 @@
+class Upsert
+  # @private
+  module Mysql2_Client
+    def chunk
+      return if rows.empty?
+      all = rows.length
+      take = all
+      while take > 1 and probably_oversize?(take)
+        take -= 1
+      end
+      if async? and take == all
+        return
+      end
+      while take > 1 and oversize?(take)
+        $stderr.puts "   Length prediction via sampling failed, shrinking" if ENV['UPSERT_DEBUG'] == 'true'
+        take -= 1
+      end
+      chunk = sql take
+      while take > 1 and chunk.bytesize > max_sql_bytesize
+        $stderr.puts "   Supposedly exact bytesize guess failed, shrinking" if ENV['UPSERT_DEBUG'] == 'true'
+        take -= 1
+        chunk = sql take
+      end
+      if chunk.bytesize > max_sql_bytesize
+        raise TooBig
+      end
+      $stderr.puts "   Chunk (#{take}/#{chunk.bytesize}) was #{(chunk.bytesize / max_sql_bytesize.to_f * 100).round}% of the max" if ENV['UPSERT_DEBUG'] == 'true'
+      @rows = rows.drop(take)
+      chunk
+    end
+
+    def execute(sql)
+      connection.query sql
+    end
+
+    def probably_oversize?(take)
+      estimate_sql_bytesize(take) > max_sql_bytesize
+    end
+
+    def oversize?(take)
+      sql_bytesize(take) > max_sql_bytesize
+    end
+
+    def columns
+      @columns ||= rows.first.columns
+    end
+
+    def insert_part
+      @insert_part ||= %{INSERT INTO "#{table_name}" (#{quote_idents(columns)}) VALUES }
+    end
+
+    def update_part
+      @update_part ||= begin
+        updaters = columns.map do |k|
+          qk = quote_ident k
+          [ qk, "VALUES(#{qk})" ].join('=')
+        end.join(',')
+        %{ ON DUPLICATE KEY UPDATE #{updaters}}
+      end
+    end
+
+    # where 2 is the parens
+    def static_sql_bytesize
+      @static_sql_bytesize ||= insert_part.bytesize + update_part.bytesize + 2
+    end
+
+    # where 3 is parens and comma
+    def variable_sql_bytesize(take)
+      rows.first(take).inject(0) { |sum, row| sum + row.values_sql_bytesize + 3 }
+    end
+
+    def estimate_variable_sql_bytesize(take)
+      p = (take / 10.0).ceil
+      10.0 * rows.sample(p).inject(0) { |sum, row| sum + row.values_sql_bytesize + 3 }
+    end
+
+    def sql_bytesize(take)
+      static_sql_bytesize + variable_sql_bytesize(take)
+    end
+
+    def estimate_sql_bytesize(take)
+      static_sql_bytesize + estimate_variable_sql_bytesize(take)
+    end
+
+    def sql(take)
+      all_value_sql = rows.first(take).map { |row| row.values_sql }
+      [ insert_part, '(', all_value_sql.join('),('), ')', update_part ].join
+    end
+
+    # since setting an option like :as => :hash actually persists that option to the client, don't pass any options
+    def max_sql_bytesize
+      @max_sql_bytesize ||= begin
+        case (row = connection.query("SHOW VARIABLES LIKE 'max_allowed_packet'").first)
+        when Array
+          row[1]
+        when Hash
+          row['Value']
+        else
+          raise "Don't know what to do if connection.query returns a #{row.class}"
+        end.to_i
+      end
+    end
+
+    def quoted_value_bytesize(v)
+      case v
+      when NilClass
+        4
+      when TrueClass
+        4
+      when FalseClass
+        5
+      when BigDecimal
+        v.to_s('F').bytesize
+      when Upsert::Binary
+        v.bytesize * 2 + 3
+      when Numeric
+        v.to_s.bytesize
+      when String
+        v.bytesize + 2
+      when Symbol
+        v.to_s.bytesize + 2
+      when Time, DateTime
+        24 + 2
+      when Date
+        10 + 2
+      else
+        raise "not sure how to get quoted length of #{v.class}: #{v.inspect}"
+      end
+    end
+
+    def quote_boolean(v)
+      v ? 'TRUE' : 'FALSE'
+    end
+
+    def quote_string(v)
+      SINGLE_QUOTE + connection.escape(v) + SINGLE_QUOTE
+    end
+
+    # This doubles the size of the representation.
+    def quote_binary(v)
+      X_AND_SINGLE_QUOTE + v.unpack("H*")[0] + SINGLE_QUOTE
+    end
+
+    # put raw binary straight into sql
+    # might work if we could get the encoding issues fixed when joining together the values for the sql
+    # alias_method :quote_binary, :quote_string
+
+    def quote_time(v)
+      quote_string v.strftime(ISO8601_DATETIME)
+    end
+
+    def quote_ident(k)
+      BACKTICK + connection.escape(k.to_s) + BACKTICK
+    end
+
+    def quote_big_decimal(v)
+      v.to_s('F')
+    end
+  end
+end

data/lib/upsert/pg_connection.rb

@@ -0,0 +1,84 @@
+require 'upsert/pg_connection/column_definition'
+
+class Upsert
+  # @private
+  module PG_Connection
+
+    attr_reader :merge_function
+
+    def chunk
+      return if rows.empty?
+      row = rows.shift
+      unless merge_function
+        create_merge_function row
+      end
+      hsh = row.to_hash
+      ordered_args = column_definitions.map do |c|
+        hsh[c.name]
+      end
+      %{SELECT #{merge_function}(#{quote_values(ordered_args)})}
+    end
+
+    def execute(sql)
+      connection.exec sql
+    end
+
+    def quote_string(v)
+      SINGLE_QUOTE + connection.escape_string(v) + SINGLE_QUOTE
+    end
+
+    def quote_binary(v)
+      E_AND_SINGLE_QUOTE + connection.escape_bytea(v) + SINGLE_QUOTE
+    end
+
+    def quote_time(v)
+      quote_string [v.strftime(ISO8601_DATETIME), sprintf(USEC_SPRINTF, v.usec)].join('.')
+    end
+
+    def quote_big_decimal(v)
+      v.to_s('F')
+    end
+
+    def quote_boolean(v)
+      v ? 'TRUE' : 'FALSE'
+    end
+
+    def quote_ident(k)
+      connection.quote_ident k.to_s
+    end
+
+    def column_definitions
+      @column_definitions ||= ColumnDefinition.all(connection, table_name)
+    end
+    
+    private
+    
+    def create_merge_function(example_row)
+      @merge_function = "pg_temp.merge_#{table_name}_#{Kernel.rand(1e11)}"
+      execute <<-EOS
+CREATE FUNCTION #{merge_function}(#{column_definitions.map { |c| "#{quote_ident(c.input_name)} #{c.sql_type} DEFAULT #{c.default || 'NULL'}" }.join(',') }) RETURNS VOID AS
+$$
+BEGIN
+  LOOP
+      -- first try to update the key
+      UPDATE #{table_name} SET #{column_definitions.map { |c| "#{quote_ident(c.name)} = #{quote_ident(c.input_name)}" }.join(',')} WHERE #{example_row.selector.keys.map { |k| "#{quote_ident(k)} = #{quote_ident([k,'input'].join('_'))}" }.join(' AND ') };
+      IF found THEN
+          RETURN;
+      END IF;
+      -- not there, so try to insert the key
+      -- if someone else inserts the same key concurrently,
+      -- we could get a unique-key failure
+      BEGIN
+          INSERT INTO #{table_name}(#{column_definitions.map { |c| quote_ident(c.name) }.join(',')}) VALUES (#{column_definitions.map { |c| quote_ident(c.input_name) }.join(',')});
+          RETURN;
+      EXCEPTION WHEN unique_violation THEN
+          -- Do nothing, and loop to try the UPDATE again.
+      END;
+  END LOOP;
+END;
+$$
+LANGUAGE plpgsql;
+EOS
+    end
+  end
+end