pass the connection pool on initialise the adapter

gocardless · Jun 27, 2024 · 1fdcc56 · 1fdcc56
1 parent f413153
commit 1fdcc56
Show file tree

Hide file tree

Showing 5 changed files with 13 additions and 80 deletions.
diff --git a/lib/que.rb b/lib/que.rb
@@ -61,10 +61,14 @@ def connection=(connection)
       self.adapter =
         if connection.to_s == "ActiveRecord"
           Adapters::ActiveRecord.new
-        elsif connection.to_s == "Que::Adapters::Yugabyte"
-          Adapters::Yugabyte.new
         else
           case connection.class.to_s
+          when "Que::Adapters::ActiveRecordWithLock" then  
+
+            Adapters::ActiveRecordWithLock.new(
+              job_connection_pool: connection.job_connection_pool, 
+              lock_connection_pool: connection.lock_connection_pool
+            )
           when "Sequel::Postgres::Database" then Adapters::Sequel.new(connection)
           when "ConnectionPool"             then Adapters::ConnectionPool.new(connection)
           when "PG::Connection"             then Adapters::PG.new(connection)

diff --git a/lib/que/adapters/base.rb b/lib/que/adapters/base.rb
@@ -9,7 +9,7 @@ module Adapters
     autoload :PG,             "que/adapters/pg"
     autoload :Pond,           "que/adapters/pond"
     autoload :Sequel,         "que/adapters/sequel"
-    autoload :Yugabyte,        "que/adapters/yugabyte"
+    autoload :ActiveRecordWithLock,        "que/adapters/active_record_with_lock"
 
     class UnavailableConnection < StandardError; end
 

diff --git a/lib/que/adapters/yugabyte.rb b/lib/que/adapters/yugabyte.rb
diff --git a/lib/que/locker.rb b/lib/que/locker.rb
@@ -153,7 +153,7 @@ def lock_job_query(queue, cursor)
 
     def handle_expired_cursors!
       @consolidated_queues.each do |queue|
-        queue_cursor_expires_at = @queue_expires_at.fetch(queue, monotonic_now)       
+        queue_cursor_expires_at = @queue_expires_at.fetch(queue, monotonic_now)  
         reset_cursor_for!(queue) if queue_cursor_expires_at < monotonic_now
       end
     end

diff --git a/spec/spec_helper.rb b/spec/spec_helper.rb
@@ -13,7 +13,7 @@
 require_relative "./helpers/sleep_job"
 require_relative "./helpers/interruptible_sleep_job"
 require_relative "./helpers/user"
-require_relative "../lib/que/adapters/yugabyte"
+
 
 def postgres_now
   ActiveRecord::Base.connection.execute("SELECT NOW();")[0]["now"]
@@ -44,23 +44,21 @@ class LockDatabaseRecord < ActiveRecord::Base
 end
 
 class YugabyteRecord < ActiveRecord::Base
-  def self.establish_lock_database_connection
     establish_connection(
       adapter: "postgresql",
       host: ENV.fetch("PGHOST", "localhost"),
       user: ENV.fetch("PGUSER", "ubuntu"),
       password: ENV.fetch("PGPASSWORD", "password"),
       database: ENV.fetch("PGDATABASE", "que-test"),
     )
-  end
-  def self.connection
-    establish_lock_database_connection.connection
-  end
 end
 
 # Make sure our test database is prepared to run Que
 if ENV['YUGABYTE_QUE_WORKER_ENABLED']
-  Que.connection = Que::Adapters::Yugabyte
+  Que.connection = Que::Adapters::ActiveRecordWithLock.new(
+                job_connection_pool: YugabyteRecord.connection_pool,
+                lock_connection_pool: LockDatabaseRecord.connection_pool,
+  )
 else
   Que.connection = ActiveRecord
 end
@@ -91,7 +89,6 @@ def self.connection
   end
 
   config.before do
-    LockDatabaseRecord.connection_pool.release_connection
     QueJob.delete_all
     FakeJob.log = []
     ExceptionalJob.log = []