From a1d4f8511b6d2d8167dadd8a4b7aeaee5ade07c7 Mon Sep 17 00:00:00 2001
From: KJ Tsanaktsidis <ktsanaktsidis@zendesk.com>
Date: Wed, 5 Apr 2023 11:51:38 +1000
Subject: [PATCH 1/2] Cache ActiveRecordShards.app_env

It can't change (except during tests), so cache it.
---
 lib/active_record_shards.rb       | 14 ++++++++------
 test/active_record_shards_test.rb |  6 +++++-
 test/helper.rb                    |  2 ++
 3 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/lib/active_record_shards.rb b/lib/active_record_shards.rb
index 8fa7b2c5..16d88449 100644
--- a/lib/active_record_shards.rb
+++ b/lib/active_record_shards.rb
@@ -18,12 +18,14 @@ class << self
   end
 
   def self.app_env
-    env = Rails.env if defined?(Rails.env)
-    env ||= RAILS_ENV if Object.const_defined?(:RAILS_ENV)
-    env ||= ENV['RAILS_ENV']
-    env ||= APP_ENV if Object.const_defined?(:APP_ENV)
-    env ||= ENV['APP_ENV']
-    env || 'development'
+    @app_env ||= begin
+      env = Rails.env if defined?(Rails.env)
+      env ||= RAILS_ENV if Object.const_defined?(:RAILS_ENV)
+      env ||= ENV['RAILS_ENV']
+      env ||= APP_ENV if Object.const_defined?(:APP_ENV)
+      env ||= ENV['APP_ENV']
+      env || 'development'
+    end
   end
 end
 
diff --git a/test/active_record_shards_test.rb b/test/active_record_shards_test.rb
index 8853afd6..a9892d2a 100644
--- a/test/active_record_shards_test.rb
+++ b/test/active_record_shards_test.rb
@@ -10,9 +10,13 @@
       end
 
       Object.send(:remove_const, 'RAILS_ENV')
+      ActiveRecordShards.instance_eval { @app_env = nil }
     end
 
-    after { Object.const_set('RAILS_ENV', 'test') }
+    after do
+      Object.const_set('RAILS_ENV', 'test')
+      ActiveRecordShards.instance_eval { @app_env = nil }
+    end
 
     describe 'Rails.env' do
       before do
diff --git a/test/helper.rb b/test/helper.rb
index e6896c0f..950dd28c 100644
--- a/test/helper.rb
+++ b/test/helper.rb
@@ -186,10 +186,12 @@ module RailsEnvSwitch
   def switch_app_env(env)
     before do
       silence_warnings { Object.const_set("RAILS_ENV", env) }
+      ActiveRecordShards.instance_eval { @app_env = nil }
       ActiveRecord::Base.establish_connection(::RAILS_ENV.to_sym)
     end
     after do
       silence_warnings { Object.const_set("RAILS_ENV", 'test') }
+      ActiveRecordShards.instance_eval { @app_env = nil }
       ActiveRecord::Base.establish_connection(::RAILS_ENV.to_sym)
       tmp_sharded_model = Class.new(ActiveRecord::Base)
       assert_equal('ars_test', tmp_sharded_model.connection.current_database)

From bfc99f0d4ce2556ba25b5fa5340581553cdbf1bd Mon Sep 17 00:00:00 2001
From: KJ Tsanaktsidis <ktsanaktsidis@zendesk.com>
Date: Wed, 5 Apr 2023 15:14:02 +1000
Subject: [PATCH 2/2] Cache the is_sharded? status of the model

Sets self.sharded depending on the superclass, if it is not otherwise
set.
---
 lib/active_record_shards/model.rb | 31 ++++++++++++++++++++-----------
 1 file changed, 20 insertions(+), 11 deletions(-)

diff --git a/lib/active_record_shards/model.rb b/lib/active_record_shards/model.rb
index 7d3c6581..e257035f 100644
--- a/lib/active_record_shards/model.rb
+++ b/lib/active_record_shards/model.rb
@@ -11,17 +11,26 @@ def not_sharded
     end
 
     def is_sharded? # rubocop:disable Naming/PredicateName
-      if self == ActiveRecord::Base
-        sharded != false && supports_sharding?
-      elsif self == base_class
-        if sharded.nil?
-          ActiveRecord::Base.is_sharded?
-        else
-          sharded != false
-        end
-      else
-        base_class.is_sharded?
-      end
+      # "sharded" here means self.sharded, but actually writing "self.sharded"
+      # doesn't work until Ruby 2.7 (and this gem currently supports 2.6) because
+      # the sharded attr_accessor is private. Private methods must be called without
+      # a receiver, but Ruby 2.7+ does allow an explicit "self" as a receiver.
+      return sharded unless sharded.nil?
+
+      # Despite self.sharded not working, self.sharded= _DOES_ work. That's an exception
+      # to the "private methods must be called with no receiver" rule (presumably
+      # because it would otherwise be ambiguous with local variable assignment).
+      self.sharded = if self == ActiveRecord::Base
+                       sharded != false && supports_sharding?
+                     elsif self == base_class
+                       if sharded.nil?
+                         ActiveRecord::Base.is_sharded?
+                       else
+                         sharded != false
+                       end
+                     else
+                       base_class.is_sharded?
+                     end
     end
 
     def on_replica_by_default?