Test connection pools and schema migrations

Author: flavorjones

GUIDE.md

@@ -69,17 +69,21 @@ TODO:
     - [ ] should error if self is not an abstract base class or if target is not tenanted abstract base class
   - [x] `tenant_config_name` and `.tenanted?`
   - [ ] `.tenanted_class` nil or the abstract base class
-  - [ ] all the creation and schema migration complications (we have existing tests for this)
-    - think about race conditions here, maybe use a file lock to figure it out
-      - running migrations (they are done in a transaction, but the second thread's migration may fail resulting in a 500?)
-      - loading schemas (if the first thread loads the schema and inserts data, can the second thread accidentally drop/load causing data loss?)
+  - [x] shared connection pools
+  - [x] all the creation and schema migration complications (we have existing tests for this)
+accidentally drop/load causing data loss?)
   - [ ] feature to turn off automatic creation/migration
     - make sure we pay attention to Rails.config.active_record.migration_error when we turn off auto-migrating
   - thinking
     - should there be a global singleton `Tenant`? I'm not sure we need it or the limitations of a global.
     - if we do, though, then `.tenanted` should set `Tenant.base_class=`
     - and we need to add checks that `.tenanted` is called only ONCE in the application
 
+- [ ] think about race conditions with database creation and schema migrations
+  - maybe use a file lock to figure it out
+  - running migrations (they are done in a transaction, but the second thread's migration may fail resulting in a 500?)
+  - loading schemas (if the first thread loads the schema and inserts data, can the second thread
+
 - database tasks
   - [ ] RootConfig should conditionally re-enable database tasks ... when AR_TENANT is present?
   - [ ] make `db:migrate:tenants` iterate over all the tenants on disk

lib/active_record/tenanted/database_configurations.rb

@@ -9,6 +9,16 @@ class RootConfig < ActiveRecord::DatabaseConfigurations::HashConfig
         def database_tasks?
           false
         end
+
+        def database_path_for(tenant_name)
+          format_specifiers = {
+            tenant: tenant_name,
+          }
+          database % format_specifiers
+        end
+      end
+
+      class TenantConfig < RootConfig
       end
     end
   end

lib/active_record/tenanted/patches.rb

@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+module ActiveRecord
+  module Tenanted
+    module Patches
+      # TODO: I think this is needed because there was no followup to rails/rails#46270.
+      # See rails/rails@901828f2 from that PR for background.
+      module DatabaseTasks
+        private def with_temporary_pool(db_config, clobber: false)
+          original_db_config = begin
+            migration_class.connection_db_config
+          rescue ActiveRecord::ConnectionNotDefined
+            nil
+          end
+
+          begin
+            pool = migration_class.connection_handler.establish_connection(db_config, clobber: clobber)
+
+            yield pool
+          ensure
+            migration_class.connection_handler.establish_connection(original_db_config, clobber: clobber) if original_db_config
+          end
+        end
+      end
+    end
+  end
+end

lib/active_record/tenanted/railtie.rb

@@ -17,6 +17,14 @@ class Railtie < ::Rails::Railtie
           prepend ActiveRecord::Tenanted::Base
         end
       end
+
+      initializer "active_record-tenanted.monkey_patches.active_record" 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
+      end
     end
   end
 end

lib/active_record/tenanted/tenant.rb

@@ -30,6 +30,42 @@ def while_tenanted(tenant_name, &block)
 
         def connection_pool
           raise NoTenantError unless current_tenant
+
+          pool = connection_handler.retrieve_connection_pool(connection_specification_name, role: current_role, shard: current_tenant, strict: false)
+
+          if pool.nil?
+            create_tenanted_pool
+            pool = connection_handler.retrieve_connection_pool(connection_specification_name, role: current_role, shard: current_tenant, strict: true)
+          end
+
+          pool
+        end
+
+        def create_tenanted_pool # :nodoc:
+          # ensure all classes use the same connection pool
+          return superclass.create_tenanted_pool unless connection_class?
+
+          tenant = current_tenant
+          base_config = ActiveRecord::Base.configurations.resolve(tenanted_config_name.to_sym)
+          tenant_name = "#{tenanted_config_name}_#{tenant}"
+          config_hash = base_config.configuration_hash.dup.tap do |hash|
+            hash[:database] = base_config.database_path_for(tenant)
+          end
+          config = Tenanted::DatabaseConfigurations::TenantConfig.new(base_config.env_name, tenant_name, config_hash)
+
+          establish_connection(config)
+          ensure_schema_migrations(config)
+        end
+
+        def ensure_schema_migrations(config) # :nodoc:
+          ActiveRecord::Tasks::DatabaseTasks.with_temporary_connection(config) do |conn|
+            pool = conn.pool
+
+            if pool.migration_context.pending_migration_versions.present?
+              ActiveRecord::Tasks::DatabaseTasks.migrate(nil)
+              # ActiveRecord::Tasks::DatabaseTasks.dump_schema(config) if Rails.env.development?
+            end
+          end
         end
       end
     end

test/scenarios/vanilla/database.yml

@@ -6,8 +6,8 @@ test:
     tenanted: true
     adapter: "sqlite3"
     database: "%{storage}/test/primary/%%{tenant}/main.sqlite3"
-    migrations_paths: "%{__dir__}/primary_migrations"
+    migrations_paths: "%{scenario}/primary_migrations"
   shared:
     adapter: "sqlite3"
     database: "%{storage}/test/shared.sqlite3"
-    migrations_paths: "%{__dir__}/shared_migrations"
+    migrations_paths: "%{scenario}/shared_migrations"

test/scenarios/vanilla/tenanted_primary.rb

@@ -9,6 +9,9 @@ class TenantedApplicationRecord < ActiveRecord::Base
 class User < TenantedApplicationRecord
 end
 
+class Post < TenantedApplicationRecord
+end
+
 class SharedApplicationRecord < ActiveRecord::Base
   self.abstract_class = true
 

test/scenarios/vanilla/tenanted_secondary.rb

@@ -9,6 +9,9 @@ class TenantedApplicationRecord < ActiveRecord::Base
 class User < TenantedApplicationRecord
 end
 
+class Post < TenantedApplicationRecord
+end
+
 class SharedApplicationRecord < ActiveRecord::Base
   primary_abstract_class
 

test/scenarios/vanilla_named_primary/database.yml

@@ -6,8 +6,8 @@ test:
     tenanted: true
     adapter: "sqlite3"
     database: "%{storage}/test/tenanted/%%{tenant}/main.sqlite3"
-    migrations_paths: "%{__dir__}/tenanted_migrations"
+    migrations_paths: "%{scenario}/tenanted_migrations"
   shared:
     adapter: "sqlite3"
     database: "%{storage}/test/shared.sqlite3"
-    migrations_paths: "%{__dir__}/shared_migrations"
+    migrations_paths: "%{scenario}/shared_migrations"

test/scenarios/vanilla_named_primary/tenanted_primary.rb

@@ -9,6 +9,9 @@ class TenantedApplicationRecord < ActiveRecord::Base
 class User < TenantedApplicationRecord
 end
 
+class Post < TenantedApplicationRecord
+end
+
 class SharedApplicationRecord < ActiveRecord::Base
   self.abstract_class = true