ar-octopus 0.8.1 → 0.10.2

data/lib/octopus/model.rb

@@ -1,70 +1,169 @@
 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?
+module Octopus
+  module Model
+    def self.extended(base)
+      base.send(:include, Octopus::ShardTracking::Attribute)
+      base.send(:include, InstanceMethods)
+      base.extend(ClassMethods)
+    end
 
-      if self != ActiveRecord::Base && self.respond_to?(:reset_table_name) && !self.custom_octopus_table_name
-        self.reset_table_name()
-      end
+    module SharedMethods
+      def using(shard)
+        if block_given?
+          raise Octopus::Exception, <<-EOF
+#{name}.using is not allowed to receive a block, it works just like a regular scope.
 
-      self.reset_column_information
-      self.instance_variable_set(:@quoted_table_name, nil)
-    end
+If you are trying to scope everything to a specific shard, use Octopus.using instead.
+          EOF
+        end
 
-    def using(shard)
-      if Octopus.enabled?
-        clean_table_name
-        Octopus::ScopeProxy.new(shard, self)
-      else
-        self
+        if Octopus.enabled?
+          Octopus::ScopeProxy.new(shard, self)
+        else
+          self
+        end
       end
     end
 
-    def hijack_initializer()
-      attr_accessor :current_shard
-      around_save :run_on_shard
+    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?, :==)
+        base.send(:alias_method, :perform_validations_without_octopus, :perform_validations)
+        base.send(:alias_method, :perform_validations, :perform_validations_with_octopus)
+      end
 
       def set_current_shard
         return unless Octopus.enabled?
+        shard = self.class.connection_proxy.current_shard
+        self.current_shard = shard if self.class.allowed_shard?(shard)
+      end
+
+      def init_with(coder)
+        obj = super
+
+        return obj unless Octopus.enabled?
+        return obj if obj.class.connection_proxy.current_model_replicated?
+
+        current_shard_value = coder['attributes']['current_shard'].value if coder['attributes']['current_shard'].present? && coder['attributes']['current_shard'].value.present?
+
+        coder['attributes'].send(:attributes).send(:values).delete('current_shard')
+        coder['attributes'].send(:attributes).send(:delegate_hash).delete('current_shard')
+
+        obj.current_shard = current_shard_value if current_shard_value.present?
+        obj
+      end
+
+      def should_set_current_shard?
+        self.respond_to?(:current_shard) && !current_shard.nil?
+      end
+
+      def equality_with_octopus(comparison_object)
+        equality_without_octopus(comparison_object) && comparison_object.current_shard.to_s == current_shard.to_s
+      end
 
-        if new_record? || self.class.connection_proxy.block
-          self.current_shard = self.class.connection_proxy.current_shard
+      def perform_validations_with_octopus(*args)
+        if Octopus.enabled? && should_set_current_shard?
+          Octopus.using(current_shard) do
+            perform_validations_without_octopus(*args)
+          end
         else
-          self.current_shard = self.class.connection_proxy.last_current_shard || self.class.connection_proxy.current_shard
+          perform_validations_without_octopus(*args)
         end
       end
-
-      after_initialize :set_current_shard
     end
 
-    def hijack_connection()
-      def self.should_use_normal_connection?
-        !Octopus.enabled? || self.custom_octopus_connection
+    module ClassMethods
+      include SharedMethods
+
+      def self.extended(base)
+        base.class_attribute(:replicated)
+        base.class_attribute(:sharded)
+        base.class_attribute(:allowed_shards)
+        base.hijack_methods
+      end
+
+      def replicated_model
+        self.replicated = true
+      end
+
+      def sharded_model
+        self.sharded = true
+      end
+
+      def allow_shard(*shards)
+        self.allowed_shards ||= []
+        self.allowed_shards += shards
+      end
+
+      def hijack_methods
+        after_initialize :set_current_shard
+
+        around_save :run_on_shard, :unless => lambda { self.class.custom_octopus_connection }
+
+        class_attribute :custom_octopus_connection
+
+        class << self
+          attr_accessor :custom_octopus_table_name
+
+          alias_method :connection_without_octopus, :connection
+          alias_method :connection, :connection_with_octopus
+
+          alias_method :connection_pool_without_octopus, :connection_pool
+          alias_method :connection_pool, :connection_pool_with_octopus
+
+          alias_method :clear_all_connections_without_octopus!, :clear_all_connections!
+          alias_method :clear_all_connections!, :clear_all_connections_with_octopus!
+
+          alias_method :clear_active_connections_without_octopus!, :clear_active_connections!
+          alias_method :clear_active_connections!, :clear_active_connections_with_octopus!
+
+          alias_method :connected_without_octopus?, :connected?
+          alias_method :connected?, :connected_with_octopus?
+
+          def table_name=(value = nil)
+            self.custom_octopus_table_name = true
+            super
+          end
+        end
+      end
+
+      def connection_proxy
+        ActiveRecord::Base.class_variable_defined?(:@@connection_proxy) &&
+          ActiveRecord::Base.class_variable_get(:@@connection_proxy) ||
+          ActiveRecord::Base.class_variable_set(:@@connection_proxy, Octopus::Proxy.new)
       end
 
-      def self.connection_proxy
-        @@connection_proxy ||= Octopus::Proxy.new
+      def should_use_normal_connection?
+        if !Octopus.enabled?
+          true
+        elsif custom_octopus_connection
+          !connection_proxy.block || !allowed_shard?(connection_proxy.current_shard)
+        end
       end
 
-      def self.connection_with_octopus
+      def allowed_shard?(shard)
+        if custom_octopus_connection
+          allowed_shards && shard && (allowed_shards.include?(shard.to_s) || allowed_shards.include?(shard.to_sym))
+        else
+          true
+        end
+      end
+
+      def connection_with_octopus
         if should_use_normal_connection?
           connection_without_octopus
         else
-          self.connection_proxy.current_model = self
-          self.connection_proxy
+          connection_proxy.current_model = self
+          connection_proxy
         end
       end
 
-      def self.connection_pool_with_octopus
+      def connection_pool_with_octopus
         if should_use_normal_connection?
           connection_pool_without_octopus
         else
@@ -72,7 +171,7 @@ module Octopus::Model
         end
       end
 
-      def self.clear_active_connections_with_octopus!
+      def clear_active_connections_with_octopus!
         if should_use_normal_connection?
           clear_active_connections_without_octopus!
         else
@@ -80,7 +179,7 @@ module Octopus::Model
         end
       end
 
-      def self.clear_all_connections_with_octopus!
+      def clear_all_connections_with_octopus!
         if should_use_normal_connection?
           clear_all_connections_without_octopus!
         else
@@ -88,7 +187,7 @@ module Octopus::Model
         end
       end
 
-      def self.connected_with_octopus?
+      def connected_with_octopus?
         if should_use_normal_connection?
           connected_without_octopus?
         else
@@ -96,100 +195,21 @@ module Octopus::Model
         end
       end
 
-      class << self
-        alias_method_chain :connection, :octopus
-        alias_method_chain :connection_pool, :octopus
-        alias_method_chain :clear_all_connections!, :octopus
-        alias_method_chain :clear_active_connections!, :octopus
-        alias_method_chain :connected?, :octopus
+      def set_table_name_with_octopus(value = nil, &block)
+        self.custom_octopus_table_name = true
+        set_table_name_without_octopus(value, &block)
       end
-    end
-  end
-
-  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?, :==)
-      base.send(:alias_method_chain, :perform_validations, :octopus)
-    end
-
-    def should_set_current_shard?
-      self.respond_to?(:current_shard) && !self.current_shard.nil?
-    end
-
-    def run_on_shard(&block)
-      if self.current_shard
-        self.class.connection_proxy.run_queries_on_shard(self.current_shard, &block)
-      else
-        yield
-      end
-    end
-
-    def equality_with_octopus(comparison_object)
-      equality_without_octopus(comparison_object) && comparison_object.current_shard == current_shard
-    end
-
-    def perform_validations_with_octopus(*args)
-      if Octopus.enabled? and should_set_current_shard?
-        Octopus.using(self.current_shard) do
-          perform_validations_without_octopus(*args)
-        end
-      else
-        perform_validations_without_octopus(*args)
+      def octopus_establish_connection(spec = ENV['DATABASE_URL'])
+        self.custom_octopus_connection = true if spec
+        establish_connection(spec)
       end
-    end
-  end
-
-  module ClassMethods
-    include SharedMethods
-
-    def self.extended(base)
-      base.class_attribute(:replicated)
-      base.class_attribute(:sharded)
-      base.hijack_methods
-    end
-
-    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
-
-        if Octopus.rails3?
-          alias_method_chain(:set_table_name, :octopus)
-        end
-
-        def table_name=(value = nil)
-          self.custom_octopus_table_name = true
-          super
-        end
+      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
-
-    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)
-      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
 

data/lib/octopus/persistence.rb

@@ -0,0 +1,45 @@
+module Octopus
+  module Persistence
+    def update_attribute(*args)
+      run_on_shard { super }
+    end
+
+    def update_attributes(*args)
+      run_on_shard { super }
+    end
+
+    def update_attributes!(*args)
+      run_on_shard { super }
+    end
+
+    def reload(*args)
+      run_on_shard { super }
+    end
+
+    def delete
+      run_on_shard { super }
+    end
+
+    def destroy
+      run_on_shard { super }
+    end
+
+    def touch(*args)
+      run_on_shard { super }
+    end
+
+    def update_column(*args)
+      run_on_shard { super }
+    end
+
+    def increment!(*args)
+      run_on_shard { super }
+    end
+
+    def decrement!(*args)
+      run_on_shard { super }
+    end
+  end
+end
+
+ActiveRecord::Base.send(:include, Octopus::Persistence)