Active Record's railtie can now eager load tenanted models at boot

Author: flavorjones

GUIDE.md

@@ -54,6 +54,7 @@ Documentation outline:
   - app.config.hosts
   - example TenantSelector proc
 - explain Tenanted options "connection_class" and "tenanted_rails_records"
+- explain why we pin `active_record.use_schema_cache_dump = true` and `active_record.check_schema_cache_dump_version = false`
 
 TODO:
 
@@ -83,8 +84,8 @@ TODO:
   - [x] all the creation and schema migration complications (we have existing tests for this)
     - [x] read and write to the schema dump file
     - [x] write to the schema cache dump file
-    - [ ] make sure we read from the schema cache dump file when untenanted
-    - [ ] test production eager loading of the schema cache from dump files
+    - [x] make sure we read from the schema cache dump file when untenanted
+    - [x] test production eager loading of the schema cache from dump files
   - [ ] feature to turn off automatic creation/migration
     - pay attention to Rails.config.active_record.migration_error when we turn off auto-migrating
   - [ ] UntenantedConnectionPool should peek at its stack and if it happened during schema cache load, output a friendly message to let people know what to do

lib/active_record/tenanted/base.rb

@@ -36,6 +36,22 @@ def subtenant_of(class_name)
         def tenanted?
           false
         end
+
+        # TODO: This monkey patch shouldn't be necessary after 8.1 lands and the need for a
+        # connection is removed. For details see https://github.com/rails/rails/pull/54348
+        def _default_attributes # :nodoc:
+          @default_attributes ||= begin
+            # I've removed the `with_connection` block here.
+            nil_connection = nil
+            attributes_hash = columns_hash.transform_values do |column|
+              ActiveModel::Attribute.from_database(column.name, column.default, type_for_column(nil_connection, column))
+            end
+
+            attribute_set = ActiveModel::AttributeSet.new(attributes_hash)
+            apply_pending_attribute_modifications(attribute_set)
+            attribute_set
+          end
+        end
       end
     end
   end

lib/active_record/tenanted/railtie.rb

@@ -39,10 +39,23 @@ class Railtie < ::Rails::Railtie
         end
       end
 
+      initializer("active_record_tenanted.active_record_schema_cache",
+                  before: "active_record.copy_schema_cache_config") do
+        # Rails must be able to load the schema for a tenanted model without a database connection
+        # (e.g., boot-time eager loading, or calling User.new to build a form). This gem relies on
+        # reading from the schema cache dump to do that.
+        #
+        # Rails defaults use_schema_cache_dump to true, but we explicitly re-set it here because if
+        # this is ever turned off, Rails will not work as expected.
+        Rails.application.config.active_record.use_schema_cache_dump = true
+
+        # The schema cache version check needs to query the database, which isn't always possible
+        # for tenanted models.
+        Rails.application.config.active_record.check_schema_cache_dump_version = false
+      end
+
       initializer "active_record-tenanted.monkey_patches" do
         ActiveSupport.on_load(:active_record) do
-          # require "rails/generators/active_record/migration.rb"
-          # ActiveRecord::Generators::Migration.prepend(ActiveRecord::Tenanted::Patches::Migration)
           ActiveRecord::Tasks::DatabaseTasks.prepend(ActiveRecord::Tenanted::Patches::DatabaseTasks)
         end
 

lib/active_record/tenanted/untenanted_connection_pool.rb

@@ -19,9 +19,12 @@ def initialize(db_config, model)
         @model = model
       end
 
-      def schema_cache
+      def schema_reflection
         schema_cache_path = ActiveRecord::Tasks::DatabaseTasks.cache_dump_filename(db_config)
-        schema_reflection = ActiveRecord::ConnectionAdapters::SchemaReflection.new(schema_cache_path)
+        ActiveRecord::ConnectionAdapters::SchemaReflection.new(schema_cache_path)
+      end
+
+      def schema_cache
         ActiveRecord::ConnectionAdapters::BoundSchemaReflection.new(schema_reflection, self)
       end
 

test/unit/base_test.rb

@@ -78,6 +78,11 @@
             User.schema_cache.columns("users")
           end
         end
+
+        test "the active record railtie will fail to eager load the schema" do
+          # the code here should mirror the "active_record.define_attribute_methods" initializer
+          assert_not(User.connection_pool.schema_reflection.cached?(User.table_name))
+        end
       end
 
       describe "when schema cache dump file exists" do
@@ -96,6 +101,15 @@
           assert_equal([ "id", "email", "created_at", "updated_at" ],
                        User.schema_cache.columns("users")&.map(&:name))
         end
+
+        test "the active record railtie will eager load the schema" do
+          # the code here should mirror the "active_record.define_attribute_methods" initializer
+          assert(User.connection_pool.schema_reflection.cached?(User.table_name))
+          assert(User.define_attribute_methods)
+
+          assert_equal([ "id", "email", "created_at", "updated_at" ],
+                       User.schema_cache.columns("users")&.map(&:name))
+        end
       end
     end
   end