add nested with connection

lib/que/adapters/active_record_with_lock.rb

@@ -4,17 +4,41 @@
 module Que
   module Adapters
     class ActiveRecordWithLock < Que::Adapters::ActiveRecord
-      LOCK_PREFIX = ENV["QUE_LOCK_PREFIX"] || 1111 # this is a random number
       def initialize(job_connection_pool:, lock_record:)
         @job_connection_pool = job_connection_pool
         @lock_record = lock_record
         super
       end
 
+      def checkout
+        #  require 'debug'; binding.break
+        checkout_lock_database_connection do |connection|
+          # require 'debug'; binding.break
+          checkout_activerecord_adapter { |conn| yield conn.raw_connection }
+        end
+      rescue *AR_UNAVAILABLE_CONNECTION_ERRORS => e
+        puts "error #{e}"
+        raise UnavailableConnection, e
+      rescue ::ActiveRecord::StatementInvalid => e
+        puts "error #{e}"
+        raise e unless AR_UNAVAILABLE_CONNECTION_ERRORS.include?(e.cause.class)
+
+        # ActiveRecord::StatementInvalid is one of the most generic exceptions AR can
+        # raise, so we catch it and only handle the specific nested exceptions.
+        raise UnavailableConnection, e.cause
+      rescue StandardError => e
+        puts "error #{e}"
+        raise e
+      end
+
       def checkout_activerecord_adapter(&block)
         @job_connection_pool.with_connection(&block)
       end
 
+      def checkout_lock_database_connection(&block)
+        @lock_record.connection_pool.with_connection(&block)
+      end
+
       def lock_database_connection
         if Thread.current[:db_connection]
           return Thread.current[:db_connection] if Thread.current[:db_connection].active?
@@ -42,6 +66,7 @@ def execute(command, params = [])
       def lock_job_with_lock_database(queue, cursor)
         result = []
         loop do
+          # require 'debug'; binding.break
           result = Que.execute(:find_job_to_lock, [queue, cursor])
 
           break if result.empty?
@@ -53,26 +78,81 @@ def lock_job_with_lock_database(queue, cursor)
       end
 
       def cleanup!
+        puts "cleanup"
         @job_connection_pool.release_connection
         @lock_record.remove_connection
       end
 
       def pg_try_advisory_lock?(job_id)
-        lock_variable = "#{LOCK_PREFIX}#{job_id}".to_i
-        lock_database_connection.execute(
-          "SELECT pg_try_advisory_lock(#{lock_variable})",
-        ).try(:first)&.fetch("pg_try_advisory_lock")
+        checkout_lock_database_connection do |conn|
+          conn.execute(
+            "SELECT pg_try_advisory_lock(#{job_id})",
+          ).try(:first)&.fetch("pg_try_advisory_lock")
+        end
       end
 
       def unlock_job(job_id)
-        lock_variable = "#{LOCK_PREFIX}#{job_id}".to_i
         # If for any reason the connection that is used to get this advisory lock
         # is corrupted, the lock on this job_id would already be released when the
         # connection holding the lock goes bad.
         # Now, if a new connection tries to release the non existing lock this would just no op
         # by returning false and return a warning "WARNING:  you don't own a lock of type ExclusiveLock"
-        lock_database_connection.execute("SELECT pg_advisory_unlock(#{lock_variable})")
+        checkout_lock_database_connection do |conn|
+          conn.execute("SELECT pg_advisory_unlock(#{job_id})")
+        end
       end
     end
   end
 end
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+