From 5b1ed5a1209fff89eeee39cc261c0c01a6e897e6 Mon Sep 17 00:00:00 2001
From: Abdelkader Boudih <terminale@gmail.com>
Date: Mon, 12 Feb 2024 03:40:08 +0100
Subject: [PATCH] chore: refractor postgresql class (#99)

---
 lib/with_advisory_lock.rb            |  1 +
 lib/with_advisory_lock/base.rb       |  6 ++--
 lib/with_advisory_lock/postgresql.rb | 53 ++++++++++++++++++++--------
 3 files changed, 41 insertions(+), 19 deletions(-)

diff --git a/lib/with_advisory_lock.rb b/lib/with_advisory_lock.rb
index 75381e2..a3b4295 100644
--- a/lib/with_advisory_lock.rb
+++ b/lib/with_advisory_lock.rb
@@ -10,6 +10,7 @@
 loader.setup
 
 module WithAdvisoryLock
+  LOCK_PREFIX_ENV = 'WITH_ADVISORY_LOCK_PREFIX'.freeze
 end
 
 ActiveSupport.on_load :active_record do
diff --git a/lib/with_advisory_lock/base.rb b/lib/with_advisory_lock/base.rb
index 35e5d2b..c9dc486 100644
--- a/lib/with_advisory_lock/base.rb
+++ b/lib/with_advisory_lock/base.rb
@@ -36,7 +36,7 @@ def initialize(connection, lock_name, options)
     end
 
     def lock_str
-      @lock_str ||= "#{ENV['WITH_ADVISORY_LOCK_PREFIX']}#{lock_name}"
+      @lock_str ||= "#{ENV[LOCK_PREFIX_ENV]}#{lock_name}"
     end
 
     def lock_stack_item
@@ -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
 
diff --git a/lib/with_advisory_lock/postgresql.rb b/lib/with_advisory_lock/postgresql.rb
index ae369cb..9ec31d6 100644
--- a/lib/with_advisory_lock/postgresql.rb
+++ b/lib/with_advisory_lock/postgresql.rb
@@ -2,42 +2,65 @@
 
 module WithAdvisoryLock
   class PostgreSQL < Base
-    # See http://www.postgresql.org/docs/9.1/static/functions-admin.html#FUNCTIONS-ADVISORY-LOCKS
+    # See https://www.postgresql.org/docs/16/functions-admin.html#FUNCTIONS-ADVISORY-LOCKS
+
+    # MRI returns 't', jruby returns true. YAY!
+    LOCK_RESULT_VALUES = ['t', true].freeze
+    PG_ADVISORY_UNLOCK = 'pg_advisory_unlock'
+    PG_TRY_ADVISORY = 'pg_try_advisory'
+    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?(advisory_try_lock_function(transaction))
     end
 
     def release_lock
       return if transaction
 
-      pg_function = "pg_advisory_unlock#{shared ? '_shared' : ''}"
-      execute_successful?(pg_function)
+      execute_successful?(advisory_unlock_function)
     rescue ActiveRecord::StatementInvalid => e
-      raise unless e.message =~ / ERROR: +current transaction is aborted,/
+      raise unless e.message =~ ERROR_MESSAGE_REGEX
 
       begin
         connection.rollback_db_transaction
-        execute_successful?(pg_function)
+        execute_successful?(advisory_unlock_function)
       ensure
         connection.begin_db_transaction
       end
     end
 
+    def advisory_try_lock_function(transaction_scope)
+      [
+        'pg_try_advisory',
+        transaction_scope ? '_xact' : nil,
+        '_lock',
+        shared ? '_shared' : nil
+      ].compact.join
+    end
+
+    def advisory_unlock_function
+      [
+        'pg_advisory_unlock',
+        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_PREFIX_ENV]
+      ].map { |ea| ea.to_i & 0x7fffffff }
     end
   end
 end