Added PgpoolNoLoadBalance.force to allow enabling pgpool_nlb inside blocks

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.

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."

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]

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