lib/penthouse/migrator.rb in penthouse-0.8.0 vs lib/penthouse/migrator.rb in penthouse-0.9.0

@@ -1,9 +1,7 @@
 require 'active_record'
-require 'set'
 require 'active_support/core_ext/module/aliasing'
-require 'active_support/core_ext/array/wrap'
 
 module Penthouse
   module Migration
     def self.included(base)
       base.alias_method_chain :announce, :penthouse
@@ -35,54 +33,79 @@
           alias_method :migrate_with_octopus, :migrate_with_penthouse
           alias_method :up_with_octopus,      :up_with_penthouse
           alias_method :down_with_octopus,    :down_with_penthouse
           alias_method :run_with_octopus,     :run_with_penthouse
         end
+
+        alias_method_chain :migrate,    :penthouse
+        alias_method_chain :migrations, :penthouse
+
+        # override any new Octopus methods with the new Penthouse ones
+        alias_method :migrate_with_octopus,    :migrate_with_penthouse
+        alias_method :migrations_with_octopus, :migrations_with_penthouse
       end
     end
 
+    # this may seem stupid but it gets around issues with Octopus
+    def migrate_with_penthouse(*args)
+      migrate_without_penthouse(*args)
+    end
+
+    # this may seem stupid but it gets around issues with Octopus
+    def migrations_with_penthouse
+      migrations_without_penthouse
+    end
+
     module ClassMethods
       def migrate_with_penthouse(migrations_paths, target_version = nil, &block)
         unless Penthouse.configuration.migrate_tenants?
           return migrate_without_penthouse(migrations_paths, target_version, &block)
         end
 
-        Penthouse.each_tenant(tenant_identifiers: tenants_to_migrate) do |tenant|
+        wrap_penthouse do
           migrate_without_penthouse(migrations_paths, target_version, &block)
         end
       end
 
       def up_with_penthouse(migrations_paths, target_version = nil, &block)
         unless Penthouse.configuration.migrate_tenants?
           return up_without_penthouse(migrations_paths, target_version, &block)
         end
 
-        Penthouse.each_tenant(tenant_identifiers: tenants_to_migrate) do |tenant|
-          up_without_penthouse(migrations_paths, target_version, &block)
+        wrap_penthouse do
+          up_without_penthouse(migrations_paths, target_version)
         end
       end
 
       def down_with_penthouse(migrations_paths, target_version = nil, &block)
         unless Penthouse.configuration.migrate_tenants?
-          return down_without_penthouse(migrations_paths, target_version, &block)
+          return down_without_penthouse(migrations_paths, target_version)
         end
 
-        Penthouse.each_tenant(tenant_identifiers: tenants_to_migrate) do |tenant|
-          down_without_penthouse(migrations_paths, target_version, &block)
+        wrap_penthouse do
+          down_without_penthouse(migrations_paths, target_version)
         end
       end
 
       def run_with_penthouse(direction, migrations_paths, target_version)
         unless Penthouse.configuration.migrate_tenants?
           return run_without_penthouse(direction, migrations_paths, target_version)
         end
 
-        Penthouse.each_tenant(tenant_identifiers: tenants_to_migrate) do |tenant|
+        wrap_penthouse do
           run_without_penthouse(direction, migrations_paths, target_version)
         end
       end
 
       private
+
+      def wrap_penthouse(&block)
+        if Penthouse.tenant
+          block.yield
+        else
+          Penthouse.each_tenant(tenant_identifiers: tenants_to_migrate, &block)
+        end
+      end
 
       def tenants_to_migrate
         return @tenants_to_migrate if defined?(@tenants_to_migrate)
         @tenants_to_migrate = begin
           if !!(t = (ENV["TENANT"] || ENV["TENANTS"]))