Eliminate queries loading dumped model schema on Postgres

jcoleman · jcoleman · commit 9ad36e067222 · 2025-01-23T21:47:23.000Z
Commit 835eb8a (rails#48743) adds a comment making it clear the intention is to avoid connections being opened/used at app boot: For resiliency, it is critical that a Rails application should be able to boot without depending on the database (or any other service) being responsive. Additionally the followup commit 35a614c (rails#48793) is even more explicit with the title: "Never connect to the database in define_attribute_methods initializer". And in commit 618db13 (rails#49452) the commit message noted: Note, this really isn't good, ideally types should be stored in the schema cache so that we don't have to extract types every time. Unfortunately neither of these changes added any test coverage for those stated invariants (they simply modified existing tests to pass). I assume that the code likely had internal testing done where it was developed in house (I think using MySQL) and was validated in production, but the lack of test coverage in Rails itself led to missing the fact that the invariant is _not_ respected by the PostgreSQL adapater because looking up the cast types (used when defining default attributes) requires a connection to be open to retrieve the type map (which in the PostgreSQL adapter is OID based rather than type name based). Opening a connection not only requires the database to be available at boot, it also triggers multiple queries to load the type map as well as configuring things like timezones which, beyond causing resiliency concerns, further slows down app boot. In this commit we store the cast type value in the schema dump so that it can be loaded along with the dumped schema file at boot rather than, even when using cached schemas, having to connect to the database to finish defining attribute methods.
diff --git a/activerecord/CHANGELOG.md b/activerecord/CHANGELOG.md
@@ -1,3 +1,11 @@
+*   Eliminate queries loading dumped schema cache on Postgres
+
+    Improve resiliency by avoiding needing to open a database connection to load the
+    type map while defining attribute methods at boot when a schema cache file is
+    configured on PostgreSQL databases.
+
+    *James Coleman*
+
 *   `ActiveRecord::Coder::JSON` can be instantiated
 
     Options can now be passed to `ActiveRecord::Coder::JSON` when instantiating the coder. This allows:
diff --git a/activerecord/lib/active_record/attributes.rb b/activerecord/lib/active_record/attributes.rb
@@ -241,10 +241,8 @@ def define_attribute(
 
       def _default_attributes # :nodoc:
         @default_attributes ||= begin
-          attributes_hash = with_connection do |connection|
-            columns_hash.transform_values do |column|
-              ActiveModel::Attribute.from_database(column.name, column.default, type_for_column(connection, column))
-            end
+          attributes_hash = columns_hash.transform_values do |column|
+            ActiveModel::Attribute.from_database(column.name, column.default, type_for_column(column))
           end
 
           attribute_set = ActiveModel::AttributeSet.new(attributes_hash)
@@ -299,7 +297,7 @@ def resolve_type_name(name, **options)
           Type.lookup(name, **options, adapter: Type.adapter_name_from(self))
         end
 
-        def type_for_column(connection, column)
+        def type_for_column(column)
           hook_attribute_type(column.name, super)
         end
     end
diff --git a/activerecord/lib/active_record/connection_adapters/abstract/database_statements.rb b/activerecord/lib/active_record/connection_adapters/abstract/database_statements.rb
@@ -624,8 +624,7 @@ def build_fixture_sql(fixtures, table_name)
 
             columns.map do |name, column|
               if fixture.key?(name)
-                type = lookup_cast_type_from_column(column)
-                with_yaml_fallback(type.serialize(fixture[name]))
+                with_yaml_fallback(column.cast_type.serialize(fixture[name]))
               else
                 default_insert_value(column)
               end
diff --git a/activerecord/lib/active_record/connection_adapters/abstract/quoting.rb b/activerecord/lib/active_record/connection_adapters/abstract/quoting.rb
@@ -114,19 +114,6 @@ def cast_bound_value(value) # :nodoc:
         value
       end
 
-      # If you are having to call this function, you are likely doing something
-      # wrong. The column does not have sufficient type information if the user
-      # provided a custom type on the class level either explicitly (via
-      # Attributes::ClassMethods#attribute) or implicitly (via
-      # AttributeMethods::Serialization::ClassMethods#serialize, +time_zone_aware_attributes+).
-      # In almost all cases, the sql type should only be used to change quoting behavior, when the primitive to
-      # represent the type doesn't sufficiently reflect the differences
-      # (varchar vs binary) for example. The type used to get this primitive
-      # should have been provided before reaching the connection adapter.
-      def lookup_cast_type_from_column(column) # :nodoc:
-        lookup_cast_type(column.sql_type)
-      end
-
       # Quotes a string, escaping any ' (single quote) and \ (backslash)
       # characters.
       def quote_string(s)
@@ -159,7 +146,8 @@ def quote_default_expression(value, column) # :nodoc:
         if value.is_a?(Proc)
           value.call
         else
-          value = lookup_cast_type(column.sql_type).serialize(value)
+          cast_type = column.respond_to?(:cast_type) ? column.cast_type : lookup_cast_type(column.sql_type)
+          value = cast_type.serialize(value)
           quote(value)
         end
       end
@@ -232,6 +220,8 @@ def type_casted_binds(binds)
           end
         end
 
+        # In some adapters this executes queries; it should not be used in any
+        # hot paths. Instead prefer using `column.cast_type`.
         def lookup_cast_type(sql_type)
           type_map.lookup(sql_type)
         end
diff --git a/activerecord/lib/active_record/connection_adapters/abstract/schema_dumper.rb b/activerecord/lib/active_record/connection_adapters/abstract/schema_dumper.rb
@@ -85,7 +85,7 @@ def schema_scale(column)
 
         def schema_default(column)
           return unless column.has_default?
-          type = @connection.lookup_cast_type_from_column(column)
+          type = column.cast_type
           default = type.deserialize(column.default)
           if default.nil?
             schema_expression(column)
diff --git a/activerecord/lib/active_record/connection_adapters/column.rb b/activerecord/lib/active_record/connection_adapters/column.rb
@@ -7,7 +7,7 @@ module ConnectionAdapters
     class Column
       include Deduplicable
 
-      attr_reader :name, :default, :sql_type_metadata, :null, :default_function, :collation, :comment
+      attr_reader :name, :cast_type, :default, :sql_type_metadata, :null, :default_function, :collation, :comment
 
       delegate :precision, :scale, :limit, :type, :sql_type, to: :sql_type_metadata, allow_nil: true
 
@@ -17,8 +17,9 @@ class Column
       # +default+ is the type-casted default value, such as +new+ in <tt>sales_stage varchar(20) default 'new'</tt>.
       # +sql_type_metadata+ is various information about the type of the column
       # +null+ determines if this column allows +NULL+ values.
-      def initialize(name, default, sql_type_metadata = nil, null = true, default_function = nil, collation: nil, comment: nil, **)
+      def initialize(name, cast_type, default, sql_type_metadata = nil, null = true, default_function = nil, collation: nil, comment: nil, **)
         @name = name.freeze
+        @cast_type = cast_type
         @sql_type_metadata = sql_type_metadata
         @null = null
         @default = default
@@ -45,6 +46,7 @@ def human_name
 
       def init_with(coder)
         @name = coder["name"]
+        @cast_type = coder["cast_type"]
         @sql_type_metadata = coder["sql_type_metadata"]
         @null = coder["null"]
         @default = coder["default"]
@@ -55,6 +57,7 @@ def init_with(coder)
 
       def encode_with(coder)
         coder["name"] = @name
+        coder["cast_type"] = @cast_type
         coder["sql_type_metadata"] = @sql_type_metadata
         coder["null"] = @null
         coder["default"] = @default
@@ -75,6 +78,7 @@ def auto_populated?
       def ==(other)
         other.is_a?(Column) &&
           name == other.name &&
+          cast_type == other.cast_type &&
           default == other.default &&
           sql_type_metadata == other.sql_type_metadata &&
           null == other.null &&
@@ -88,6 +92,7 @@ def hash
         Column.hash ^
           name.hash ^
           name.encoding.hash ^
+          cast_type.hash ^
           default.hash ^
           sql_type_metadata.hash ^
           null.hash ^
@@ -114,7 +119,7 @@ def deduplicated
 
     class NullColumn < Column
       def initialize(name, **)
-        super(name, nil)
+        super(name, nil, nil)
       end
     end
   end
diff --git a/activerecord/lib/active_record/connection_adapters/mysql/schema_statements.rb b/activerecord/lib/active_record/connection_adapters/mysql/schema_statements.rb
@@ -209,6 +209,7 @@ def new_column_from_field(table_name, field, _definitions)
 
             MySQL::Column.new(
               field["Field"],
+              lookup_cast_type(type_metadata.sql_type),
               default,
               type_metadata,
               field["Null"] == "YES",
diff --git a/activerecord/lib/active_record/connection_adapters/postgresql/quoting.rb b/activerecord/lib/active_record/connection_adapters/postgresql/quoting.rb
@@ -153,14 +153,14 @@ def quoted_binary(value) # :nodoc:
           "'#{escape_bytea(value.to_s)}'"
         end
 
+        # `column` may be either an instance of Column or ColumnDefinition.
         def quote_default_expression(value, column) # :nodoc:
           if value.is_a?(Proc)
             value.call
           elsif column.type == :uuid && value.is_a?(String) && value.include?("()")
             value # Does not quote function default values for UUID columns
           elsif column.respond_to?(:array?)
-            type = lookup_cast_type_from_column(column)
-            quote(type.serialize(value))
+            quote(column.cast_type.serialize(value))
           else
             super
           end
@@ -186,11 +186,6 @@ def type_cast(value) # :nodoc:
           end
         end
 
-        def lookup_cast_type_from_column(column) # :nodoc:
-          verify! if type_map.nil?
-          type_map.lookup(column.oid, column.fmod, column.sql_type)
-        end
-
         private
           def lookup_cast_type(sql_type)
             super(query_value("SELECT #{quote(sql_type)}::regtype::oid", "SCHEMA").to_i)
diff --git a/activerecord/lib/active_record/connection_adapters/postgresql/schema_statements.rb b/activerecord/lib/active_record/connection_adapters/postgresql/schema_statements.rb
@@ -1001,6 +1001,7 @@ def new_column_from_field(table_name, field, _definitions)
 
             PostgreSQL::Column.new(
               column_name,
+              get_oid_type(oid.to_i, fmod.to_i, column_name, type),
               default_value,
               type_metadata,
               !notnull,
diff --git a/activerecord/lib/active_record/connection_adapters/sqlite3/schema_statements.rb b/activerecord/lib/active_record/connection_adapters/sqlite3/schema_statements.rb
@@ -157,6 +157,7 @@ def new_column_from_field(table_name, field, definitions)
 
             Column.new(
               field["name"],
+              lookup_cast_type(field["type"]),
               default_value,
               type_metadata,
               field["notnull"].to_i == 0,
diff --git a/activerecord/lib/active_record/connection_adapters/sqlite3_adapter.rb b/activerecord/lib/active_record/connection_adapters/sqlite3_adapter.rb
@@ -654,8 +654,7 @@ def copy_table(from, to, options = {})
                 column_options[:stored] = column.virtual_stored?
                 column_options[:type] = column.type
               elsif column.has_default?
-                type = lookup_cast_type_from_column(column)
-                default = type.deserialize(column.default)
+                default = column.cast_type.deserialize(column.default)
                 default = -> { column.default_function } if default.nil?
 
                 unless column.auto_increment?
diff --git a/activerecord/lib/active_record/model_schema.rb b/activerecord/lib/active_record/model_schema.rb
@@ -619,8 +619,8 @@ def compute_table_name
           end
         end
 
-        def type_for_column(connection, column)
-          type = connection.lookup_cast_type_from_column(column)
+        def type_for_column(column)
+          type = column.cast_type
 
           if immutable_strings_by_default && type.respond_to?(:to_immutable_string)
             type = type.to_immutable_string
diff --git a/activerecord/lib/active_record/type_caster/connection.rb b/activerecord/lib/active_record/type_caster/connection.rb
@@ -19,7 +19,7 @@ def type_for_attribute(attr_name)
         if schema_cache.data_source_exists?(table_name)
           column = schema_cache.columns_hash(table_name)[attr_name.to_s]
           if column
-            type = @klass.with_connection { |connection| connection.lookup_cast_type_from_column(column) }
+            type = column.cast_type
           end
         end
 
diff --git a/activerecord/test/cases/adapter_test.rb b/activerecord/test/cases/adapter_test.rb
@@ -20,7 +20,7 @@ def test_type_map_is_ractor_shareable
       # the type map is Ractor shareable.
       @connection.tables.each do |table|
         @connection.columns(table).each do |column|
-          assert_ractor_shareable @connection.lookup_cast_type_from_column(column)
+          assert_ractor_shareable @connection.send(:lookup_cast_type, column.sql_type)
         end
       end
     end
diff --git a/activerecord/test/cases/attribute_methods_test.rb b/activerecord/test/cases/attribute_methods_test.rb
@@ -1135,6 +1135,33 @@ def name
     assert_equal "New topic", topic.title_alias_to_be_undefined
   end
 
+  test "#define_attribute_methods doesn't connect to the database when schema cache is present" do
+    with_temporary_connection_pool do
+      if in_memory_db?
+        # Separate connections to an in-memory database create an entirely new database,
+        # with an empty schema etc, so we just stub out this schema on the fly.
+        ActiveRecord::Base.connection_pool.with_connection do |connection|
+          connection.create_table :tasks do |t|
+            t.datetime :starting
+            t.datetime :ending
+          end
+        end
+      end
+
+      @target.table_name = "tasks"
+
+      @target.connection_pool.schema_cache.load!
+      @target.connection_pool.schema_cache.add("tasks")
+      @target.connection_pool.disconnect!
+
+      assert_no_queries(include_schema: true) do
+        @target.define_attribute_methods
+      end
+    ensure
+      ActiveRecord::Base.connection_pool.disconnect!
+    end
+  end
+
   test "define_attribute_method works with both symbol and string" do
     klass = Class.new(ActiveRecord::Base)
     klass.table_name = "foo"
@@ -1604,4 +1631,11 @@ def #{method_signature}
         end
       private_method
     end
+
+    def with_temporary_connection_pool(&block)
+      pool_config = ActiveRecord::Base.lease_connection.pool.pool_config
+      new_pool = ActiveRecord::ConnectionAdapters::ConnectionPool.new(pool_config)
+
+      pool_config.stub(:pool, new_pool, &block)
+    end
 end
diff --git a/activerecord/test/cases/connection_adapters/schema_cache_test.rb b/activerecord/test/cases/connection_adapters/schema_cache_test.rb
@@ -99,8 +99,9 @@ def test_yaml_dump_and_load
         # Load the cache.
         cache = load_bound_reflection(tempfile.path)
 
-        assert_no_queries do
+        assert_no_queries(include_schema: true) do
           assert_equal 3, cache.columns("courses").size
+          assert_equal 3, cache.columns("courses").map(&:cast_type).compact.size
           assert_equal 3, cache.columns_hash("courses").size
           assert cache.data_source_exists?("courses")
           assert_equal "id", cache.primary_keys("courses")
@@ -137,8 +138,9 @@ def test_yaml_dump_and_load_with_gzip
           YAML.respond_to?(:unsafe_load) ? YAML.unsafe_load(gz.read) : YAML.load(gz.read)
         end
 
-        assert_no_queries do
+        assert_no_queries(include_schema: true) do
           assert_equal 3, cache.columns(@connection, "courses").size
+          assert_equal 3, cache.columns(@connection, "courses").map(&:cast_type).compact.size
           assert_equal 3, cache.columns_hash(@connection, "courses").size
           assert cache.data_source_exists?(@connection, "courses")
           assert_equal "id", cache.primary_keys(@connection, "courses")
diff --git a/activerecord/test/cases/migration/change_schema_test.rb b/activerecord/test/cases/migration/change_schema_test.rb
@@ -70,8 +70,8 @@ def test_create_table_with_defaults
         five = columns.detect { |c| c.name == "five" } unless mysql
 
         assert_equal "hello", one.default
-        assert_equal true, connection.lookup_cast_type_from_column(two).deserialize(two.default)
-        assert_equal false, connection.lookup_cast_type_from_column(three).deserialize(three.default)
+        assert_equal true, two.cast_type.deserialize(two.default)
+        assert_equal false, three.cast_type.deserialize(three.default)
         assert_equal "1", four.default
         assert_equal "hello", five.default unless mysql
       end
diff --git a/activerecord/test/cases/migration/columns_test.rb b/activerecord/test/cases/migration/columns_test.rb
@@ -195,19 +195,19 @@ def test_change_column
 
         old_columns = connection.columns(TestModel.table_name)
         assert old_columns.find { |c|
-          default = connection.lookup_cast_type_from_column(c).deserialize(c.default)
+          default = c.cast_type.deserialize(c.default)
           c.name == "approved" && c.type == :boolean && default == true
         }
 
         change_column :test_models, :approved, :boolean, default: false
         new_columns = connection.columns(TestModel.table_name)
 
         assert_not new_columns.find { |c|
-          default = connection.lookup_cast_type_from_column(c).deserialize(c.default)
+          default = c.cast_type.deserialize(c.default)
           c.name == "approved" && c.type == :boolean && default == true
         }
         assert new_columns.find { |c|
-          default = connection.lookup_cast_type_from_column(c).deserialize(c.default)
+          default = c.cast_type.deserialize(c.default)
           c.name == "approved" && c.type == :boolean && default == false
         }
         change_column :test_models, :approved, :boolean, default: true
diff --git a/activerecord/test/cases/schema_loading_test.rb b/activerecord/test/cases/schema_loading_test.rb
@@ -42,6 +42,38 @@ def test_model_with_changed_custom_lock_is_loaded_twice
     assert_equal 2, klass.load_schema_calls
   end
 
+  def test_schema_loading_doesnt_query_when_schema_cache_is_loaded
+    with_temporary_connection_pool do
+      if in_memory_db?
+        # Separate connections to an in-memory database create an entirely new database,
+        # with an empty schema etc, so we just stub out this schema on the fly.
+        ActiveRecord::Base.connection_pool.with_connection do |connection|
+          connection.create_table :tasks do |t|
+            t.datetime :starting
+            t.datetime :ending
+          end
+        end
+      end
+
+      klass = define_model do |c|
+        c.table_name = :tasks
+      end
+
+      klass.connection_pool.schema_cache.load!
+      klass.connection_pool.schema_cache.add("tasks")
+      klass.connection_pool.disconnect!
+      klass.send(:reload_schema_from_cache)
+
+
+      assert_no_queries(include_schema: true) do
+        klass.load_schema
+      end
+      assert_equal 1, klass.load_schema_calls
+    ensure
+      ActiveRecord::Base.connection_pool.disconnect!
+    end
+  end
+
   private
     def define_model
       Class.new(ActiveRecord::Base) do
@@ -50,4 +82,11 @@ def define_model
         yield self if block_given?
       end
     end
+
+    def with_temporary_connection_pool(&block)
+      pool_config = ActiveRecord::Base.lease_connection.pool.pool_config
+      new_pool = ActiveRecord::ConnectionAdapters::ConnectionPool.new(pool_config)
+
+      pool_config.stub(:pool, new_pool, &block)
+    end
 end
diff --git a/activerecord/test/support/fake_adapter.rb b/activerecord/test/support/fake_adapter.rb
