From e85697e6c5a8f5dc989b629f8ab30244f97a1acf Mon Sep 17 00:00:00 2001
From: Bruno Campos <brunocampos@gmail.com>
Date: Mon, 2 Mar 2026 19:03:09 -0300
Subject: [PATCH] Added PgpoolNoLoadBalance.force to allow enabling pgpool_nlb
 inside blocks

---
 README.md                                             |  9 +++++++++
 lib/pgpool_no_load_balance.rb                         | 11 +++++++++++
 .../connection_adapters/postgresql_adapter.rb         |  2 +-
 lib/pgpool_no_load_balance/arel/select_manager.rb     |  2 +-
 4 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index 0f0a44d..48015e7 100644
--- a/README.md
+++ b/README.md
@@ -35,6 +35,15 @@ irb(main):001:0> User.pgpool_nlb.all
   /*NO LOAD BALANCE*/ SELECT "users".* FROM "users" LIMIT $1  [["LIMIT", 11]]
 ```
 
+### blocks
+
+Use `PgpoolNoLoadBalance.force` block to force all read queries to run with `pgpool_nlb` enabled inside the block
+
+```rb
+irb(main):001:0> PgpoolNoLoadBalance.force { User.all }
+  /*NO LOAD BALANCE*/ SELECT "users".* FROM "users" LIMIT $1  [["LIMIT", 11]]
+```
+
 ### unscope
 
 Can remove the scope with the unscope method.
diff --git a/lib/pgpool_no_load_balance.rb b/lib/pgpool_no_load_balance.rb
index a6a6694..a59eebb 100644
--- a/lib/pgpool_no_load_balance.rb
+++ b/lib/pgpool_no_load_balance.rb
@@ -11,6 +11,17 @@ module PgpoolNoLoadBalance
 
   class PostgreSQLAdapterMissing < StandardError; end
 
+  def self.force
+    Thread.current[:pgpool_nlb_force] = true
+    yield
+  ensure
+    Thread.current[:pgpool_nlb_force] = false
+  end
+
+  def self.force?
+    !!Thread.current[:pgpool_nlb_force]
+  end
+
   def self.setup!
     unless ::ActiveRecord::Base.respond_to?(:postgresql_connection)
       raise PostgreSQLAdapterMissing, "No postgresql adapter specified by 'config/database.yml', or 'ActiveRecord::Base.establish_connection' method is not called."
diff --git a/lib/pgpool_no_load_balance/active_record/connection_adapters/postgresql_adapter.rb b/lib/pgpool_no_load_balance/active_record/connection_adapters/postgresql_adapter.rb
index 6478184..5628ebe 100644
--- a/lib/pgpool_no_load_balance/active_record/connection_adapters/postgresql_adapter.rb
+++ b/lib/pgpool_no_load_balance/active_record/connection_adapters/postgresql_adapter.rb
@@ -11,7 +11,7 @@ def execute(sql, name = nil, pgpool_nlb: false)
 
         def to_sql_and_binds(arel_or_sql_string, binds = [], preparable = nil) # :nodoc:
           sql, binds, preparable = super
-          if arel_or_sql_string.respond_to?(:pgpool_nlb?) && arel_or_sql_string.pgpool_nlb?
+          if PgpoolNoLoadBalance.force? || (arel_or_sql_string.respond_to?(:pgpool_nlb?) && arel_or_sql_string.pgpool_nlb?)
             sql = "#{NLB_COMMENT} #{sql}"
           end
           [sql.freeze, binds, preparable]
diff --git a/lib/pgpool_no_load_balance/arel/select_manager.rb b/lib/pgpool_no_load_balance/arel/select_manager.rb
index fdad2f9..3284ee3 100644
--- a/lib/pgpool_no_load_balance/arel/select_manager.rb
+++ b/lib/pgpool_no_load_balance/arel/select_manager.rb
@@ -7,7 +7,7 @@ def pgpool_nlb(value = true)
       end
 
       def pgpool_nlb?
-        @pgpool_nlb_flag
+        @pgpool_nlb_flag || PgpoolNoLoadBalance.force?
       end
     end
   end