From 256de6f389fff5070c693d1cb1458b3270b4209c Mon Sep 17 00:00:00 2001
From: Felipe Sateler <fsateler@gmail.com>
Date: Mon, 13 Apr 2020 14:47:22 -0400
Subject: [PATCH] Avoid getting a connection to the database before it's
 necessary

We init the first time we request a connection. On code reloads, we
"forget" about our init state and we call it again.

After init, ensure the search path is set properly
---
 lib/apartment/adapters/abstract_adapter.rb   |  6 +++++
 lib/apartment/adapters/postgresql_adapter.rb |  5 ++++
 lib/apartment/railtie.rb                     | 27 +++++++++-----------
 lib/apartment/tenant.rb                      | 19 ++++++++++----
 spec/examples/generic_adapter_examples.rb    |  2 +-
 5 files changed, 38 insertions(+), 21 deletions(-)

diff --git a/lib/apartment/adapters/abstract_adapter.rb b/lib/apartment/adapters/abstract_adapter.rb
index c9a548aa..b73a09ff 100644
--- a/lib/apartment/adapters/abstract_adapter.rb
+++ b/lib/apartment/adapters/abstract_adapter.rb
@@ -35,6 +35,12 @@ def create(tenant)
         end
       end
 
+      #   Initialize Apartment config options such as excluded_models
+      #
+      def init
+        process_excluded_models
+      end
+
       #   Note alias_method here doesn't work with inheritence apparently ??
       #
       def current
diff --git a/lib/apartment/adapters/postgresql_adapter.rb b/lib/apartment/adapters/postgresql_adapter.rb
index 3e98b627..6aa81d6d 100644
--- a/lib/apartment/adapters/postgresql_adapter.rb
+++ b/lib/apartment/adapters/postgresql_adapter.rb
@@ -39,6 +39,11 @@ def reset
         Apartment.connection.schema_search_path = full_search_path
       end
 
+      def init
+        super
+        Apartment.connection.schema_search_path = full_search_path
+      end
+
       def current
         @current || default_tenant
       end
diff --git a/lib/apartment/railtie.rb b/lib/apartment/railtie.rb
index 7ff83b1a..5f2eaf9d 100644
--- a/lib/apartment/railtie.rb
+++ b/lib/apartment/railtie.rb
@@ -24,26 +24,23 @@ class Railtie < Rails::Railtie
       ActiveRecord::Migrator.migrations_paths = Rails.application.paths['db/migrate'].to_a
     end
 
-    #   Hook into ActionDispatch::Reloader to ensure Apartment is properly initialized
-    #   Note that this doens't entirely work as expected in Development, because this is called before classes are reloaded
-    #   See the middleware/console declarations below to help with this. Hope to fix that soon.
-    #
+    # Make sure Apartment is reconfigured when the code is reloaded
     config.to_prepare do
-      next if ARGV.any? { |arg| arg =~ /\Aassets:(?:precompile|clean)\z/ }
-      next if ARGV.any? { |arg| arg == 'webpacker:compile' }
+      Apartment::Tenant.reinitialize
+    end
 
-      begin
-        Apartment.connection_class.connection_pool.with_connection do
-          Apartment::Tenant.init
-        end
-      rescue ::ActiveRecord::NoDatabaseError, PG::ConnectionBad
-        # Since `db:create` and other tasks invoke this block from Rails 5.2.0,
-        # we need to swallow the error to execute `db:create` properly.
-        Rails.logger.warn do
-          'Failed to initialize Apartment because a database connection could not be established.'
+    #
+    # Ensure that Apartment::Tenant.init is called when
+    # a new connection is requested.
+    #
+    module ApartmentInitializer
+      def connection
+        super.tap do
+          Apartment::Tenant.init_once
         end
       end
     end
+    ActiveRecord::Base.singleton_class.prepend ApartmentInitializer
 
     #
     #   Ensure rake tasks are loaded
diff --git a/lib/apartment/tenant.rb b/lib/apartment/tenant.rb
index bc69f9c9..17560b27 100644
--- a/lib/apartment/tenant.rb
+++ b/lib/apartment/tenant.rb
@@ -10,15 +10,23 @@ module Tenant
     extend Forwardable
 
     def_delegators :adapter, :create, :drop, :switch, :switch!, :current, :each,
-                   :reset, :set_callback, :seed, :current_tenant,
+                   :reset, :init, :set_callback, :seed, :current_tenant,
                    :default_tenant, :environmentify
 
     attr_writer :config
 
-    #   Initialize Apartment config options such as excluded_models
-    #
-    def init
-      adapter.process_excluded_models
+    def init_once
+      return if @already_initialized
+
+      # To avoid infinite loops in work init is doing,
+      # we need to set @already_initialized to true
+      # before init is called
+      @already_initialized = true
+      init
+    end
+
+    def reinitialize
+      @already_initialized = false
     end
 
     #   Fetch the proper multi-tenant adapter based on Rails config
@@ -53,6 +61,7 @@ def adapter
     #
     def reload!(config = nil)
       Thread.current[:apartment_adapter] = nil
+      reinitialize
       @config = config
     end
 
diff --git a/spec/examples/generic_adapter_examples.rb b/spec/examples/generic_adapter_examples.rb
index 252ac47c..514f8eb6 100644
--- a/spec/examples/generic_adapter_examples.rb
+++ b/spec/examples/generic_adapter_examples.rb
@@ -24,7 +24,7 @@
                                              .instance_variable_get(:@queue)
                                              .size
 
-        expect(num_available_connections).to eq(1)
+        expect(num_available_connections).to eq(0)
       end
     end
   end