require 'active_support/deprecation' module Octopus module Model def self.extended(base) base.send(:include, Octopus::ShardTracking::Attribute) base.send(:include, InstanceMethods) base.extend(ClassMethods) 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. If you are trying to scope everything to a specific shard, use Octopus.using instead. EOF end if Octopus.enabled? Octopus::ScopeProxy.new(shard, self) else self 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, :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 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 perform_validations_without_octopus(*args) end end end 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 should_use_normal_connection? if !Octopus.enabled? true elsif custom_octopus_connection !connection_proxy.block || !allowed_shard?(connection_proxy.current_shard) end end 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 connection_proxy.current_model = self connection_proxy end end def connection_pool_with_octopus if should_use_normal_connection? connection_pool_without_octopus else connection_proxy.connection_pool end end def clear_active_connections_with_octopus! if should_use_normal_connection? clear_active_connections_without_octopus! else connection_proxy.clear_active_connections! end end def clear_all_connections_with_octopus! if should_use_normal_connection? clear_all_connections_without_octopus! else connection_proxy.clear_all_connections! end end def connected_with_octopus? if should_use_normal_connection? connected_without_octopus? else connection_proxy.connected? 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 end ActiveRecord::Base.extend(Octopus::Model)