wip: patch for lease_connection

aktsk · Sep 3, 2024 · fc17846 · fc17846
1 parent b73827d
commit fc17846
Showing 1 changed file with 10 additions and 2 deletions.
diff --git a/lib/octoball/connection_adapters.rb b/lib/octoball/connection_adapters.rb
@@ -26,15 +26,23 @@ def with_connection(prevent_permanent_checkout: false)
 
       def active_connection?
         conn = connection_lease.connection
-        conn.current_shard = conn.shard
+        conn.current_shard = conn.shard if conn
         conn
       end
 
       def active_connection
         conn = connection_lease.connection
-        conn.current_shard = conn.shard
+        conn.current_shard = conn.shard if conn
         conn
       end
+
+      def lease_connection
+        lease = connection_lease
+        lease.sticky = true
+        lease.connection ||= checkout
+        lease.connection.current_shard = lease.connection.shard
+        lease.connection
+      end
     end
 
     ::ActiveRecord::ConnectionAdapters::ConnectionPool.prepend(ConnectionPoolSetCurrentShard)