migration.rb in ar-octopus-0.10.0

- old
+ new

@@ -60,28 +60,30 @@
 end
 
 module Octopus
   module Migrator
     def self.included(base)
-      base.extend(ClassMethods)
+      unless Octopus.atleast_rails52?
+        base.extend(ClassMethods)
 
-      base.class_eval do
-        class << self
-          alias_method :migrate_without_octopus, :migrate
-          alias_method :migrate, :migrate_with_octopus
+        base.class_eval do
+          class << self
+              alias_method :migrate_without_octopus, :migrate
+              alias_method :migrate, :migrate_with_octopus
+              
+              alias_method :up_without_octopus, :up
+              alias_method :up, :up_with_octopus
 
-          alias_method :up_without_octopus, :up
-          alias_method :up, :up_with_octopus
+              alias_method :down_without_octopus, :down
+              alias_method :down, :down_with_octopus
 
-          alias_method :down_without_octopus, :down
-          alias_method :down, :down_with_octopus
+              alias_method :run_without_octopus, :run
+              alias_method :run, :run_with_octopus
 
-          alias_method :run_without_octopus, :run
-          alias_method :run, :run_with_octopus
-
-          alias_method :rollback_without_octopus, :rollback
-          alias_method :rollback, :rollback_with_octopus
+              alias_method :rollback_without_octopus, :rollback
+              alias_method :rollback, :rollback_with_octopus
+          end
         end
       end
 
       base.send :alias_method, :run_without_octopus, :run
       base.send :alias_method, :run, :run_with_octopus
@@ -90,77 +92,116 @@
       base.send :alias_method, :migrate, :migrate_with_octopus
 
       base.send :alias_method, :migrations_without_octopus, :migrations
       base.send :alias_method, :migrations, :migrations_with_octopus
     end
+    if Octopus.atleast_rails52?
+      ### Post RAILS 5.2 Migration methods
+      
+      def run_with_octopus(&block)
+        return run_without_octopus(&block) unless connection.is_a?(Octopus::Proxy)
+        shards = migrations.map(&:shards).flatten.map(&:to_s)
+        connection.send_queries_to_multiple_shards(shards) do
+          run_without_octopus(&block)
+        end
+        rescue ActiveRecord::UnknownMigrationVersionError => e
+          raise unless migrations(true).detect { |m| m.version == e.version }
+      end
 
-    def run_with_octopus(&block)
-      run_without_octopus(&block)
-    rescue ActiveRecord::UnknownMigrationVersionError => e
-      raise unless migrations(true).detect { |m| m.version == e.version }
-    end
+      def migrate_with_octopus(&block)
+        return migrate_without_octopus(&block) unless connection.is_a?(Octopus::Proxy)
+        shards = migrations.map(&:shards).flatten.map(&:to_s)
+        connection.send_queries_to_multiple_shards(shards) do
+          migrate_without_octopus(&block)
+        end
+        rescue ActiveRecord::UnknownMigrationVersionError => e
+          raise unless migrations(true).detect { |m| m.version == e.version }
+      end
 
-    def migrate_with_octopus(&block)
-      migrate_without_octopus(&block)
-    rescue ActiveRecord::UnknownMigrationVersionError => e
-      raise unless migrations(true).detect { |m| m.version == e.version }
-    end
+      def migrations_with_octopus(shard_agnostic = true)
+        migrations = migrations_without_octopus
+        return migrations if !connection.is_a?(Octopus::Proxy) || shard_agnostic
 
-    def migrations_with_octopus(shard_agnostic = false)
-      connection = ActiveRecord::Base.connection
-      migrations = migrations_without_octopus
-      return migrations if !connection.is_a?(Octopus::Proxy) || shard_agnostic
+        migrations.select { |m| m.shards.include?(connection.current_shard.to_sym) }
+      end
+      
+      private 
+      
+      def connection
+        ActiveRecord::Base.connection
+      end
+      
+    else 
+      ### Pre RAILS 5.2 Migration methods
+      
+      def run_with_octopus(&block)
+        run_without_octopus(&block)
+      rescue ActiveRecord::UnknownMigrationVersionError => e
+        raise unless migrations(true).detect { |m| m.version == e.version }
+      end
 
-      migrations.select { |m| m.shards.include?(connection.current_shard.to_sym) }
-    end
+      def migrate_with_octopus(&block)
+        migrate_without_octopus(&block)
+      rescue ActiveRecord::UnknownMigrationVersionError => e
+        raise unless migrations(true).detect { |m| m.version == e.version }
+      end
 
-    module ClassMethods
-      def migrate_with_octopus(migrations_paths, target_version = nil, &block)
-        return migrate_without_octopus(migrations_paths, target_version, &block) unless connection.is_a?(Octopus::Proxy)
+      def migrations_with_octopus(shard_agnostic = false)
+        connection = ActiveRecord::Base.connection
+        migrations = migrations_without_octopus
+        return migrations if !connection.is_a?(Octopus::Proxy) || shard_agnostic
 
-        connection.send_queries_to_multiple_shards(connection.shard_names) do
-          migrate_without_octopus(migrations_paths, target_version, &block)
-        end
+        migrations.select { |m| m.shards.include?(connection.current_shard.to_sym) }
       end
 
-      def up_with_octopus(migrations_paths, target_version = nil, &block)
-        return up_without_octopus(migrations_paths, target_version, &block) unless connection.is_a?(Octopus::Proxy)
-        return up_without_octopus(migrations_paths, target_version, &block) unless connection.current_shard.to_s == Octopus.master_shard.to_s
+      module ClassMethods
+        def migrate_with_octopus(migrations_paths, target_version = nil, &block)
+          return migrate_without_octopus(migrations_paths, target_version, &block) unless connection.is_a?(Octopus::Proxy)
 
-        connection.send_queries_to_multiple_shards(connection.shard_names) do
-          up_without_octopus(migrations_paths, target_version, &block)
+          connection.send_queries_to_multiple_shards(connection.shard_names) do
+            migrate_without_octopus(migrations_paths, target_version, &block)
+          end
         end
-      end
 
-      def down_with_octopus(migrations_paths, target_version = nil, &block)
-        return down_without_octopus(migrations_paths, target_version, &block) unless connection.is_a?(Octopus::Proxy)
-        return down_without_octopus(migrations_paths, target_version, &block) unless connection.current_shard.to_s == Octopus.master_shard.to_s
+        def up_with_octopus(migrations_paths, target_version = nil, &block)
+          return up_without_octopus(migrations_paths, target_version, &block) unless connection.is_a?(Octopus::Proxy)
+          return up_without_octopus(migrations_paths, target_version, &block) unless connection.current_shard.to_s == Octopus.master_shard.to_s
 
-        connection.send_queries_to_multiple_shards(connection.shard_names) do
-          down_without_octopus(migrations_paths, target_version, &block)
+          connection.send_queries_to_multiple_shards(connection.shard_names) do
+            up_without_octopus(migrations_paths, target_version, &block)
+          end
         end
-      end
 
-      def run_with_octopus(direction, migrations_paths, target_version)
-        return run_without_octopus(direction, migrations_paths, target_version) unless connection.is_a?(Octopus::Proxy)
+        def down_with_octopus(migrations_paths, target_version = nil, &block)
+          return down_without_octopus(migrations_paths, target_version, &block) unless connection.is_a?(Octopus::Proxy)
+          return down_without_octopus(migrations_paths, target_version, &block) unless connection.current_shard.to_s == Octopus.master_shard.to_s
 
-        connection.send_queries_to_multiple_shards(connection.shard_names) do
-          run_without_octopus(direction, migrations_paths, target_version)
+          connection.send_queries_to_multiple_shards(connection.shard_names) do
+            down_without_octopus(migrations_paths, target_version, &block)
+          end
         end
-      end
 
-      def rollback_with_octopus(migrations_paths, steps = 1)
-        return rollback_without_octopus(migrations_paths, steps) unless connection.is_a?(Octopus::Proxy)
+        def run_with_octopus(direction, migrations_paths, target_version)
+          return run_without_octopus(direction, migrations_paths, target_version) unless connection.is_a?(Octopus::Proxy)
 
-        connection.send_queries_to_multiple_shards(connection.shard_names) do
-          rollback_without_octopus(migrations_paths, steps)
+          connection.send_queries_to_multiple_shards(connection.shard_names) do
+            run_without_octopus(direction, migrations_paths, target_version)
+          end
         end
-      end
 
-      private
+        def rollback_with_octopus(migrations_paths, steps = 1)
+          return rollback_without_octopus(migrations_paths, steps) unless connection.is_a?(Octopus::Proxy)
 
-      def connection
-        ActiveRecord::Base.connection
+          connection.send_queries_to_multiple_shards(connection.shard_names) do
+            rollback_without_octopus(migrations_paths, steps)
+          end
+        end
+
+        private
+
+        def connection
+          ActiveRecord::Base.connection
+        end
       end
     end
   end
 end
 