Skip to content

Commit

Permalink
chore: refractor postgresql class
Browse files Browse the repository at this point in the history
  • Loading branch information
seuros committed Feb 12, 2024
1 parent a078c31 commit d6d41c8
Show file tree
Hide file tree
Showing 2 changed files with 30 additions and 23 deletions.
4 changes: 1 addition & 3 deletions lib/with_advisory_lock/base.rb
Original file line number Diff line number Diff line change
Expand Up @@ -56,9 +56,7 @@ def already_locked?
def with_advisory_lock_if_needed(&block)
if disable_query_cache
return lock_and_yield do
connection.uncached do
yield
end
connection.uncached(&block)
end
end

Expand Down
49 changes: 29 additions & 20 deletions lib/with_advisory_lock/postgresql.rb
Original file line number Diff line number Diff line change
@@ -1,43 +1,52 @@
# frozen_string_literal: true

module WithAdvisoryLock
class PostgreSQL < Base
# See http://www.postgresql.org/docs/9.1/static/functions-admin.html#FUNCTIONS-ADVISORY-LOCKS
LOCK_RESULT_VALUES = ['t', true].freeze
LOCK_ENV_NAME = 'WITH_ADVISORY_LOCK_PREFIX'.freeze
PG_ADVISORY_UNLOCK = 'pg_advisory_unlock'.freeze
PG_TRY_ADVISORY = 'try_advisory'.freeze
ERROR_MESSAGE_REGEX = / ERROR: +current transaction is aborted,/

def try_lock
pg_function = "pg_try_advisory#{transaction ? '_xact' : ''}_lock#{shared ? '_shared' : ''}"
execute_successful?(pg_function)
execute_successful?(pg_function_name(PG_TRY_ADVISORY, transaction))
end

def release_lock
return if transaction

pg_function = "pg_advisory_unlock#{shared ? '_shared' : ''}"
execute_successful?(pg_function)
rescue ActiveRecord::StatementInvalid => e
raise unless e.message =~ / ERROR: +current transaction is aborted,/

begin
execute_successful?(pg_function_name(PG_ADVISORY_UNLOCK, false))
rescue ActiveRecord::StatementInvalid => e
raise unless e.message =~ ERROR_MESSAGE_REGEX
connection.rollback_db_transaction
execute_successful?(pg_function)
retry
ensure
connection.begin_db_transaction
end
end

def pg_function_name(base_name, transaction_scope)
[
base_name,
transaction_scope ? '_xact' : nil,
shared ? '_shared' : nil
].compact.join
end

def execute_successful?(pg_function)
result = connection.select_value(prepare_sql(pg_function))
LOCK_RESULT_VALUES.include?(result)
end

def prepare_sql(pg_function)
comment = lock_name.to_s.gsub(%r{(/\*)|(\*/)}, '--')
sql = "SELECT #{pg_function}(#{lock_keys.join(',')}) AS #{unique_column_name} /* #{comment} */"
result = connection.select_value(sql)
# MRI returns 't', jruby returns true. YAY!
['t', true].include?(result)
"SELECT #{pg_function}(#{lock_keys.join(',')}) AS #{unique_column_name} /* #{comment} */"
end

# PostgreSQL wants 2 32bit integers as the lock key.
def lock_keys
@lock_keys ||= [stable_hashcode(lock_name), ENV['WITH_ADVISORY_LOCK_PREFIX']].map do |ea|
# pg advisory args must be 31 bit ints
ea.to_i & 0x7fffffff
end
@lock_keys ||= [
stable_hashcode(lock_name),
ENV[LOCK_ENV_NAME]
].map { |ea| ea.to_i & 0x7fffffff }
end
end
end

0 comments on commit d6d41c8

Please sign in to comment.