Flexible switching with rails 5.2

lib/apartment.rb

@@ -10,12 +10,12 @@ class << self
 
     ACCESSOR_METHODS = [
       :use_sql, :seed_after_create, :tenant_decorator,
-      :force_reconnect_on_switch, :pool_per_config
+      :force_reconnect_on_switch, :pool_per_config, :default_tenant
     ]
     WRITER_METHODS   = [
       :tenant_names, :database_schema_file, :excluded_models,
       :persistent_schemas, :connection_class, :tld_length, :db_migrate_tenants,
-      :seed_data_file, :default_tenant
+      :seed_data_file
     ]
     OTHER_METHODS    = [:tenant_resolver, :resolver_class]
 
@@ -59,10 +59,6 @@ def excluded_models
       @excluded_models || []
     end
 
-    def default_tenant
-      @default_tenant || tenant_resolver.init_config
-    end
-
     def persistent_schemas
       @persistent_schemas || []
     end
@@ -90,6 +86,17 @@ def reset
 
       Thread.current[:_apartment_connection_specification_name] = nil
     end
+
+    def clear_connections
+      connection_class.clear_all_connections!
+      connection_handler.tap do |ch|
+        ch.send(:owner_to_pool).each_key do |k|
+          ch.remove_connection(k) if k =~ /^_apartment/
+        end
+      end
+      Thread.current[:_apartment_connection_specification_name] = nil
+      Apartment::Tenant.reload!
+    end
   end
 
   # Exceptions

lib/apartment/adapters/abstract_adapter.rb

@@ -30,8 +30,10 @@ def create(tenant)
             config = config_for(tenant)
             difference = current_difference_from(config)
 
-            if difference[:host]
-              connection_switch(config, without_keys: [:database, :schema_search_path])
+            if difference[:host] || difference[:port]
+              temp_config = config.dup
+              temp_config[:database] = "postgres"
+              connection_switch!(temp_config)
             end
 
             create_tenant!(config)
@@ -42,6 +44,8 @@ def create(tenant)
             seed_data if Apartment.seed_after_create
 
             yield if block_given?
+          rescue *rescuable_exceptions => exception
+            raise_create_tenant_error!(tenant, exception)
           ensure
             switch!(previous_tenant) rescue reset
           end
@@ -54,8 +58,10 @@ def drop(tenant)
         config = config_for(tenant)
         difference = current_difference_from(config)
 
-        if difference[:host]
-          connection_switch(config, without_keys: [:database])
+        if difference[:host] || difference[:port]
+          temp_config = config.dup
+          temp_config[:database] = "postgres"
+          connection_switch!(temp_config)
         end
 
         unless database_exists?(config[:database])
@@ -65,6 +71,8 @@ def drop(tenant)
         Apartment.connection.drop_database(config[:database])
 
         @current = tenant
+      rescue *rescuable_exceptions => exception
+        raise_drop_tenant_error!(tenant, exception)
       ensure
         switch!(previous_tenant) rescue reset
       end
@@ -116,12 +124,10 @@ def setup_connection_specification_name
         Apartment.connection_class.connection_specification_name = nil
         Apartment.connection_class.instance_eval do
           def connection_specification_name
-            if !defined?(@connection_specification_name) || @connection_specification_name.nil?
-              apartment_spec_name = Thread.current[:_apartment_connection_specification_name]
-              return apartment_spec_name ||
-                (self == ActiveRecord::Base ? "primary" : superclass.connection_specification_name)
-            end
-            @connection_specification_name
+            return :_apartment_excluded if @connection_specification_name == :_apartment_excluded
+
+            return Thread.current[:_apartment_connection_specification_name] ||
+              (self == ActiveRecord::Base ? "primary" : superclass.connection_specification_name)
           end
         end
       end
@@ -155,6 +161,7 @@ def import_database_schema
       def seed_data
         silence_warnings{ load_or_abort(Apartment.seed_data_file) } if Apartment.seed_data_file
       end
+      alias seed seed_data
 
       def load_or_abort(file)
         if File.exist?(file)
@@ -164,9 +171,25 @@ def load_or_abort(file)
         end
       end
 
+      def rescuable_exceptions
+        [ActiveRecord::ActiveRecordError] + Array(rescue_from)
+      end
+
+      def rescue_from
+        []
+      end
+
       def raise_connect_error!(tenant, exception)
         raise TenantNotFound, "Error while connecting to tenant #{tenant}: #{exception.message}"
       end
+
+      def raise_create_tenant_error!(tenant, exception)
+        raise TenantExists, "Error while creating tenant #{tenant}: #{ exception.message }"
+      end
+
+      def raise_drop_tenant_error!(tenant, exception)
+        raise TenantNotFound, "Error while dropping tenant #{tenant}: #{ exception.message }"
+      end
     end
   end
 end

lib/apartment/adapters/postgresql_adapter.rb

@@ -4,12 +4,9 @@
 module Apartment
   module Adapters
     class PostgresqlAdapter < AbstractAdapter
-      # -- ABSTRACT OVERRIDES --
       def drop(tenant)
-        raise NotImplementedError,
-          "Please use either drop_database or drop_schema for PG adapter"
+        drop_database(tenant)
       end
-      # -- END ABSTRACT OVERRIDES --
 
       def drop_database(tenant)
         # Apartment.connection.select_all "select pg_terminate_backend(pg_stat_activity.pid) from pg_stat_activity where datname='#{tenant}' AND state='idle';"
@@ -42,7 +39,7 @@ def switch_tenant(config)
         difference = config.select{ |k, v| current_config[k] != v }
 
         # PG doesn't have the ability to switch DB without reconnecting
-        if difference[:host] || difference[:database]
+        if difference[:host] || difference[:database] || difference[:port]
           connection_switch!(config)
         else
           simple_switch(config) if difference[:schema_search_path]
@@ -78,11 +75,16 @@ def connection_specification_name(config)
         if Apartment.pool_per_config
           "_apartment_#{config.hash}".to_sym
         else
-          host_hash = Digest::MD5.hexdigest(config[:host] || config[:url] || "127.0.0.1")
+          value = "#{config[:host]}:#{config[:port]}" || config[:url] || "127.0.0.1"
+          host_hash = Digest::MD5.hexdigest(value)
           "_apartment_#{host_hash}_#{config[:adapter]}_#{config[:database]}".to_sym
         end
       end
 
+      def rescue_from
+        PG::Error
+      end
+
       private
         def database_exists?(database)
           result = Apartment.connection.exec_query(<<-SQL).try(:first)

lib/apartment/migrator.rb

@@ -10,24 +10,42 @@ def migrate(database)
       Tenant.switch(database) do
         version = ENV["VERSION"] ? ENV["VERSION"].to_i : nil
 
-        ActiveRecord::Migrator.migrate(ActiveRecord::Migrator.migrations_paths, version) do |migration|
-          ENV["SCOPE"].blank? || (ENV["SCOPE"] == migration.scope)
+        migration_scope_block = -> (migration) { ENV["SCOPE"].blank? || (ENV["SCOPE"] == migration.scope) }
+
+        if activerecord_below_5_2?
+          ActiveRecord::Migrator.migrate(ActiveRecord::Migrator.migrations_paths, version, &migration_scope_block)
+        else
+          ActiveRecord::Base.connection.migration_context.migrate(version, &migration_scope_block)
         end
       end
     end
 
     # Migrate up/down to a specific version
     def run(direction, database, version)
       Tenant.switch(database) do
-        ActiveRecord::Migrator.run(direction, ActiveRecord::Migrator.migrations_paths, version)
+        if activerecord_below_5_2?
+          ActiveRecord::Migrator.run(direction, ActiveRecord::Migrator.migrations_paths, version)
+        else
+          ActiveRecord::Base.connection.migration_context.run(direction, version)
+        end
       end
     end
 
     # rollback latest migration `step` number of times
     def rollback(database, step = 1)
       Tenant.switch(database) do
-        ActiveRecord::Migrator.rollback(ActiveRecord::Migrator.migrations_paths, step)
+        if activerecord_below_5_2?
+          ActiveRecord::Migrator.rollback(ActiveRecord::Migrator.migrations_paths, step)
+        else
+          ActiveRecord::Base.connection.migration_context.rollback(step)
+        end
       end
     end
+
+    private
+
+    def activerecord_below_5_2?
+      ActiveRecord.version.release() < Gem::Version.new('5.2.0')
+    end
   end
 end

lib/apartment/railtie.rb

@@ -28,9 +28,14 @@ def self.prep
     #   See the middleware/console declarations below to help with this. Hope to fix that soon.
     #
     config.to_prepare do
-      unless ARGV.any? { |arg| arg =~ /\Aassets:(?:precompile|clean)\z/ }
+      next if ARGV.any? { |arg| arg =~ /\Aassets:(?:precompile|clean)\z/ }
+
+      begin
         Apartment::Tenant.init
         Apartment.connection_class.clear_active_connections!
+      rescue ::ActiveRecord::NoDatabaseError
+        # Since `db:create` and other tasks invoke this block from Rails 5.2.0,
+        # we need to swallow the error to execute `db:create` properly.
       end
     end
 

lib/apartment/resolvers/config.rb

@@ -0,0 +1,23 @@
+require 'apartment/resolvers/abstract'
+
+module Apartment
+  module Resolvers
+    class Config < Abstract
+      def resolve(tenant)
+        return init_config.dup if !tenant || tenant == Apartment.default_tenant
+
+        database_config(tenant)
+      end
+
+      private
+
+      def database_config(tenant)
+        ActiveRecord::Base.configurations[config_name(tenant)].symbolize_keys
+      end
+
+      def config_name(tenant)
+        "#{Rails.env}_#{tenant}"
+      end
+    end
+  end
+end

lib/apartment/tasks/enhancements.rb

@@ -3,33 +3,54 @@
 
 module Apartment
   class RakeTaskEnhancer
-
-    TASKS = %w(db:migrate db:rollback db:migrate:up db:migrate:down db:migrate:redo db:seed)
-
+
+    module TASKS
+      ENHANCE_BEFORE = %w(db:drop)
+      ENHANCE_AFTER  = %w(db:migrate db:rollback db:migrate:up db:migrate:down db:migrate:redo db:seed)
+      freeze
+    end
+
     # This is a bit convoluted, but helps solve problems when using Apartment within an engine
     # See spec/integration/use_within_an_engine.rb
-    
+
     class << self
       def enhance!
-        TASKS.each do |name|
+        return unless should_enhance?
+
+        # insert task before
+        TASKS::ENHANCE_BEFORE.each do |name|
           task = Rake::Task[name]
-          task.enhance do
-            if should_enhance?
-              enhance_task(task)
-            end
-          end
+          enhance_before_task(task)
         end
+
+        # insert task after
+        TASKS::ENHANCE_AFTER.each do |name|
+          task = Rake::Task[name]
+          enhance_after_task(task)
+        end
+
       end
-    
+
       def should_enhance?
         Apartment.db_migrate_tenants
       end
-
-      def enhance_task(task)
-        Rake::Task[task.name.sub(/db:/, 'apartment:')].invoke
+
+      def enhance_before_task(task)
+        task.enhance([inserted_task_name(task)])
+      end
+
+      def enhance_after_task(task)
+        task.enhance do
+          Rake::Task[inserted_task_name(task)].invoke
+        end
       end
+
+      def inserted_task_name(task)
+        task.name.sub(/db:/, 'apartment:')
+      end
+
     end
-    
+
   end
 end
 

lib/apartment/tenant.rb

@@ -25,9 +25,7 @@ def init
     #
     def adapter
       Thread.current[:apartment_adapter] ||= begin
-        config = Apartment.default_tenant
-
-        adapter_name = "#{config[:adapter]}_adapter"
+        adapter_name = "postgresql_adapter"
 
         begin
           require "apartment/adapters/#{adapter_name}"

lib/tasks/apartment.rake

@@ -3,14 +3,31 @@ require 'apartment/migrator'
 apartment_namespace = namespace :apartment do
 
   desc "Create all tenants"
-  task create: 'db:migrate' do
+  task :create do
+    Apartment::Tenant.init
     tenants.each do |tenant|
       begin
-        quietly { Apartment::Tenant.create(tenant) }
+        puts("Creating #{tenant} tenant")
+        Apartment::Tenant.create(tenant)
       rescue Apartment::TenantExists => e
         puts e.message
       end
     end
+    Apartment.clear_connections
+  end
+
+  desc "Drop all tenants"
+  task :drop do
+    tenants.each do |tenant|
+      begin
+        puts("Dropping #{tenant} tenant")
+        Apartment::Tenant.drop(tenant)
+      rescue Apartment::TenantNotFound => e
+        puts e.message
+      end
+    end
+
+    Apartment.clear_connections
   end
 
   desc "Migrate all tenants"
@@ -39,6 +56,7 @@ apartment_namespace = namespace :apartment do
         puts e.message
       end
     end
+    Apartment.clear_connections
   end
 
   desc "Rolls the migration back to the previous version (specify steps w/ STEP=n) across all tenants."