upsert 0.3.4 → 0.4.0

data/CHANGELOG

@@ -1,3 +1,15 @@
+0.4.0 / 2012-09-04
+
+* Bug fixes
+
+  * Don't raise TooBig - rely on Mysql2 to complain about oversized packets
+
+* Enhancements
+
+  * Re-use PostgreSQL merge functions across connections, even outside of batch mode. Huzzah!
+  * For MySQL, increase speed for one-off upserts by not checking packet size
+  * Allow configuring Upsert.logger. Defaults to Rails.logger or Logger.new($stderr). If you set env var UPSERT_DEBUG=true then it will set log level to debug.
+
 0.3.4 / 2012-07-03
 
 * Bug fixes

data/README.md

@@ -28,13 +28,13 @@ Rows are buffered in memory until it's efficient to send them to the database.
 
 Tested to be much about 85% faster on PostgreSQL and 50% faster on MySQL than comparable methods (see the tests).
 
-### Gotchas
+## Gotchas
 
-#### Undefined behavior if you use this without properly defining UNIQUE indexes
+### Undefined behavior without real UNIQUE indexes
 
 Make sure you're upserting against either primary key columns or columns with UNIQUE indexes or both.
 
-#### Columns are set based on the first row you pass
+### Columns are set based on the first row you pass
 
 Currently, the first row you pass in determines the columns that will be used. That's useful for mass importing of many rows with the same columns, but is surprising if you're trying to use a single `Upsert` object to add arbitrary data. For example, this won't work:
 
@@ -52,12 +52,9 @@ You would need to use a new `Upsert` object. On the other hand, this is totally
 
 Pull requests for any of these would be greatly appreciated:
 
-1. Fix SQLite tests.
-2. For PG, be smarter about when you create functions - try to re-use them within a connection.
-3. Provide optional SQL logging.
-4. Provide `require 'upsert/debug'` that will make sure you are selecting on columns that have unique indexes
-5. Make `Upsert` instances accept arbitrary columns, which is what people probably expect.
-6. Naming suggestions: should "document" be called "setters" or "attributes"?
+1. Provide `require 'upsert/debug'` that will make sure you are selecting on columns that have unique indexes
+1. Make `Upsert` instances accept arbitrary columns, which is what people probably expect. (this should work on PG already)
+1. Naming suggestions: should "document" be called "setters" or "attributes"?
 
 ## Real-world usage
 

data/Rakefile

@@ -1,25 +1,20 @@
 #!/usr/bin/env rake
 require "bundler/gem_tasks"
 
-require 'rake'
-require 'rake/testtask'
-Rake::TestTask.new(:_test) do |test|
-  test.libs << 'lib' << 'test'
-  test.pattern = 'test/**/test_*.rb'
-  test.verbose = true
-end
-
-task :test_each_db_adapter do
-  %w{ active_record_upsert mysql2 sqlite pg active_record_connection_adapter }.each do |database|
+task :rspec_all_databases do
+  require 'posix-spawn'
+  %w{ postgresql mysql2 sqlite3 }.each do |adapter|
     puts
-    puts "#{'*'*10} Running #{database} tests"
+    puts '#'*50
+    puts "# Running specs against #{adapter}"
+    puts '#'*50
     puts
-    puts `rake _test TEST=test/test_#{database}.rb`
+    pid = POSIX::Spawn.spawn({'ADAPTER' => adapter}, 'rspec', '--format', 'documentation', File.expand_path('../spec', __FILE__))
+    Process.waitpid pid
   end
 end
 
-task :default => :test_each_db_adapter
-task :test => :test_each_db_adapter
+task :default => :rspec_all_databases
 
 require 'yard'
 YARD::Rake::YardocTask.new

data/lib/upsert.rb

@@ -1,14 +1,39 @@
 require 'bigdecimal'
+require 'thread'
+require 'logger'
 
 require 'upsert/version'
 require 'upsert/binary'
+require 'upsert/buffer'
+require 'upsert/connection'
 require 'upsert/row'
-require 'upsert/mysql2_client'
-require 'upsert/pg_connection'
-require 'upsert/sqlite3_database'
 
 class Upsert
   class << self
+    # What logger to use.
+    # @return [#info,#warn,#debug]
+    attr_writer :logger
+    
+    # The current logger
+    # @return [#info,#warn,#debug]
+    def logger
+      @logger || Thread.exclusive do
+        @logger ||= if defined?(::Rails) and (rails_logger = ::Rails.logger)
+          rails_logger
+        elsif defined?(::ActiveRecord) and ::ActiveRecord.const_defined?(:Base) and (ar_logger = ::ActiveRecord::Base.logger)
+          ar_logger
+        else
+          my_logger = Logger.new $stderr
+          my_logger.level = Logger::INFO
+          my_logger
+        end
+        if ENV['UPSERT_DEBUG'] == 'true'
+          @logger.level = Logger::DEBUG
+        end
+        @logger
+      end
+    end
+
     # @param [String] v A string containing binary data that should be inserted/escaped as such.
     #
     # @return [Upsert::Binary]
@@ -20,8 +45,6 @@ class Upsert
     #
     # @note Buffered in memory until it's efficient to send to the server a packet.
     #
-    # @raise [Upsert::TooBig] If any row is too big to fit inside a single packet.
-    #
     # @return [nil]
     #
     # @example Many at once
@@ -31,19 +54,15 @@ class Upsert
     #   end
     def batch(connection, table_name)
       upsert = new connection, table_name
-      upsert.async!
+      upsert.buffer.async!
       yield upsert
-      upsert.sync!
+      upsert.buffer.sync!
     end
 
     # @deprecated Use .batch instead.
     alias :stream :batch
   end
 
-  # Raised if a query would be too large to send in a single packet.
-  class TooBig < RuntimeError
-  end
-
   SINGLE_QUOTE = %{'}
   DOUBLE_QUOTE = %{"}
   BACKTICK = %{`}
@@ -54,10 +73,10 @@ class Upsert
   ISO8601_DATE = '%F'
   NULL_WORD = 'NULL'
 
-  # @return [Mysql2::Client,Sqlite3::Database,PG::Connection,#raw_connection]
+  # @return [Upsert::Connection]
   attr_reader :connection
 
-  # @return [String,Symbol]
+  # @return [String]
   attr_reader :table_name
 
   # @private
@@ -66,17 +85,11 @@ class Upsert
   # @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)
-    @table_name = table_name
-    @buffer = []
-
-    @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+/, '_'))
+    @table_name = table_name.to_s
+    raw_connection = connection.respond_to?(:raw_connection) ? connection.raw_connection : connection
+    n = raw_connection.class.name.gsub(/\W+/, '_')
+    @connection = Connection.const_get(n).new self, raw_connection
+    @buffer = Buffer.const_get(n).new self
   end
 
   # Upsert a row given a selector and a document.
@@ -86,8 +99,6 @@ class Upsert
   # @param [Hash] selector Key-value pairs that will be used to find or create a row.
   # @param [Hash] document Key-value pairs that will be set on the row, whether it previously existed or not.
   #
-  # @raise [Upsert::TooBig] If any row is too big to fit inside a single packet.
-  #
   # @return [nil]
   #
   # @example One at a time
@@ -95,54 +106,12 @@ class Upsert
   #   upsert.row({:name => 'Jerry'}, :breed => 'beagle')
   #   upsert.row({:name => 'Pierre'}, :breed => 'tabby')
   def row(selector, document = {})
-    buffer.push Row.new(self, selector, document)
-    if sql = chunk
-      execute sql
-    end
+    buffer << Row.new(self, selector, document)
     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_WORD
-    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
+  def quoted_table_name
+    @quoted_table_name ||= connection.quote_ident table_name
   end
 end

data/lib/upsert/buffer.rb

@@ -0,0 +1,36 @@
+require 'upsert/buffer/mysql2_client'
+require 'upsert/buffer/pg_connection'
+require 'upsert/buffer/sqlite3_database'
+
+class Upsert
+  # @private
+  class Buffer
+    attr_reader :parent
+    attr_reader :rows
+
+    def initialize(parent)
+      @parent = parent
+      @rows = []
+    end
+
+    def <<(row)
+      rows << row
+      ready
+    end
+
+    def async?
+      !!@async
+    end
+
+    def async!
+      @async = true
+    end
+
+    def sync!
+      @async = false
+      until rows.empty?
+        ready
+      end
+    end
+  end
+end

data/lib/upsert/buffer/mysql2_client.rb

@@ -0,0 +1,67 @@
+class Upsert
+  # @private
+  class Buffer
+    class Mysql2_Client < Buffer
+      def ready
+        return if rows.empty?
+        c = parent.connection
+        if not async?
+          c.execute sql
+          rows.clear
+          return
+        end
+        @cumulative_sql_bytesize ||= static_sql_bytesize
+        new_row = rows.pop
+        d = new_row.values_sql_bytesize + 3 # ),(
+        if @cumulative_sql_bytesize + d > max_sql_bytesize
+          c.execute sql
+          rows.clear
+          @cumulative_sql_bytesize = static_sql_bytesize + d
+        else
+          @cumulative_sql_bytesize += d
+        end
+        rows << new_row
+        nil
+      end
+
+      def columns
+        @columns ||= rows.first.columns
+      end
+
+      def insert_part
+        @insert_part ||= begin
+          connection = parent.connection
+          columns_sql = columns.map { |k| connection.quote_ident(k) }.join(',')
+          %{INSERT INTO #{parent.quoted_table_name} (#{columns_sql}) VALUES }
+        end
+      end
+
+      def update_part
+        @update_part ||= begin
+          connection = parent.connection
+          updaters = columns.map do |k|
+            qk = connection.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
+
+      def sql
+        all_value_sql = rows.map { |row| row.values_sql }
+        retval = [ insert_part, '(', all_value_sql.join('),('), ')', update_part ].join
+        retval
+      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 ||= parent.connection.database_variable_get(:MAX_ALLOWED_PACKET).to_i
+      end
+    end
+  end
+end

data/lib/upsert/buffer/pg_connection.rb

@@ -0,0 +1,54 @@
+require 'upsert/buffer/pg_connection/merge_function'
+
+class Upsert
+  class Buffer
+    # @private
+    class PG_Connection < Buffer
+      def ready
+        return if rows.empty?
+        row = rows.shift
+        MergeFunction.execute(self, row)
+      end
+
+      def clear_database_functions
+        connection = parent.connection
+        # http://stackoverflow.com/questions/7622908/postgresql-drop-function-without-knowing-the-number-type-of-parameters
+        connection.execute <<-EOS
+CREATE OR REPLACE FUNCTION pg_temp.upsert_delfunc(text)
+  RETURNS void AS
+$BODY$
+DECLARE
+   _sql text;
+BEGIN
+
+FOR _sql IN
+   SELECT 'DROP FUNCTION ' || quote_ident(n.nspname)
+                    || '.' || quote_ident(p.proname)
+           || '(' || pg_catalog.pg_get_function_identity_arguments(p.oid) || ');'
+   FROM   pg_catalog.pg_proc p
+   LEFT   JOIN pg_catalog.pg_namespace n ON n.oid = p.pronamespace
+   WHERE  p.proname = $1
+   AND    pg_catalog.pg_function_is_visible(p.oid) -- you may or may not want this
+LOOP
+   EXECUTE _sql;
+END LOOP;
+
+END;
+$BODY$
+  LANGUAGE plpgsql;
+EOS
+        res = connection.execute(%{SELECT proname FROM pg_proc WHERE proname LIKE 'upsert_%'})
+        res.each do |row|
+          k = row['proname']
+          next if k == 'upsert_delfunc'
+          Upsert.logger.info %{[upsert] Dropping function #{k.inspect}}
+          connection.execute %{SELECT pg_temp.upsert_delfunc('#{k}')}
+        end
+      end
+    end
+
+    # @private
+    # backwards compatibility - https://github.com/seamusabshere/upsert/issues/2
+    PGconn = PG_Connection
+  end
+end

data/lib/upsert/buffer/pg_connection/merge_function.rb

@@ -0,0 +1,138 @@
+require 'digest/md5'
+
+class Upsert
+  # @private
+  class Buffer
+    class PG_Connection < Buffer
+      class MergeFunction
+        class << self
+          def execute(buffer, row)
+            first_try = true
+            begin
+              buffer.parent.connection.execute sql(buffer, row)
+            rescue PG::Error => pg_error
+              if first_try and pg_error.message =~ /function upsert_(.+) does not exist/
+                Upsert.logger.info %{[upsert] Function #{"upsert_#{$1}".inspect} went missing, trying to recreate}
+                first_try = false
+                @lookup.clear
+                retry
+              else
+                raise pg_error
+              end
+            end
+          end
+
+          def sql(buffer, row)
+            merge_function = lookup buffer, row
+            %{SELECT #{merge_function.name}(#{merge_function.values_sql(row)})}
+          end
+
+          def unique_key(table_name, selector, columns)
+            [
+              table_name,
+              selector.join(','),
+              columns.join(',')
+            ].join '/'
+          end
+
+          def lookup(buffer, row)
+            @lookup ||= {}
+            s = row.selector.keys
+            c = row.columns
+            @lookup[unique_key(buffer.parent.table_name, s, c)] ||= new(buffer, s, c)
+          end
+        end
+
+        attr_reader :buffer
+        attr_reader :selector
+        attr_reader :columns
+
+        def initialize(buffer, selector, columns)
+          @buffer = buffer
+          @selector = selector
+          @columns = columns
+          create!
+        end
+
+        def name
+          @name ||= "upsert_#{Digest::MD5.hexdigest(unique_key)}"
+        end
+
+        def values_sql(row)
+          ordered_args = columns.map do |k|
+            row.quoted_value(k) || NULL_WORD
+          end.join(',')
+        end
+
+        private
+
+        def unique_key
+          @unique_key ||= MergeFunction.unique_key buffer.parent.table_name, selector, columns
+        end
+
+        def connection
+          buffer.parent.connection
+        end
+
+        def quoted_table_name
+          buffer.parent.quoted_table_name
+        end
+
+        ColumnDefinition = Struct.new(:quoted_name, :quoted_input_name, :sql_type, :default)
+
+        # activerecord-3.2.5/lib/active_record/connection_adapters/postgresql_adapter.rb#column_definitions
+        def get_column_definitions
+          res = connection.execute <<-EOS
+SELECT a.attname AS name, format_type(a.atttypid, a.atttypmod) AS sql_type, d.adsrc AS default
+FROM pg_attribute a LEFT JOIN pg_attrdef d
+  ON a.attrelid = d.adrelid AND a.attnum = d.adnum
+WHERE a.attrelid = '#{quoted_table_name}'::regclass
+  AND a.attnum > 0 AND NOT a.attisdropped
+EOS
+          unsorted = res.select do |row|
+            columns.include? row['name']
+          end.inject({}) do |memo, row|
+            k = row['name']
+            memo[k] = ColumnDefinition.new connection.quote_ident(k), connection.quote_ident("#{k}_input"), row['sql_type'], row['default']
+            memo
+          end
+          columns.map do |k|
+            unsorted[k]
+          end
+        end
+
+        # the "canonical example" from http://www.postgresql.org/docs/9.1/static/plpgsql-control-structures.html#PLPGSQL-UPSERT-EXAMPLE
+        def create!
+          Upsert.logger.info "[upsert] Creating or replacing database function #{name.inspect} on table #{buffer.parent.table_name.inspect} for selector #{selector.map(&:inspect).join(', ')} and columns #{columns.map(&:inspect).join(', ')}"
+          column_definitions = get_column_definitions
+          connection.execute <<-EOS
+CREATE OR REPLACE FUNCTION #{name}(#{column_definitions.map { |c| "#{c.quoted_input_name} #{c.sql_type} DEFAULT #{c.default || 'NULL'}" }.join(',') }) RETURNS VOID AS
+$$
+BEGIN
+  LOOP
+      -- first try to update the key
+      UPDATE #{quoted_table_name} SET #{column_definitions.map { |c| "#{c.quoted_name} = #{c.quoted_input_name}" }.join(',')}
+          WHERE #{selector.map { |k| "#{connection.quote_ident(k)} = #{connection.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 #{quoted_table_name}(#{column_definitions.map { |c| c.quoted_name }.join(',')}) VALUES (#{column_definitions.map { |c| c.quoted_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
+  end
+end