em-pg-client-helper 1.0.0 → 1.1.0

data/lib/em-pg-client-helper.rb

@@ -47,6 +47,106 @@ module PG::EM::Client::Helper
 		db.exec_defer(*insert_sql(tbl, params))
 	end
 
+	# @macro upsert_params
+	#
+	#   @param tbl [#to_s] The name of the table on which to operate.
+	#
+	#   @param key [#to_s, Array<#to_s>] A field (or list of fields) which
+	#     are the set of values that uniquely identify the record to be
+	#     updated, if it exists.  You only need to specify the field names
+	#     here, as the values which will be used in the query will be taken
+	#     from the `data`.
+	#
+	#   @param data [Hash<#to_s, Object>] The fields and values to insert into
+	#     the database, or to set in the existing record.
+	#
+	#   @raise [ArgumentError] if a field is specified in `key` but which
+	#     does not exist in `data`.
+	#
+	#
+	# An "upsert" is a kind of crazy hybrid "update if the record exists,
+	# insert it if it doesn't" query.  It isn't part of the SQL standard,
+	# but it is such a common idiom that we're keen to support it.
+	#
+	# The trick is that it's actually two queries in one.  We try to do an
+	# `UPDATE` first, and if that doesn't actually update anything, then we
+	# try an `INSERT`.  Since it is two separate queries, though, there is still
+	# a small chance that the query will fail with a `PG::UniqueViolation`, so
+	# your code must handle that.
+	#
+	# @!macro upsert_params
+	#
+	# @return [Array<String, Array<Object>>] A two-element array, the first
+	#   of which is a string containing the literal SQL to be executed, while
+	#   the second element is an array containing the values, in order
+	#   corresponding to the placeholders in the SQL.
+	#
+	def upsert_sql(tbl, key, data)
+		tbl = quote_identifier(tbl)
+		insert_keys = data.keys.map { |k| quote_identifier(k.to_s) }
+		unique_keys = (key.is_a?(Array) ? key : [key])
+		unique_keys.map! { |k| quote_identifier(k.to_s) }
+		update_keys = insert_keys - unique_keys
+
+		unless (bad_keys = unique_keys - insert_keys).empty?
+			raise ArgumentError,
+			      "These field(s) were mentioned in the key list, but were not in the data set: #{bad_keys.inspect}"
+		end
+
+		values = data.values
+		# field-to-placeholder mapping
+		i = 0
+		fp_map = Hash[insert_keys.map { |k| i += 1; [k, "$#{i}"] }]
+
+		update_values = update_keys.map { |k| "#{k}=#{fp_map[k]}" }.join(',')
+		select_values = unique_keys.map { |k| "#{k}=#{fp_map[k]}" }.join(' AND ')
+		update_query = "UPDATE #{tbl} SET #{update_values} WHERE #{select_values} RETURNING *"
+
+		insert_query = "INSERT INTO #{tbl} (#{fp_map.keys.join(',')}) SELECT #{fp_map.values.join(',')}"
+
+		["WITH upsert AS (#{update_query}) #{insert_query} WHERE NOT EXISTS (SELECT * FROM upsert)",
+		 data.values
+		]
+	end
+
+	# Run an upsert query.
+	#
+	# @see #upsert_sql
+	#
+	# Apply an upsert (update-or-insert) against a given database connection or
+	# connection pool, handling the (rarely needed) unique violation that can
+	# result.
+	#
+	# @param db [PG::EM::Client, PG::EM::ConnectionPool] the connection
+	#   against which all database operations will be run.
+	#
+	# @!macro upsert_params
+	#
+	# @return [EM::Deferrable] the deferrable in which the query is being
+	#   called; this means you should attach the code to run after the query
+	#   completes with `#callback`, and you can attach an error handler with
+	#   `#errback` if you like.
+	#
+	def db_upsert(db, tbl, key, data)
+		q = upsert_sql(tbl, key, data)
+
+		::EM::DefaultDeferrable.new.tap do |df|
+			db.exec_defer(*q).callback do
+				df.succeed
+			end.errback do |ex|
+				if ex.is_a?(PG::UniqueViolation)
+					db.exec_defer(*q).callback do
+						df.succeed
+					end.errback do |ex|
+						df.fail(ex)
+					end
+				else
+					df.fail(ex)
+				end
+			end
+		end
+	end
+
 	# Execute code in a transaction.
 	#
 	# Calling this method opens up a transaction (by executing `BEGIN`), and

data/lib/em-pg-client-helper/transaction.rb

@@ -117,6 +117,19 @@ class PG::EM::Client::Helper::Transaction
 		exec(*insert_sql(tbl, params), &blk)
 	end
 
+	# Run an upsert inside a transaction.
+	#
+	# @see {PG::EM::Client::Helper#upsert_sql} for all the parameters.
+	#
+	# @return [EM::Deferrable]
+	def upsert(*args, &blk)
+		db_upsert(@conn, *args).tap do |df|
+			df.callback(&blk) if block_given?
+		end.errback do |ex|
+			rollback(ex)
+		end
+	end
+
 	# Execute an arbitrary block of SQL in `sql` within the transaction.
 	# If you need to pass dynamic values to the query, those should be
 	# given in `values`, and referenced in `sql` as `$1`, `$2`, etc.  The

data/spec/db_transaction_spec.rb

@@ -3,36 +3,6 @@ require_relative './spec_helper'
 describe "PG::EM::Client::Helper#db_transaction" do
 	let(:mock_conn) { double(PG::EM::Client) }
 
-	def expect_query_failure(q, args=[], err=nil, exec_time = 0.001)
-		err ||= RuntimeError.new("Dummy failure")
-		expect_query(q, args, exec_time, :fail, err)
-	end
-
-	def expect_query(q, args=[], exec_time = 0.001, disposition = :succeed, *disp_opts)
-		df = EM::DefaultDeferrable.new
-
-		expect(mock_conn)
-		  .to receive(:exec_defer)
-		  .with(*[q, args].compact)
-		  .and_return(df)
-		  .ordered
-
-		EM.add_timer(exec_time) do
-			df.__send__(disposition, *disp_opts)
-		end
-	end
-
-	def in_transaction(*args, &blk)
-		db_transaction(mock_conn, *args, &blk).callback { EM.stop }.errback { EM.stop }
-	end
-
-	def in_em
-		EM.run do
-			EM.add_timer(0.5) { EM.stop; raise "test timeout" }
-			yield
-		end
-	end
-
 	it "runs a BEGIN/COMMIT cycle by default" do
 		in_em do
 			expect_query("BEGIN")

data/spec/db_transaction_upsert_spec.rb

@@ -0,0 +1,87 @@
+require_relative './spec_helper'
+
+describe "PG::EM::Client::Helper#db_transaction#upsert" do
+	let(:mock_conn) { double(PG::EM::Client) }
+	let(:mock_query) do
+		[
+		 'WITH upsert AS (UPDATE "foo" SET "bar"=$1 WHERE "wombat"=$2 RETURNING *) ' +
+		 'INSERT INTO "foo" ("bar","wombat") SELECT $1,$2 ' +
+		 'WHERE NOT EXISTS (SELECT * FROM upsert)',
+		 ["baz", 42]
+		]
+	end
+	let(:mock_query_opts) do
+		["foo", :wombat, { :bar => "baz", :wombat => 42 }]
+	end
+
+	it "runs a simple UPSERT correctly" do
+		in_em do
+			expect_query("BEGIN")
+			expect_query(*mock_query)
+			expect_query("COMMIT")
+			in_transaction do |txn|
+				txn.upsert(*mock_query_opts) do
+					txn.commit
+				end
+			end
+		end
+	end
+
+	it "rolls back after a failed attempt" do
+		in_em do
+			expect_query("BEGIN")
+			expect_query_failure(*mock_query)
+			expect_query("ROLLBACK")
+			in_transaction do |txn|
+				txn.upsert(*mock_query_opts) do
+					txn.commit
+				end
+			end
+		end
+	end
+
+	it "tries again if the first failure was a PG::UniqueViolation" do
+		in_em do
+			expect_query("BEGIN")
+			expect_query_failure(*mock_query, PG::UniqueViolation.new("OMFG"))
+			expect_query(*mock_query)
+			expect_query("COMMIT")
+
+			in_transaction do |txn|
+				txn.upsert(*mock_query_opts) do
+					txn.commit
+				end
+			end
+		end
+	end
+
+	it "explodes if the second time fails" do
+		in_em do
+			expect_query("BEGIN")
+			expect_query_failure(*mock_query, PG::UniqueViolation.new("OMFG"))
+			expect_query_failure(*mock_query)
+			expect_query("ROLLBACK")
+
+			in_transaction do |txn|
+				txn.upsert(*mock_query_opts) do
+					txn.commit
+				end
+			end
+		end
+	end
+
+	it "explodes if the second failure was a PG::UniqueViolation" do
+		in_em do
+			expect_query("BEGIN")
+			expect_query_failure(*mock_query, PG::UniqueViolation.new("OMFG"))
+			expect_query_failure(*mock_query, PG::UniqueViolation.new("OMFG"))
+			expect_query("ROLLBACK")
+
+			in_transaction do |txn|
+				txn.upsert(*mock_query_opts) do
+					txn.commit
+				end
+			end
+		end
+	end
+end

data/spec/db_upsert_spec.rb

@@ -0,0 +1,58 @@
+require_relative './spec_helper'
+
+describe "PG::EM::Client::Helper#db_upsert" do
+	let(:mock_query) do
+		[
+		 'WITH upsert AS (UPDATE "foo" SET "bar"=$1 WHERE "wombat"=$2 RETURNING *) ' +
+		 'INSERT INTO "foo" ("bar","wombat") SELECT $1,$2 ' +
+		 'WHERE NOT EXISTS (SELECT * FROM upsert)',
+		 ["baz", 42]
+		]
+	end
+
+	it "works first time" do
+		expect(db = double).
+		  to receive(:exec_defer).
+		  with(*mock_query).
+		  and_return(::EM::DefaultDeferrable.new.tap { |df| df.succeed })
+
+		df = db_upsert(db, "foo", :wombat, :bar => "baz", :wombat => 42)
+		expect(df.instance_variable_get(:@deferred_status)).to eq(:succeeded)
+	end
+
+	it "works on retry" do
+		expect(db = double).
+		  to receive(:exec_defer).
+		  with(*mock_query).
+		  ordered.
+		  and_return(
+		    ::EM::DefaultDeferrable.new.tap do |df|
+		      df.fail(PG::UniqueViolation.new("OMFG"))
+		    end
+		  )
+
+		expect(db).
+		  to receive(:exec_defer).
+		  with(*mock_query).
+		  ordered.
+		  and_return(::EM::DefaultDeferrable.new.tap { |df| df.succeed })
+
+		df = db_upsert(db, "foo", :wombat, :bar => "baz", :wombat => 42)
+		expect(df.instance_variable_get(:@deferred_status)).to eq(:succeeded)
+	end
+
+	it "fails on a different error" do
+		expect(db = double).
+		  to receive(:exec_defer).
+		  with(*mock_query).
+		  ordered.
+		  and_return(
+		    ::EM::DefaultDeferrable.new.tap do |df|
+		      df.fail(PG::UndefinedTable.new("OMFG"))
+		    end
+		  )
+
+		df = db_upsert(db, "foo", :wombat, :bar => "baz", :wombat => 42)
+		expect(df.instance_variable_get(:@deferred_status)).to eq(:failed)
+	end
+end

data/spec/spec_helper.rb

@@ -5,6 +5,7 @@ Spork.prefork do
 	Bundler.setup(:default, :development, :test)
 	require 'rspec/core'
 	require 'rspec/mocks'
+	require 'txn_helper'
 
 	require 'pry'
 
@@ -15,6 +16,8 @@ Spork.prefork do
 		config.expect_with :rspec do |c|
 			c.syntax = :expect
 		end
+
+		config.include TxnHelper
 	end
 end
 

data/spec/txn_helper.rb

@@ -0,0 +1,31 @@
+module TxnHelper
+	def expect_query_failure(q, args=[], err=nil, exec_time = 0.001)
+		err ||= RuntimeError.new("A mock exec_defer failure")
+		expect_query(q, args, exec_time, :fail, err)
+	end
+
+	def expect_query(q, args=[], exec_time = 0.001, disposition = :succeed, *disp_opts)
+		df = EM::DefaultDeferrable.new
+
+		expect(mock_conn)
+		  .to receive(:exec_defer)
+		  .with(*[q, args].compact)
+		  .and_return(df)
+		  .ordered
+
+		EM.add_timer(exec_time) do
+			df.__send__(disposition, *disp_opts)
+		end
+	end
+
+	def in_transaction(*args, &blk)
+		db_transaction(mock_conn, *args, &blk).callback { EM.stop }.errback { EM.stop }
+	end
+
+	def in_em
+		EM.run do
+			EM.add_timer(0.5) { EM.stop; raise "test timeout" }
+			yield
+		end
+	end
+end

data/spec/upsert_sql_spec.rb

@@ -0,0 +1,32 @@
+require_relative './spec_helper'
+
+describe "PG::EM::Client::Helper#upsert_sql" do
+	it "assembles a simple query correctly" do
+		expect(upsert_sql("foo", :wombat, :bar => "baz", :wombat => 42)).
+		  to eq([
+		    'WITH upsert AS (UPDATE "foo" SET "bar"=$1 WHERE "wombat"=$2 RETURNING *) ' +
+		    'INSERT INTO "foo" ("bar","wombat") SELECT $1,$2 ' +
+		    'WHERE NOT EXISTS (SELECT * FROM upsert)',
+		    ["baz", 42]
+		  ])
+	end
+
+	it "assembles a more complex query correctly" do
+		expect(
+		  upsert_sql(
+		    "user_store_data",
+		    [:id, :platform, :profile_type],
+		    :id           => 42,
+		    :platform     => 'wooden',
+		    :profile_type => 'xyzzy',
+		    :data         => "ohai!"
+		  )
+		).to eq([
+		    'WITH upsert AS (UPDATE "user_store_data" SET "data"=$4 ' +
+		    'WHERE "id"=$1 AND "platform"=$2 AND "profile_type"=$3 RETURNING *) ' +
+		    'INSERT INTO "user_store_data" ("id","platform","profile_type","data") ' +
+		    'SELECT $1,$2,$3,$4 WHERE NOT EXISTS (SELECT * FROM upsert)',
+		    [42, "wooden", "xyzzy", "ohai!"]
+		  ])
+	end
+end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: em-pg-client-helper
 version: !ruby/object:Gem::Version
-  version: 1.0.0
+  version: 1.1.0
   prerelease: 
 platform: ruby
 authors:
@@ -9,7 +9,7 @@ authors:
 autorequire: 
 bindir: bin
 cert_chain: []
-date: 2015-01-22 00:00:00.000000000 Z
+date: 2015-01-23 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: git-version-bump
@@ -245,8 +245,12 @@ files:
 - lib/em-pg-client-helper/transaction.rb
 - spec/db_insert_spec.rb
 - spec/db_transaction_spec.rb
+- spec/db_transaction_upsert_spec.rb
+- spec/db_upsert_spec.rb
 - spec/insert_sql_spec.rb
 - spec/spec_helper.rb
+- spec/txn_helper.rb
+- spec/upsert_sql_spec.rb
 homepage: http://github.com/mpalmer/em-pg-client-helper
 licenses: []
 post_install_message: 
@@ -261,7 +265,7 @@ required_ruby_version: !ruby/object:Gem::Requirement
       version: '0'
       segments:
       - 0
-      hash: 4034289824364546696
+      hash: 3057886313740167473
 required_rubygems_version: !ruby/object:Gem::Requirement
   none: false
   requirements:
@@ -270,7 +274,7 @@ required_rubygems_version: !ruby/object:Gem::Requirement
       version: '0'
       segments:
       - 0
-      hash: 4034289824364546696
+      hash: 3057886313740167473
 requirements: []
 rubyforge_project: 
 rubygems_version: 1.8.23