Merge pull request #311 from zendesk/ktsanaktsidis/optimise_is_sharded

Cache the result of is_sharded?
zendesk · Apr 6, 2023 · 5669646 · 5669646
2 parents 14d4f62 + bfc99f0
commit 5669646
Show file tree

Hide file tree

Showing 4 changed files with 35 additions and 18 deletions.
diff --git 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/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?

diff --git 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
@@ -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)