Merge pull request rails#54333 from jcoleman/eliminate-queries-when-loading-dumped-model-schema

rafaelfranca · web-flow · commit e82b645262b2 · 2025-01-23T15:18:12.000-08:00
Eliminate queries loading dumped model schema on Postgres
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/testing/query_assertions.rb b/activerecord/lib/active_record/testing/query_assertions.rb
@@ -11,12 +11,18 @@ module QueryAssertions
       #   # Check for any number of queries
       #   assert_queries_count { Post.first }
       #
-      # If the +:include_schema+ option is provided, any queries (including schema related) are counted.
+      # Any unmaterialized transactions will be materialized to ensure only
+      # queries attempted intside the block are counted.
+      #
+      # If the +:include_schema+ option is provided, any queries (including
+      # schema related) are counted. Setting this option also skips leasing a
+      # connection to materialize pending transactions since we want to count
+      # queries executed at connection open (e.g., type map).
       #
       #   assert_queries_count(1, include_schema: true) { Post.columns }
       #
       def assert_queries_count(count = nil, include_schema: false, &block)
-        ActiveRecord::Base.lease_connection.materialize_transactions
+        ActiveRecord::Base.lease_connection.materialize_transactions unless include_schema
 
         counter = SQLCounter.new
         ActiveSupport::Notifications.subscribed(counter, "sql.active_record") do
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
diff --git a/activerecord/test/cases/migration/change_schema_test.rb b/activerecord/test/cases/migration/change_schema_test.rb
diff --git a/activerecord/test/cases/migration/columns_test.rb b/activerecord/test/cases/migration/columns_test.rb
diff --git a/activerecord/test/cases/schema_loading_test.rb b/activerecord/test/cases/schema_loading_test.rb
diff --git a/activerecord/test/support/fake_adapter.rb b/activerecord/test/support/fake_adapter.rb
