ar-octopus 0.4.0 → 0.5.0

data/lib/octopus/has_and_belongs_to_many_association.rb

@@ -4,12 +4,12 @@ module Octopus::HasAndBelongsToManyAssociation
       alias_method_chain :insert_record, :octopus
     end
   end
-  
+
   def insert_record_with_octopus(record, force = true, validate = true)
     if should_wrap_the_connection?
-      Octopus.using(@owner.current_shard) { insert_record_without_octopus(record, force, validate) } 
-    else 
-      insert_record_without_octopus(record, force, validate) 
+      Octopus.using(@owner.current_shard) { insert_record_without_octopus(record, force, validate) }
+    else
+      insert_record_without_octopus(record, force, validate)
     end
   end
 end

data/lib/octopus/logger.rb

@@ -1,13 +1,18 @@
 require "logger"
 
 class Octopus::Logger < Logger
+  def initialize(logdev, shift_age = 0, shift_size = 1048576)
+    ActiveSupport::Deprecation.warn "Octopus::Logger is deprecated and will be removed in Octopus 0.6.x", caller
+    super
+  end
+
   def format_message(severity, timestamp, progname, msg)
-    str = super 
-    
+    str = super
+
     if ActiveRecord::Base.connection.respond_to?(:current_shard)
-      str += "Shard: #{ActiveRecord::Base.connection.current_shard} -" 
+      str += "Shard: #{ActiveRecord::Base.connection.current_shard} -"
     end
-    
+
     str
   end
 end

data/lib/octopus/migration.rb

@@ -1,73 +1,178 @@
-module Octopus::Migration  
-  def self.extended(base)
-    class << base
-      def announce_with_octopus(message)
-        announce_without_octopus("#{message} - #{get_current_shard}")
-      end
-      
-      alias_method_chain :migrate, :octopus
-      alias_method_chain :announce, :octopus
-      attr_accessor :current_shard
+require "set"
+require "active_support/core_ext/module/aliasing"
+
+if Octopus.rails2?
+  require "active_support/core_ext/array/wrapper"
+else
+  require "active_support/core_ext/array/wrap"
+end
+
+module Octopus::Migration
+  module InstanceOrClassMethods
+    def announce_with_octopus(message)
+      announce_without_octopus("#{message} - #{get_current_shard}")
+    end
+
+    def get_current_shard
+      "Shard: #{connection.current_shard}" if connection.respond_to?(:current_shard)
     end
   end
 
-  def using(*args)
-    if self.connection().is_a?(Octopus::Proxy)
-      args.each do |shard|
-        self.connection().check_schema_migrations(shard)
+  include InstanceOrClassMethods if Octopus.rails31? || Octopus.rails32?
+
+  def self.included(base)
+    base.send(:extend, ClassMethods)
+
+    if Octopus.rails31? || Octopus.rails32?
+      base.alias_method_chain :announce, :octopus
+    else
+      base.class_eval do
+        class << self
+          alias_method_chain :announce, :octopus
+        end
       end
+    end
+
+    base.class_attribute :current_shard, :current_group, :instance_reader => false, :instance_writer => false
+  end
+
+  module ClassMethods
+    include InstanceOrClassMethods unless Octopus.rails31? || Octopus.rails32?
+
+    def using(*args)
+      return self unless connection.is_a?(Octopus::Proxy)
 
-      self.connection().block = true
       self.current_shard = args
-      self.connection().current_shard = args        
+      self
     end
 
-    return self
-  end
+    def using_group(*groups)
+      return self unless connection.is_a?(Octopus::Proxy)
+
+      self.current_group = groups
+      self
+    end
 
-  def using_group(*args)
-    if self.connection().is_a?(Octopus::Proxy)
-      args.each do |group_shard|
-        shards = self.connection().instance_variable_get(:@groups)[group_shard] || []
+    def shards
+      shards = Set.new
 
-        shards.each do |shard|
-          self.connection().check_schema_migrations(shard)
+      if groups = current_group
+        Array.wrap(groups).each do |group|
+          group_shards = connection.shards_for_group(group)
+          shards.merge(group_shards) if group_shards
         end
+      elsif shard = current_shard
+        shards.merge(Array.wrap(shard))
       end
 
-      self.connection().block = true
-      self.connection().current_group = args
+      shards.to_a.presence || [:master]
     end
-    
-    return self
   end
-  
-  def get_current_shard
-    "Shard: #{ActiveRecord::Base.connection.current_shard()}" if ActiveRecord::Base.connection.respond_to?(:current_shard)
+end
+
+module Octopus::Migrator
+  def self.included(base)
+    base.send(:extend, ClassMethods)
+
+    base.class_eval do
+      class << self
+        alias_method_chain :migrate, :octopus
+        alias_method_chain :up, :octopus
+        alias_method_chain :down, :octopus
+        alias_method_chain :run, :octopus
+      end
+    end
+
+    base.alias_method_chain :run, :octopus
+    base.alias_method_chain :migrate, :octopus
+    base.alias_method_chain :migrations, :octopus
+  end
+
+  def run_with_octopus(&block)
+    run_without_octopus(&block)
+  rescue ActiveRecord::UnknownMigrationVersionError => e
+    raise unless migrations(true).find {|m| m.version == e.version}
+  end
+
+  def migrate_with_octopus(&block)
+    migrate_without_octopus(&block)
+  rescue ActiveRecord::UnknownMigrationVersionError => e
+    raise unless migrations(true).find {|m| m.version == e.version}
   end
 
+  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
 
-  def migrate_with_octopus(direction)
-    conn = ActiveRecord::Base.connection
-    return migrate_without_octopus(direction) unless conn.is_a?(Octopus::Proxy)
-    self.connection().current_shard = self.current_shard if self.current_shard != nil
-    
-    groups = conn.instance_variable_get(:@groups)
-    
-    begin
-      if conn.current_group.is_a?(Array)
-        conn.current_group.each { |group| conn.send_queries_to_multiple_shards(groups[group]) { migrate_without_octopus(direction) } } 
-      elsif conn.current_group.is_a?(Symbol)       
-        conn.send_queries_to_multiple_shards(groups[conn.current_group]) { migrate_without_octopus(direction) }     
-      elsif conn.current_shard.is_a?(Array)
-        conn.send_queries_to_multiple_shards(conn.current_shard) { migrate_without_octopus(direction) }     
-      else
-        migrate_without_octopus(direction)
+    migrations.select {|m| m.shards.include?(connection.current_shard.to_sym)}
+  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)
+
+      connection.send_queries_to_multiple_shards(connection.shard_names) do
+        migrate_without_octopus(migrations_paths, target_version, &block)
+      end
+    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 == :master
+
+      connection.send_queries_to_multiple_shards(connection.shard_names) do
+        up_without_octopus(migrations_paths, target_version, &block)
+      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 == :master
+
+      connection.send_queries_to_multiple_shards(connection.shard_names) do
+        down_without_octopus(migrations_paths, target_version, &block)
+      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)
+
+      connection.send_queries_to_multiple_shards(connection.shard_names) do
+        run_without_octopus(direction, migrations_paths, target_version)
       end
-    ensure
-      conn.clean_proxy()
     end
+
+    private
+    def connection
+      ActiveRecord::Base.connection
+    end
+  end
+end
+
+module Octopus::MigrationProxy
+  def shards
+    if Octopus.rails31? || Octopus.rails32?
+      migration.class.shards
+    else
+      migration.shards
+    end
+  end
+end
+
+module Octopus::UnknownMigrationVersionError
+  def self.included(base)
+    base.alias_method_chain :initialize, :octopus
+    base.send(:attr_accessor, :version)
+  end
+
+  def initialize_with_octopus(version)
+    @version = version
+    initialize_without_octopus(version)
   end
 end
 
-ActiveRecord::Migration.extend(Octopus::Migration)
+ActiveRecord::Migration.send(:include, Octopus::Migration)
+ActiveRecord::Migrator.send(:include, Octopus::Migrator)
+ActiveRecord::MigrationProxy.send(:include, Octopus::MigrationProxy)
+ActiveRecord::UnknownMigrationVersionError.send(:include, Octopus::UnknownMigrationVersionError)

data/lib/octopus/model.rb

@@ -1,15 +1,19 @@
-module Octopus::Model  
-  def self.extended(base) 
+require 'active_support/deprecation'
+
+module Octopus::Model
+  def self.extended(base)
     base.send(:include, InstanceMethods)
     base.extend(ClassMethods)
     base.hijack_connection()
+    base.hijack_initializer()
   end
 
   module SharedMethods
     def clean_table_name
       return unless self.connection_proxy.should_clean_table_name?
-      if self != ActiveRecord::Base && self.respond_to?(:reset_table_name) && !self.read_inheritable_attribute(:set_table_name)
-        self.reset_table_name() 
+
+      if self != ActiveRecord::Base && self.respond_to?(:reset_table_name) && !self.custom_octopus_table_name
+        self.reset_table_name()
       end
 
       if Octopus.rails3?
@@ -19,14 +23,12 @@ module Octopus::Model
     end
 
     def using(shard)
-      return self if defined?(::Rails) && !Octopus.environments.include?(Rails.env.to_s)
-      
-      hijack_initializer() if !respond_to?(:set_current_shard)
-      clean_table_name()
-
-      self.connection_proxy.using_enabled = true
-
-      return Octopus::ScopeProxy.new(shard, self)
+      if Octopus.enabled?
+        clean_table_name
+        Octopus::ScopeProxy.new(shard, self)
+      else
+        self
+      end
     end
 
     def hijack_initializer()
@@ -34,10 +36,12 @@ module Octopus::Model
       before_save :reload_connection
 
       def set_current_shard
+        return unless Octopus.enabled?
+
         if new_record? || self.class.connection_proxy.block
-          self.current_shard = self.class.connection_proxy.current_shard    
+          self.current_shard = self.class.connection_proxy.current_shard
         else
-          self.current_shard = self.class.connection_proxy.last_current_shard  
+          self.current_shard = self.class.connection_proxy.last_current_shard || self.class.connection_proxy.current_shard
         end
       end
 
@@ -52,22 +56,33 @@ module Octopus::Model
 
     def hijack_connection()
       def self.should_use_normal_connection?
-        (defined?(Rails) && Octopus.config() && !Octopus.environments.include?(Rails.env.to_s)) || self.read_inheritable_attribute(:establish_connection)
+        !Octopus.enabled? || self.custom_octopus_connection
       end
-      
+
       def self.connection_proxy
-        Thread.current[:connection_proxy] ||= Octopus::Proxy.new(Octopus.config())
+        @@connection_proxy ||= Octopus::Proxy.new
       end
 
-      def self.connection_with_octopus()
-        return connection_without_octopus() if should_use_normal_connection?
+      def self.connection_with_octopus
+        if should_use_normal_connection?
+          connection_without_octopus
+        else
+          self.connection_proxy.current_model = self
+          self.connection_proxy
+        end
+      end
 
-        self.connection_proxy().current_model = self
-        self.connection_proxy()
+      def self.connection_pool_with_octopus
+        if should_use_normal_connection?
+          connection_pool_without_octopus
+        else
+          connection_proxy.connection_pool
+        end
       end
 
       class << self
         alias_method_chain :connection, :octopus
+        alias_method_chain :connection_pool, :octopus
       end
     end
   end
@@ -75,36 +90,85 @@ module Octopus::Model
   module InstanceMethods
     include SharedMethods
 
+    def self.included(base)
+      base.send(:alias_method, :equality_without_octopus, :==)
+      base.send(:alias_method, :==, :equality_with_octopus)
+      base.send(:alias_method, :eql?, :==)
+    end
+
     def should_set_current_shard?
       self.respond_to?(:current_shard) && !self.current_shard.nil?
     end
 
+    def reload_connection_safe(&block)
+      return yield unless should_set_current_shard?
+      original = self.class.connection_proxy.current_shard
+      self.class.connection_proxy.current_shard = self.current_shard
+      result = yield
+      self.class.connection_proxy.current_shard = original
+      result
+    end
+
     def reload_connection()
-      self.class.connection_proxy.current_shard = self.current_shard() if should_set_current_shard?
+      return unless should_set_current_shard?
+      self.class.connection_proxy.current_shard = self.current_shard
+    end
+
+    def equality_with_octopus(comparison_object)
+      equality_without_octopus(comparison_object) && comparison_object.current_shard == current_shard
     end
   end
 
   module ClassMethods
     include SharedMethods
 
-    def replicated_model()
-      write_inheritable_attribute(:replicated, true)
+    def self.extended(base)
+      base.class_attribute(:replicated)
+      base.class_attribute(:sharded)
+      base.hijack_methods
     end
-    
-    def sharded_model()
-      write_inheritable_attribute(:sharded, true)      
-    end    
-    
-    def octopus_establish_connection(spec = nil)
-      write_inheritable_attribute(:establish_connection, true)      
+
+    def replicated_model
+      self.replicated = true
+    end
+
+    def sharded_model
+      self.sharded = true
+    end
+
+    def hijack_methods
+      class << self
+        attr_accessor :custom_octopus_connection
+        attr_accessor :custom_octopus_table_name
+
+        alias_method_chain(:set_table_name, :octopus)
+
+        if Octopus.rails32?
+          def table_name=(value = nil)
+            self.custom_octopus_table_name = true
+            super
+          end
+        end
+      end
+    end
+
+    def set_table_name_with_octopus(value = nil, &block)
+      self.custom_octopus_table_name = true
+      set_table_name_without_octopus(value, &block)
+    end
+
+    def octopus_establish_connection(spec = ENV['DATABASE_URL'])
+      self.custom_octopus_connection = true if spec
       establish_connection(spec)
     end
 
-    def octopus_set_table_name(value = nil, &block)
-      write_inheritable_attribute(:set_table_name, true)      
-      set_table_name(value, &block)
+    def octopus_set_table_name(value = nil)
+      ActiveSupport::Deprecation.warn "Calling `octopus_set_table_name` is deprecated and will be removed in Octopus 1.0.", caller
+      set_table_name(value)
     end
   end
 end
 
 ActiveRecord::Base.extend(Octopus::Model)
+
+class OctopusModel < ActiveRecord::Base; end;