lockable.rb in good_job-2.9.5

- old
+ new

@@ -146,10 +146,11 @@
       #   end
       def with_advisory_lock(column: _advisory_lockable_column, function: advisory_lockable_function, unlock_session: false)
         raise ArgumentError, "Must provide a block" unless block_given?
 
         records = advisory_lock(column: column, function: function).to_a
+
         begin
           unscoped { yield(records) }
         ensure
           if unlock_session
             advisory_unlock_session
@@ -159,10 +160,57 @@
             end
           end
         end
       end
 
+      # Acquires an advisory lock on this record if it is not already locked by
+      # another database session. Be careful to ensure you release the lock when
+      # you are done with {#advisory_unlock_key} to release all remaining locks.
+      # @param key [String, Symbol] Key to Advisory Lock against
+      # @param function [String, Symbol] Postgres Advisory Lock function name to use
+      # @return [Boolean] whether the lock was acquired.
+      def advisory_lock_key(key, function: advisory_lockable_function)
+        query = if function.include? "_try_"
+                  <<~SQL.squish
+                    SELECT #{function}(('x'||substr(md5($1::text), 1, 16))::bit(64)::bigint) AS locked
+                  SQL
+                else
+                  <<~SQL.squish
+                    SELECT #{function}(('x'||substr(md5($1::text), 1, 16))::bit(64)::bigint)::text AS locked
+                  SQL
+                end
+
+        binds = [
+          ActiveRecord::Relation::QueryAttribute.new('key', key, ActiveRecord::Type::String.new),
+        ]
+        locked = connection.exec_query(pg_or_jdbc_query(query), 'GoodJob::Lockable Advisory Lock', binds).first['locked']
+        return locked unless block_given?
+        return nil unless locked
+
+        begin
+          yield
+        ensure
+          advisory_unlock_key(key, function: advisory_unlockable_function(function))
+        end
+      end
+
+      # Releases an advisory lock on this record if it is locked by this database
+      # session. Note that advisory locks stack, so you must call
+      # {#advisory_unlock} and {#advisory_lock} the same number of times.
+      # @param key [String, Symbol] Key to lock against
+      # @param function [String, Symbol] Postgres Advisory Lock function name to use
+      # @return [Boolean] whether the lock was released.
+      def advisory_unlock_key(key, function: advisory_unlockable_function)
+        query = <<~SQL.squish
+          SELECT #{function}(('x'||substr(md5($1::text), 1, 16))::bit(64)::bigint) AS unlocked
+        SQL
+        binds = [
+          ActiveRecord::Relation::QueryAttribute.new('key', key, ActiveRecord::Type::String.new),
+        ]
+        connection.exec_query(pg_or_jdbc_query(query), 'GoodJob::Lockable Advisory Unlock', binds).first['unlocked']
+      end
+
       def _advisory_lockable_column
         advisory_lockable_column || primary_key
       end
 
       def supports_cte_materialization_specifiers?
@@ -203,52 +251,32 @@
     # all remaining locks).
     # @param key [String, Symbol] Key to Advisory Lock against
     # @param function [String, Symbol] Postgres Advisory Lock function name to use
     # @return [Boolean] whether the lock was acquired.
     def advisory_lock(key: lockable_key, function: advisory_lockable_function)
-      query = if function.include? "_try_"
-                <<~SQL.squish
-                  SELECT #{function}(('x'||substr(md5($1::text), 1, 16))::bit(64)::bigint) AS locked
-                SQL
-              else
-                <<~SQL.squish
-                  SELECT #{function}(('x'||substr(md5($1::text), 1, 16))::bit(64)::bigint)::text AS locked
-                SQL
-              end
-
-      binds = [
-        ActiveRecord::Relation::QueryAttribute.new('key', key, ActiveRecord::Type::String.new),
-      ]
-      self.class.connection.exec_query(pg_or_jdbc_query(query), 'GoodJob::Lockable Advisory Lock', binds).first['locked']
+      self.class.advisory_lock_key(key, function: function)
     end
 
     # Releases an advisory lock on this record if it is locked by this database
     # session. Note that advisory locks stack, so you must call
     # {#advisory_unlock} and {#advisory_lock} the same number of times.
     # @param key [String, Symbol] Key to lock against
     # @param function [String, Symbol] Postgres Advisory Lock function name to use
     # @return [Boolean] whether the lock was released.
     def advisory_unlock(key: lockable_key, function: self.class.advisory_unlockable_function(advisory_lockable_function))
-      query = <<~SQL.squish
-        SELECT #{function}(('x'||substr(md5($1::text), 1, 16))::bit(64)::bigint) AS unlocked
-      SQL
-      binds = [
-        ActiveRecord::Relation::QueryAttribute.new('key', key, ActiveRecord::Type::String.new),
-      ]
-      self.class.connection.exec_query(pg_or_jdbc_query(query), 'GoodJob::Lockable Advisory Unlock', binds).first['unlocked']
+      self.class.advisory_unlock_key(key, function: function)
     end
 
     # Acquires an advisory lock on this record or raises
     # {RecordAlreadyAdvisoryLockedError} if it is already locked by another
     # database session.
     # @param key [String, Symbol] Key to lock against
     # @param function [String, Symbol] Postgres Advisory Lock function name to use
     # @raise [RecordAlreadyAdvisoryLockedError]
     # @return [Boolean] +true+
     def advisory_lock!(key: lockable_key, function: advisory_lockable_function)
-      result = advisory_lock(key: key, function: function)
-      result || raise(RecordAlreadyAdvisoryLockedError)
+      self.class.advisory_lock_key(key, function: function) || raise(RecordAlreadyAdvisoryLockedError)
     end
 
     # Acquires an advisory lock on this record and safely releases it after the
     # passed block is completed. If the record is locked by another database
     # session, this raises {RecordAlreadyAdvisoryLockedError}.
@@ -264,12 +292,14 @@
     #   end
     def with_advisory_lock(key: lockable_key, function: advisory_lockable_function)
       raise ArgumentError, "Must provide a block" unless block_given?
 
       advisory_lock!(key: key, function: function)
-      yield
-    ensure
-      advisory_unlock(key: key, function: self.class.advisory_unlockable_function(function)) unless $ERROR_INFO.is_a? RecordAlreadyAdvisoryLockedError
+      begin
+        yield
+      ensure
+        advisory_unlock(key: key, function: self.class.advisory_unlockable_function(function))
+      end
     end
 
     # Tests whether this record has an advisory lock on it.
     # @param key [String, Symbol] Key to test lock against
     # @return [Boolean]