Merge pull request rails#47972 from Shopify/ar-get-version-schema-cache

Active Record: assign connection pool before checking version

Author: byroot

activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb

@@ -684,9 +684,10 @@ def remove_connection_from_thread_cache(conn, owner_thread = conn.owner)
         alias_method :release, :remove_connection_from_thread_cache
 
         def new_connection
-          Base.public_send(db_config.adapter_method, db_config.configuration_hash).tap do |conn|
-            conn.check_version
-          end
+          connection = Base.public_send(db_config.adapter_method, db_config.configuration_hash)
+          connection.pool = self
+          connection.check_version
+          connection
         end
 
         # If the pool is not at a <tt>@size</tt> limit, establish new connection. Connecting

activerecord/test/cases/connection_adapters/connection_handler_test.rb

@@ -354,9 +354,9 @@ def test_forked_child_doesnt_mangle_parent_connection
 
           pid = fork {
             rd.close
-            if ActiveRecord::Base.connection.active?
-              wr.write Marshal.dump ActiveRecord::Base.connection.object_id
-            end
+            wr.write Marshal.dump [
+              ActiveRecord::Base.connection.object_id,
+            ]
             wr.close
 
             exit # allow finalizers to run
@@ -365,7 +365,8 @@ def test_forked_child_doesnt_mangle_parent_connection
           wr.close
 
           Process.waitpid pid
-          assert_not_equal object_id, Marshal.load(rd.read)
+          child_id = Marshal.load(rd.read)
+          assert_not_equal object_id, child_id
           rd.close
 
           assert_equal 3, ActiveRecord::Base.connection.select_value("SELECT COUNT(*) FROM people")
@@ -385,11 +386,11 @@ def test_forked_child_recovers_from_disconnected_parent
 
               pid = fork {
                 rd.close
-                if ActiveRecord::Base.connection.active?
-                  pair = [ActiveRecord::Base.connection.object_id,
-                          ActiveRecord::Base.connection.select_value("SELECT COUNT(*) FROM people")]
-                  wr.write Marshal.dump pair
-                end
+                wr.write Marshal.dump [
+                  !!ActiveRecord::Base.connection.active?,
+                  ActiveRecord::Base.connection.object_id,
+                  ActiveRecord::Base.connection.select_value("SELECT COUNT(*) FROM people"),
+                ]
                 wr.close
 
                 exit # allow finalizers to run
@@ -401,8 +402,9 @@ def test_forked_child_recovers_from_disconnected_parent
             wr.close
 
             Process.waitpid outer_pid
-            child_id, child_count = Marshal.load(rd.read)
+            active, child_id, child_count = Marshal.load(rd.read)
 
+            assert_equal false, active
             assert_not_equal object_id, child_id
             rd.close
 

activerecord/test/cases/connection_pool_test.rb

@@ -83,6 +83,18 @@ def test_with_connection
         assert_equal 0, active_connections(pool).size
       end
 
+      def test_new_connection_no_query
+        skip("Can't test with in-memory dbs") if in_memory_db?
+        assert_equal 0, pool.connections.size
+        pool.with_connection { |_conn| } # warm the schema cache
+        pool.flush(0)
+        assert_equal 0, pool.connections.size
+
+        assert_no_queries do
+          pool.with_connection { |_conn| }
+        end
+      end
+
       def test_active_connection_in_use
         assert_not_predicate pool, :active_connection?
         main_thread = pool.connection

activerecord/test/cases/reaper_test.rb

@@ -180,6 +180,7 @@ def new_conn_in_thread(pool)
 
           child = Thread.new do
             conn = pool.checkout
+            conn.query("SELECT 1") # ensure connected
             event.set
             Thread.stop
           end

railties/test/application/initializers/frameworks_test.rb

@@ -219,7 +219,7 @@ def show
     test "can boot with an unhealthy database" do
       rails %w(generate model post title:string)
 
-      switch_env("DATABASE_URL", "mysql2://127.0.0.1:1") do
+      with_unhealthy_database do
         require "#{app_path}/config/environment"
       end
     end
@@ -261,16 +261,18 @@ def show
         config.eager_load = true
       RUBY
 
-      switch_env("DATABASE_URL", "mysql2://127.0.0.1:1") do
-        # The existing schema cache dump will contain ActiveRecord::ConnectionAdapters::SQLite3::Column objects
-        require "active_record/connection_adapters/sqlite3/column"
-
+      with_unhealthy_database do
         require "#{app_path}/config/environment"
 
-        assert_nil ActiveRecord::Base.connection_pool.schema_cache
+        assert_not_nil ActiveRecord::Base.connection_pool.schema_cache
+
         assert_raises ActiveRecord::ConnectionNotEstablished do
           ActiveRecord::Base.connection.execute("SELECT 1")
         end
+
+        assert_raises ActiveRecord::ConnectionNotEstablished do
+          ActiveRecord::Base.connection_pool.schema_cache.columns("posts")
+        end
       end
     end
 
@@ -296,10 +298,7 @@ def show
         config.active_record.check_schema_cache_dump_version = false
       RUBY
 
-      switch_env("DATABASE_URL", "mysql2://127.0.0.1:1") do
-        # The existing schema cache dump will contain ActiveRecord::ConnectionAdapters::SQLite3::Column objects
-        require "active_record/connection_adapters/sqlite3/column"
-
+      with_unhealthy_database do
         require "#{app_path}/config/environment"
 
         assert ActiveRecord::Base.connection_pool.schema_cache.data_sources("posts")

railties/test/isolation/abstract_unit.rb

@@ -227,6 +227,23 @@ def multi_db_database_configs
       YAML
     end
 
+    def with_unhealthy_database(&block)
+      # The existing schema cache dump will contain ActiveRecord::ConnectionAdapters::SQLite3Adapter objects
+      require "active_record/connection_adapters/sqlite3_adapter"
+
+      # We need to change the `database_version` to match what is expected for MySQL
+      dump_path = File.join(app_path, "db/schema_cache.yml")
+      if File.exist?(dump_path)
+        schema_cache = ActiveRecord::ConnectionAdapters::SchemaCache.load_from(dump_path)
+        schema_cache.connection = Struct.new(:schema_version).new(schema_cache.version)
+        schema_cache.instance_variable_set(:@database_version, ActiveRecord::ConnectionAdapters::AbstractAdapter::Version.new("8.8.8"))
+        File.write(dump_path, YAML.dump(schema_cache))
+      end
+
+      # We load the app while pointing at a non-existing MySQL server
+      switch_env("DATABASE_URL", "mysql2://127.0.0.1:1", &block)
+    end
+
     # Make a very basic app, without creating the whole directory structure.
     # This is faster and simpler than the method above.
     def make_basic_app