ar-octopus 0.8.5 → 0.8.6

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA1:
-  metadata.gz: a2e171072013526c5597b2d233d983f129a84cec
-  data.tar.gz: 0c05c12e05f5dfb64aaf6f15e0accb59c2174a53
+  metadata.gz: 28a5543a3a9dc4dafa60e8d52c6b35d8064180eb
+  data.tar.gz: 1212d62583a96614b636058f7919c75fdaa0df14
 SHA512:
-  metadata.gz: 19f50eb815a6273a87430354c5124096d90d561a8b27c1bfe0073a5a81f3a2242c71f099b88db0bab031018c56a8662ed587122412a0042ba543bc41c6b0d745
-  data.tar.gz: 11828df91744a59445fe6a4add2c6081971f724ed9394b96eb12b36c6a09e69db2546d7dc1eadc56c19074ed5de1d7cd28f39a38c0bd7d3ad88030238a46a98e
+  metadata.gz: 3b8085b83e646f726fc68553aaec14cb0663d3bcca5e2412b8059712e815737c6e6b23a284a9944e37fccfa19a916b157c0efefc7cb0c039af37a65d09f538f6
+  data.tar.gz: e3bde2a286faef9f0c11531f73fe7cfaf5351a8a9cc354fc9fbf60921901eb300c82ba0ca27a11cae7d92b4f3fc06432573ebabfa42dedbca9f39391ca229aaa

data/.travis.yml

@@ -5,9 +5,10 @@ before_script:
   - "bundle exec rake db:prepare"
 rvm:
   - 2.0.0
-  - 2.1.5
-  - 2.2.0
+  - 2.1.6
+  - 2.2.2
 gemfile:
+  - gemfiles/rails32.gemfile
   - gemfiles/rails4.gemfile
   - gemfiles/rails41.gemfile
   - gemfiles/rails42.gemfile
@@ -15,7 +16,3 @@ notifications:
   recipients:
     - gabriel.sobrinho@gmail.com
     - thiago.pradi@gmail.com
-matrix:
-  include:
-    - rvm: 1.9.3
-      gemfile: gemfiles/rails32.gemfile

data/README.mkdn

@@ -1,6 +1,6 @@
 # Octopus  - Easy Database Sharding for ActiveRecord
 
-<a href='http://www.pledgie.com/campaigns/20950'><img alt='Click here to lend your support to: Octopus and make a donation at www.pledgie.com !' src='http://www.pledgie.com/campaigns/20950.png?skin_name=chrome' border='0' /></a> [![Build Status](https://travis-ci.org/tchandy/octopus.png)](https://travis-ci.org/tchandy/octopus) [![Code Climate](https://codeclimate.com/badge.png)](https://codeclimate.com/github/tchandy/octopus)
+<a href='http://www.pledgie.com/campaigns/20950'><img alt='Click here to lend your support to: Octopus and make a donation at www.pledgie.com !' src='http://www.pledgie.com/campaigns/20950.png?skin_name=chrome' border='0' /></a> [![Build Status](https://travis-ci.org/thiagopradi/octopus.png)](https://travis-ci.org/thiagopradi/octopus) [![Code Climate](https://codeclimate.com/badge.png)](https://codeclimate.com/github/thiagopradi/octopus)
 
 Octopus is a better way to do Database Sharding in ActiveRecord. Sharding allows multiple databases in the same rails application. While there are several projects that implement Sharding (e.g. DbCharmer, DataFabric, MultiDb), each project has its own limitations. The main goal of octopus project is to provide a better way of doing Database Sharding.
 
@@ -15,15 +15,15 @@ Octopus supports:
 - Tools to manage database configurations. (soon)
 
 ### Replication
-When using replication, all writes queries will be sent to master, and read queries to slaves. More info could be found at: <a href="http://wiki.github.com/tchandy/octopus/replication"> Wiki</a>
+When using replication, all writes queries will be sent to master, and read queries to slaves. More info could be found at: <a href="http://wiki.github.com/thiagopradi/octopus/replication"> Wiki</a>
 
 ### Sharding
- When using sharding, you need to specify which shard to send the query. Octopus supports selecting the shard inside a controller, or manually in each object. More could be found at <a href="http://wiki.github.com/tchandy/octopus/sharding"> Wiki</a>
+ When using sharding, you need to specify which shard to send the query. Octopus supports selecting the shard inside a controller, or manually in each object. More could be found at <a href="http://wiki.github.com/thiagopradi/octopus/sharding"> Wiki</a>
 
 ### Replication + Sharding
- When using replication and sharding concurrently, you must specify a shard, and can optionally specify a <a href="https://github.com/tchandy/octopus/wiki/Slave-Groups">slave group</a>.
+ When using replication and sharding concurrently, you must specify a shard, and can optionally specify a <a href="https://github.com/thiagopradi/octopus/wiki/Slave-Groups">slave group</a>.
  All write queries will be sent to each shard's master. If the slave group is specified read queries will be sent to slaves in it, or else to shard's master.
- More info could be found at <a href="https://github.com/tchandy/octopus/wiki/Slave-Groups"> Wiki</a>
+ More info could be found at <a href="https://github.com/thiagopradi/octopus/wiki/Slave-Groups"> Wiki</a>
 
 ## Install
 
@@ -55,7 +55,7 @@ going forward.
 
 ## How to use Octopus?
 
-First, you need to create a config file, shards.yml, inside your config/ directory. to see the syntax and how this file should look, please checkout <a href="http://wiki.github.com/tchandy/octopus/config-file">this page on wiki</a>.
+First, you need to create a config file, shards.yml, inside your config/ directory. to see the syntax and how this file should look, please checkout <a href="http://wiki.github.com/thiagopradi/octopus/config-file">this page on wiki</a>.
 
 ### Syntax
 
@@ -73,6 +73,20 @@ Octopus.using(:slave_two) do
 end
 ```
 
+If you want to use the same code for all shards or all shards in a specific group (for example in `db/seeds.rb`), you can use this syntax.
+
+```ruby
+# This will return a list of the given block's results, per shard.
+Octopus.using_all do
+  User.create_from_csv!
+end
+
+# This will return a list of the given block's results, per shard in history_shards group.
+Octopus.using_group(:history_shards) do
+  HistoryCategory.create_from_csv!
+end
+```
+
 Each model instance knows which shard it came from so this will work automatically:
 
 ```ruby
@@ -131,6 +145,19 @@ octopus:
 
 There is no need for a corresponding `default_migration_shard` - simply define that database to be your master. You might want this setting if all of your databases have identical schemas, but are not replicated.
 
+You can configure a master shard for the rails application, to connect to when rails is going up. The safest would be to configure this to the shard specified in `database.yml` (some things still use it).
+
+```yaml
+octopus:
+  master_shard: <%= ENV['SHARD'] || 'shard1' %>
+```
+
+Then you can use the `SHARD` environment variable to override the `master_shard` specified in `config/shards.yml`, useful for running rake tasks.
+
+```bash
+SHARD=shard1 rake db:setup && SHARD=shard2 rake db:setup
+```
+
 ### Rails Controllers
 
 If you want to send a specified action, or all actions from a controller, to a specific shard,  use this syntax:
@@ -145,8 +172,8 @@ class ApplicationController < ActionController::Base
 end
 ```
 
- To see the complete list of features and syntax, please check out our <a href="http://wiki.github.com/tchandy/octopus/"> Wiki</a>
-Want to see sample rails applications using octopus features? please check it out: <a href="http://github.com/tchandy/octopus_sharding_example">Sharding Example</a> and <a href="http://github.com/tchandy/octopus_replication_example">Replication Example</a>. Also, we have an example that shows how to use Octopus without Rails: <a href="http://github.com/tchandy/octopus_sinatra"> Octopus + Sinatra Example</a>.
+ To see the complete list of features and syntax, please check out our <a href="http://wiki.github.com/thiagopradi/octopus/"> Wiki</a>
+Want to see sample rails applications using octopus features? please check it out: <a href="http://github.com/thiagopradi/octopus_sharding_example">Sharding Example</a> and <a href="http://github.com/thiagopradi/octopus_replication_example">Replication Example</a>. Also, we have an example that shows how to use Octopus without Rails: <a href="http://github.com/thiagopradi/octopus_sinatra"> Octopus + Sinatra Example</a>.
 
 ## Mixing Octopus with the Rails multiple database model
 If you want to set a custom connection to a specific model, use the syntax `octopus_establish_connection` syntax:
@@ -174,13 +201,13 @@ CustomConnectedModel.using(:my_shard).first
 CustomConnectedModel.using(:some_other_shard).first
 ```
 
-This can be useful if you have a model that lives in a separate database and would like to add sharding or replication to it. For other use cases, you may be better off with <a href="https://github.com/tchandy/octopus/wiki/Slave-Groups">slave groups</a>.
+This can be useful if you have a model that lives in a separate database and would like to add sharding or replication to it. For other use cases, you may be better off with <a href="https://github.com/thiagopradi/octopus/wiki/Slave-Groups">slave groups</a>.
 
 ## Contributing with Octopus
 Contributors are welcome! To run the test suite, you need mysql, postgresql and sqlite3 installed. This is what you need to setup your Octopus development environment:
 
 ```bash
-git clone http://github.com/tchandy/octopus.git
+git clone http://github.com/thiagopradi/octopus.git
 cd octopus
 bundle install
 bundle exec rake db:prepare
@@ -200,7 +227,7 @@ cucumber
 If you are having issues running the octopus spec suite, verify your database users and passwords match those inside the config files and your permissions are correct.
 
 ## Contributors:
-- <a href="https://github.com/tchandy/octopus/contributors">All Contributors</a>
+- <a href="https://github.com/thiagopradi/octopus/contributors">All Contributors</a>
 
 ## Mailing List:
 - <a href="http://groups.google.com/group/octopus-activerecord/">Octopus Mailing List</a>

data/ar-octopus.gemspec

@@ -25,7 +25,7 @@ Gem::Specification.new do |s|
   s.add_dependency 'activesupport', '>= 3.2.0'
 
   s.add_development_dependency 'appraisal', '>= 0.3.8'
-  s.add_development_dependency 'mysql2', '> 0.3'
+  s.add_development_dependency 'mysql2', '~> 0.3.18'
   s.add_development_dependency 'pg', '>= 0.11.0'
   s.add_development_dependency 'rake', '>= 0.8.7'
   s.add_development_dependency 'rspec', '>= 3'

data/lib/octopus.rb

@@ -11,12 +11,12 @@ module Octopus
   end
 
   def self.rails_env
-    @rails_env ||= self.rails? ? Rails.env.to_s : 'shards'
+    @rails_env ||= defined?(::Rails.env) ? Rails.env.to_s : 'shards'
   end
 
   def self.config
     @config ||= begin
-      file_name = Octopus.directory + '/config/shards.yml'
+      file_name = File.join(Octopus.directory, 'config/shards.yml').to_s
 
       if File.exist?(file_name) || File.symlink?(file_name)
         config ||= HashWithIndifferentAccess.new(YAML.load(ERB.new(File.read(file_name)).result))[Octopus.env]
@@ -28,13 +28,25 @@ module Octopus
     end
   end
 
+  def self.load_balancer=(balancer)
+    @load_balancer = balancer
+  end
+
+  def self.load_balancer
+    @load_balancer ||= Octopus::LoadBalancing::RoundRobin
+  end
+
+  def self.master_shard
+    ((config && config[:master_shard]) || :master).to_sym
+  end
+
   # Public: Whether or not Octopus is configured and should hook into the
   # current environment. Checks the environments config option for the Rails
   # environment by default.
   #
   # Returns a boolean
   def self.enabled?
-    if defined?(::Rails)
+    if defined?(::Rails.env)
       Octopus.environments.include?(Rails.env.to_s)
     else
       # TODO: This doens't feel right but !Octopus.config.blank? is breaking a
@@ -46,7 +58,7 @@ module Octopus
   # Returns the Rails.root_to_s when you are using rails
   # Running the current directory in a generic Ruby process
   def self.directory
-    @directory ||= defined?(Rails) ?  Rails.root.to_s : Dir.pwd
+    @directory ||= defined?(::Rails.root) ? Rails.root.to_s : Dir.pwd
   end
 
   # This is the default way to do Octopus Setup
@@ -90,14 +102,10 @@ module Octopus
     rails4? && ActiveRecord::VERSION::MINOR >= 1
   end
 
-  def self.rails?
-    defined?(Rails)
-  end
-
   attr_writer :logger
 
   def self.logger
-    if defined?(Rails)
+    if defined?(Rails.logger)
       @logger ||= Rails.logger
     else
       @logger ||= Logger.new($stderr)
@@ -119,15 +127,37 @@ module Octopus
     end
   end
 
+  def self.using_group(group, &block)
+    conn = ActiveRecord::Base.connection
+
+    if conn.is_a?(Octopus::Proxy)
+      conn.send_queries_to_group(group, &block)
+    else
+      yield
+    end
+  end
+
+  def self.using_all(&block)
+    conn = ActiveRecord::Base.connection
+
+    if conn.is_a?(Octopus::Proxy)
+      conn.send_queries_to_all_shards(&block)
+    else
+      yield
+    end
+  end
+
   def self.fully_replicated(&_block)
-    old_fully_replicated = Thread.current['octopus.fully_replicated']
-    Thread.current['octopus.fully_replicated'] = true
+    old_fully_replicated = Thread.current[Octopus::Proxy::FULLY_REPLICATED_KEY]
+    Thread.current[Octopus::Proxy::FULLY_REPLICATED_KEY] = true
     yield
   ensure
-    Thread.current['octopus.fully_replicated'] = old_fully_replicated
+    Thread.current[Octopus::Proxy::FULLY_REPLICATED_KEY] = old_fully_replicated
   end
 end
 
+require 'octopus/exception'
+
 require 'octopus/shard_tracking'
 require 'octopus/shard_tracking/attribute'
 require 'octopus/shard_tracking/dynamic'
@@ -143,7 +173,7 @@ require 'octopus/log_subscriber'
 require 'octopus/abstract_adapter'
 require 'octopus/singular_association'
 
-require 'octopus/railtie' if defined?(::Rails)
+require 'octopus/railtie' if defined?(::Rails::Railtie)
 
 require 'octopus/proxy'
 require 'octopus/collection_proxy'

data/lib/octopus/abstract_adapter.rb

@@ -20,20 +20,16 @@ module Octopus
         end
       end
 
-      def self.included(base)
-        base.alias_method_chain :initialize, :octopus_shard
-      end
-
       def octopus_shard
         @config[:octopus_shard]
       end
 
-      def initialize_with_octopus_shard(*args)
-        initialize_without_octopus_shard(*args)
+      def initialize(*args)
+        super
         @instrumenter = InstrumenterDecorator.new(self, @instrumenter)
       end
     end
   end
 end
 
-ActiveRecord::ConnectionAdapters::AbstractAdapter.send(:include, Octopus::AbstractAdapter::OctopusShard)
+ActiveRecord::ConnectionAdapters::AbstractAdapter.send(:prepend, Octopus::AbstractAdapter::OctopusShard)

data/lib/octopus/association.rb

@@ -2,6 +2,7 @@ module Octopus
   module Association
     def self.included(base)
       base.send(:include, Octopus::ShardTracking::Dynamic)
+      base.sharded_methods :target_scope
     end
 
     def current_shard

data/lib/octopus/association_shard_tracking.rb

@@ -4,37 +4,6 @@ module Octopus
       base.send(:include, InstanceMethods)
     end
 
-    module QueryOnCurrentShard
-      METHODS = %w(
-        all
-        average
-        count
-        empty?
-        exists?
-        find
-        find_by_sql
-        first
-        last
-        maximum
-        minimum
-        pluck
-        scoping
-        size
-        sum
-        to_a
-      )
-
-      METHODS.each do |m|
-        define_method m.to_sym do |*args, &block|
-          if self.respond_to?(:proxy_association) && proxy_association
-            proxy_association.owner.run_on_shard { super(*args, &block) }
-          else
-            super(*args, &block)
-          end
-        end
-      end
-    end
-
     module InstanceMethods
       def connection_on_association=(record)
         return unless ::Octopus.enabled?
@@ -56,14 +25,7 @@ module Octopus
         end
         super
       end
-    else
-      def has_many(association_id, options = {}, &extension)
-        default_octopus_opts(options)
-        super
-      end
-    end
 
-    if Octopus.rails4?
       def has_and_belongs_to_many(association_id, scope = nil, options = {}, &extension)
         if options == {} && scope.is_a?(Hash)
           default_octopus_opts(scope)
@@ -72,7 +34,13 @@ module Octopus
         end
         super
       end
+
     else
+      def has_many(association_id, options = {}, &extension)
+        default_octopus_opts(options)
+        super
+      end
+
       def has_and_belongs_to_many(association_id, options = {}, &extension)
         default_octopus_opts(options)
         super
@@ -80,23 +48,8 @@ module Octopus
     end
 
     def default_octopus_opts(options)
-      if options[:before_add].is_a?(Array)
-        options[:before_add] << :connection_on_association=
-      elsif options[:before_add].is_a?(Symbol)
-        options[:before_add] = [:connection_on_association=, options[:before_add]]
-      else
-        options[:before_add] = :connection_on_association=
-      end
-
-      if options[:before_remove].is_a?(Array)
-        options[:before_remove] << :connection_on_association=
-      elsif options[:before_remove].is_a?(Symbol)
-        options[:before_remove] = [:connection_on_association=, options[:before_remove]]
-      else
-        options[:before_remove] = :connection_on_association=
-      end
-
-      options[:extend] = [Octopus::AssociationShardTracking::QueryOnCurrentShard, options[:extend]].flatten.compact
+      options[:before_add] = [ :connection_on_association=, options[:before_add] ].compact.flatten
+      options[:before_remove] = [ :connection_on_association=, options[:before_remove] ].compact.flatten
     end
   end
 end

data/lib/octopus/collection_proxy.rb

@@ -5,6 +5,9 @@ module Octopus
       base.sharded_methods :any?, :build, :count, :create, :create!, :concat, :delete, :delete_all,
                            :destroy, :destroy_all, :empty?, :find, :first, :include?, :last, :length,
                            :many?, :pluck, :replace, :select, :size, :sum, :to_a, :uniq
+      if Octopus.rails3?
+        base.sharded_methods :scoped
+      end
     end
 
     def current_shard

data/lib/octopus/exception.rb

@@ -0,0 +1,4 @@
+module Octopus
+  class Exception < ::Exception
+  end
+end

data/lib/octopus/load_balancing/round_robin.rb

@@ -11,7 +11,7 @@ module Octopus
       end
 
       # Returns the next available slave in the pool
-      def next
+      def next(options)
         @slaves_list[@slave_index = (@slave_index + 1) % @slaves_list.length]
       end
     end

data/lib/octopus/migration.rb

@@ -51,7 +51,7 @@ module Octopus
           shards.merge(Array.wrap(shard))
         end
 
-        shards.to_a.presence || [:master]
+        shards.to_a.presence || [Octopus.master_shard]
       end
     end
   end
@@ -107,7 +107,7 @@ module Octopus
 
       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
+        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
           up_without_octopus(migrations_paths, target_version, &block)
@@ -116,7 +116,7 @@ module Octopus
 
       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
+        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
           down_without_octopus(migrations_paths, target_version, &block)

data/lib/octopus/model.rb

@@ -21,6 +21,14 @@ module Octopus
       end
 
       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?
           clean_table_name
           Octopus::ScopeProxy.new(shard, self)
@@ -57,7 +65,7 @@ module Octopus
       end
 
       def equality_with_octopus(comparison_object)
-        equality_without_octopus(comparison_object) && comparison_object.current_shard == current_shard
+        equality_without_octopus(comparison_object) && comparison_object.current_shard.to_s == current_shard.to_s
       end
 
       def perform_validations_with_octopus(*args)
@@ -95,11 +103,12 @@ module Octopus
       end
 
       def hijack_methods
-        around_save :run_on_shard, :unless => -> { self.class.custom_octopus_connection }
+        around_save :run_on_shard, :unless => lambda { self.class.custom_octopus_connection }
         after_initialize :set_current_shard
 
+        class_attribute :custom_octopus_connection
+
         class << self
-          attr_accessor :custom_octopus_connection
           attr_accessor :custom_octopus_table_name
 
           alias_method_chain :connection, :octopus

data/lib/octopus/proxy.rb

@@ -6,6 +6,15 @@ module Octopus
   class Proxy
     attr_accessor :config, :sharded
 
+    CURRENT_MODEL_KEY = 'octopus.current_model'.freeze
+    CURRENT_SHARD_KEY = 'octopus.current_shard'.freeze
+    CURRENT_GROUP_KEY = 'octopus.current_group'.freeze
+    CURRENT_SLAVE_GROUP_KEY = 'octopus.current_slave_group'.freeze
+    CURRENT_LOAD_BALANCE_OPTIONS_KEY = 'octopus.current_load_balance_options'.freeze
+    BLOCK_KEY = 'octopus.block'.freeze
+    LAST_CURRENT_SHARD_KEY = 'octopus.last_current_shard'.freeze
+    FULLY_REPLICATED_KEY = 'octopus.fully_replicated'.freeze
+
     def initialize(config = Octopus.config)
       initialize_shards(config)
       initialize_replication(config) if !config.nil? && config['replicated']
@@ -73,7 +82,7 @@ module Octopus
         end
       end
 
-      @shards[:master] ||= ActiveRecord::Base.connection_pool_without_octopus
+      @shards[:master] ||= ActiveRecord::Base.connection_pool_without_octopus if Octopus.master_shard == :master
     end
 
     def initialize_replication(config)
@@ -86,29 +95,30 @@ module Octopus
 
       @slaves_list = @shards.keys.map(&:to_s).sort
       @slaves_list.delete('master')
-      @slaves_load_balancer = Octopus::LoadBalancing::RoundRobin.new(@slaves_list)
+      @slaves_load_balancer = Octopus.load_balancer.new(@slaves_list)
     end
 
     def current_model
-      Thread.current['octopus.current_model']
+      Thread.current[CURRENT_MODEL_KEY]
     end
 
     def current_model=(model)
-      Thread.current['octopus.current_model'] = model.is_a?(ActiveRecord::Base) ? model.class : model
+      Thread.current[CURRENT_MODEL_KEY] = model.is_a?(ActiveRecord::Base) ? model.class : model
     end
 
     def current_shard
-      Thread.current['octopus.current_shard'] ||= :master
+      Thread.current[CURRENT_SHARD_KEY] ||= Octopus.master_shard
     end
 
     def current_shard=(shard_symbol)
-      self.current_slave_group = nil
       if shard_symbol.is_a?(Array)
+        self.current_slave_group = nil
         shard_symbol.each { |symbol| fail "Nonexistent Shard Name: #{symbol}" if @shards[symbol].nil? }
       elsif shard_symbol.is_a?(Hash)
         hash = shard_symbol
         shard_symbol = hash[:shard]
         slave_group_symbol = hash[:slave_group]
+        load_balance_options = hash[:load_balance_options]
 
         if shard_symbol.nil? && slave_group_symbol.nil?
           fail 'Neither shard or slave group must be specified'
@@ -123,17 +133,18 @@ module Octopus
               (@shards_slave_groups.try(:[], shard_symbol).nil? && @slave_groups[slave_group_symbol].nil?)
             fail "Nonexistent Slave Group Name: #{slave_group_symbol} in shards config: #{@shards_config.inspect}"
           end
-          self.current_slave_group = slave_group_symbol
         end
+        self.current_slave_group = slave_group_symbol
+        self.current_load_balance_options = load_balance_options
       else
         fail "Nonexistent Shard Name: #{shard_symbol}" if @shards[shard_symbol].nil?
       end
 
-      Thread.current['octopus.current_shard'] = shard_symbol
+      Thread.current[CURRENT_SHARD_KEY] = shard_symbol
     end
 
     def current_group
-      Thread.current['octopus.current_group']
+      Thread.current[CURRENT_GROUP_KEY]
     end
 
     def current_group=(group_symbol)
@@ -142,35 +153,44 @@ module Octopus
         fail "Nonexistent Group Name: #{group}" unless has_group?(group)
       end
 
-      Thread.current['octopus.current_group'] = group_symbol
+      Thread.current[CURRENT_GROUP_KEY] = group_symbol
     end
 
     def current_slave_group
-      Thread.current['octopus.current_slave_group']
+      Thread.current[CURRENT_SLAVE_GROUP_KEY]
     end
 
     def current_slave_group=(slave_group_symbol)
-      Thread.current['octopus.current_slave_group'] = slave_group_symbol
+      Thread.current[CURRENT_SLAVE_GROUP_KEY] = slave_group_symbol
+      Thread.current[CURRENT_LOAD_BALANCE_OPTIONS_KEY] = nil if slave_group_symbol.nil?
+    end
+
+    def current_load_balance_options
+      Thread.current[CURRENT_LOAD_BALANCE_OPTIONS_KEY]
+    end
+
+    def current_load_balance_options=(options)
+      Thread.current[CURRENT_LOAD_BALANCE_OPTIONS_KEY] = options
     end
 
     def block
-      Thread.current['octopus.block']
+      Thread.current[BLOCK_KEY]
     end
 
     def block=(block)
-      Thread.current['octopus.block'] = block
+      Thread.current[BLOCK_KEY] = block
     end
 
     def last_current_shard
-      Thread.current['octopus.last_current_shard']
+      Thread.current[LAST_CURRENT_SHARD_KEY]
     end
 
     def last_current_shard=(last_current_shard)
-      Thread.current['octopus.last_current_shard'] = last_current_shard
+      Thread.current[LAST_CURRENT_SHARD_KEY] = last_current_shard
     end
 
     def fully_replicated?
-      @fully_replicated || Thread.current['octopus.fully_replicated']
+      @fully_replicated || Thread.current[FULLY_REPLICATED_KEY]
     end
 
     # Public: Whether or not a group exists with the given name converted to a
@@ -202,7 +222,7 @@ module Octopus
     # reconnect, but in Rails 3.1 the flag prevents this.
     def safe_connection(connection_pool)
       connection_pool.automatic_reconnect ||= true
-      if !connection_pool.connected? && @shards[:master].connection.query_cache_enabled
+      if !connection_pool.connected? && @shards[Octopus.master_shard].connection.query_cache_enabled
         connection_pool.connection.enable_query_cache!
       end
       connection_pool.connection
@@ -229,13 +249,23 @@ module Octopus
     end
 
     def send_queries_to_multiple_shards(shards, &block)
-      shards.each do |shard|
+      shards.map do |shard|
         run_queries_on_shard(shard, &block)
       end
     end
 
+    def send_queries_to_group(group, &block)
+      using_group(group) do
+        send_queries_to_multiple_shards(shards_for_group(group), &block)
+      end
+    end
+
+    def send_queries_to_all_shards(&block)
+      send_queries_to_multiple_shards(shard_names.uniq { |shard_name| @shards[shard_name] }, &block)
+    end
+
     def clean_connection_proxy
-      self.current_shard = :master
+      self.current_shard = Octopus.master_shard
       self.current_model = nil
       self.current_group = nil
       self.block = nil
@@ -249,7 +279,7 @@ module Octopus
 
     def transaction(options = {}, &block)
       if !sharded && current_model_replicated?
-        run_queries_on_shard(:master) do
+        run_queries_on_shard(Octopus.master_shard) do
           select_connection.transaction(options, &block)
         end
       else
@@ -411,9 +441,9 @@ module Octopus
 
     def send_queries_to_selected_slave(method, *args, &block)
       if current_model.replicated || fully_replicated?
-        selected_slave = @slaves_load_balancer.next
+        selected_slave = @slaves_load_balancer.next current_load_balance_options
       else
-        selected_slave = :master
+        selected_slave = Octopus.master_shard
       end
 
       send_queries_to_slave(selected_slave, method, *args, &block)
@@ -439,7 +469,7 @@ module Octopus
     # Temporarily switch `current_shard` to the next slave in a slave group and send queries to it
     # while preserving `current_shard`
     def send_queries_to_balancer(balancer, method, *args, &block)
-      send_queries_to_slave(balancer.next, method, *args, &block)
+      send_queries_to_slave(balancer.next(current_load_balance_options), method, *args, &block)
     end
 
     # Temporarily switch `current_shard` to the specified slave and send queries to it
@@ -468,6 +498,9 @@ module Octopus
     # Temporarily switch `current_shard` and run the block
     def using_shard(shard, &_block)
       older_shard = current_shard
+      older_slave_group = current_slave_group
+      older_load_balance_options = current_load_balance_options
+
 
       begin
         unless current_model && !current_model.allowed_shard?(shard)
@@ -476,6 +509,20 @@ module Octopus
         yield
       ensure
         self.current_shard = older_shard
+        self.current_slave_group = older_slave_group
+        self.current_load_balance_options = older_load_balance_options
+      end
+    end
+
+    # Temporarily switch `current_group` and run the block
+    def using_group(group, &_block)
+      older_group = current_group
+
+      begin
+        self.current_group = group
+        yield
+      ensure
+        self.current_group = older_group
       end
     end
 

data/lib/octopus/railtie.rb

@@ -1,6 +1,4 @@
 begin
-  require 'rails/railtie'
-
   module Octopus
     class Railtie < Rails::Railtie
       rake_tasks do

data/lib/octopus/relation_proxy.rb

@@ -17,7 +17,11 @@ module Octopus
     end
 
     def method_missing(method, *args, &block)
-      run_on_shard { @ar_relation.public_send(method, *args, &block) }
+      if block
+        @ar_relation.public_send(method, *args, &block)
+      else
+        run_on_shard { @ar_relation.public_send(method, *args) }
+      end
     end
 
     def ==(other)

data/lib/octopus/slave_group.rb

@@ -3,11 +3,11 @@ module Octopus
     def initialize(slaves)
       slaves = HashWithIndifferentAccess.new(slaves)
       slaves_list = slaves.values
-      @load_balancer = Octopus::LoadBalancing::RoundRobin.new(slaves_list)
+      @load_balancer = Octopus.load_balancer.new(slaves_list)
     end
 
-    def next
-      @load_balancer.next
+    def next(options)
+      @load_balancer.next options
     end
   end
 end

data/lib/octopus/version.rb

@@ -1,3 +1,3 @@
 module Octopus
-  VERSION = '0.8.5'
+  VERSION = '0.8.6'
 end

data/spec/octopus/association_shard_tracking_spec.rb

@@ -630,6 +630,12 @@ describe Octopus::AssociationShardTracking, :shards => [:brazil, :master, :canad
         expect(@brazil_client.comments.count).to eq(2)
       end
 
+      it 'group + count' do
+        expect(@brazil_client.comments.group(:id).count.length).to eq(1)
+        _cmt = @brazil_client.comments.create(:name => 'Builded Comment')
+        expect(@brazil_client.comments.group(:id).count.length).to eq(2)
+      end
+
       it 'size' do
         expect(@brazil_client.comments.size).to eq(1)
         _cmt = @brazil_client.comments.create(:name => 'Builded Comment')

data/spec/octopus/model_spec.rb

@@ -2,6 +2,10 @@ require 'spec_helper'
 
 describe Octopus::Model do
   describe '#using method' do
+    it 'raise when Model#using receives a block' do
+      expect { User.using(:master) { true } }.to raise_error(Octopus::Exception, /User\.using is not allowed to receive a block/)
+    end
+
     it 'should allow to send a block to the master shard' do
       Octopus.using(:master) do
         User.create!(:name => 'Block test')
@@ -16,6 +20,16 @@ describe Octopus::Model do
       expect(User.using('canada').find_by_name('Rafael Pilha')).not_to be_nil
     end
 
+    it 'should allow comparison of a string shard name with symbol shard name' do
+      u = User.using('canada').create!(:name => 'Rafael Pilha')
+      expect(u).to eq(User.using(:canada).find_by_name('Rafael Pilha'))
+    end
+
+    it 'should allow comparison of a symbol shard name with string shard name' do
+      u = User.using(:canada).create!(:name => 'Rafael Pilha')
+      expect(u).to eq(User.using('canada').find_by_name('Rafael Pilha'))
+    end
+
     it 'should allow to pass a string as the shard name to a block' do
       Octopus.using('canada') do
         User.create!(:name => 'Rafael Pilha')
@@ -81,6 +95,21 @@ describe Octopus::Model do
       expect(ActiveRecord::Base.connection.current_shard).to eq(:master)
     end
 
+    it 'should ensure that the connection will be cleaned with custom master' do
+      OctopusHelper.using_environment :octopus do
+        Octopus.config[:master_shard] = :brazil
+        expect(ActiveRecord::Base.connection.current_shard).to eq(:brazil)
+        expect do
+          Octopus.using(:canada) do
+            fail 'Some Exception'
+          end
+        end.to raise_error
+
+        expect(ActiveRecord::Base.connection.current_shard).to eq(:brazil)
+        Octopus.config[:master_shard] = nil
+      end
+    end
+
     it 'should allow creating more than one user' do
       User.using(:canada).create([{ :name => 'America User 1' }, { :name => 'America User 2' }])
       User.create!(:name => 'Thiago')
@@ -99,6 +128,15 @@ describe Octopus::Model do
       expect(User.connection.current_shard).to eq(:master)
     end
 
+    it 'should clean #current_shard from proxy when using execute' do
+      OctopusHelper.using_environment :octopus do
+        Octopus.config[:master_shard] = :brazil
+        User.using(:canada).connection.execute('select * from users limit 1;')
+        expect(User.connection.current_shard).to eq(:brazil)
+        Octopus.config[:master_shard] = nil
+      end
+    end
+
     it 'should allow scoping dynamically' do
       User.using(:canada).using(:master).using(:canada).create!(:name => 'oi')
       expect(User.using(:canada).using(:master).count).to eq(0)
@@ -263,6 +301,12 @@ describe Octopus::Model do
       expect(CustomConnection.connection.current_database).to eq('octopus_shard_2')
     end
 
+    it 'reuses parent model connection' do
+      klass = Class.new(CustomConnection)
+
+      expect(klass.connection).to be klass.connection
+    end
+
     it 'should not mess with custom connection table names' do
       expect(Advert.connection.current_database).to eq('octopus_shard_1')
       Advert.create!(:name => 'Teste')

data/spec/octopus/proxy_spec.rb

@@ -211,6 +211,14 @@ describe Octopus::Proxy do
         expect(proxy.shard_name).to eq(:master)
       end
 
+      it 'when current_shard is empty with custom master' do
+        OctopusHelper.using_environment :octopus do
+          Octopus.config[:master_shard] = :brazil
+          expect(proxy.shard_name).to eq(:brazil)
+          Octopus.config[:master_shard] = nil
+        end
+      end
+
       it 'when current_shard is a single shard' do
         proxy.current_shard = :canada
         expect(proxy.shard_name).to eq(:canada)

data/spec/octopus/relation_proxy_spec.rb

@@ -68,14 +68,23 @@ describe Octopus::RelationProxy do
       it 'uses the correct shard' do
         expect(Item.using(:brazil).count).to eq(0)
         _clients_on_brazil = Client.using(:brazil).all
-        Client.using(:brazil) do
+        Octopus.using(:brazil) do
           expect(@relation.count).to eq(1)
         end
       end
 
+      it 'uses the correct shard in block when method_missing is triggered on CollectionProxy objects' do
+        Octopus.using(:brazil) do
+          @client.items.each do |item|
+            expect(item.current_shard).to eq(:canada)
+            expect(ActiveRecord::Base.connection.current_shard).to eq(:brazil)
+          end
+        end
+      end
+
       it 'lazily evaluates on the correct shard' do
         expect(Item.using(:brazil).count).to eq(0)
-        Client.using(:brazil) do
+        Octopus.using(:brazil) do
           expect(@relation.select(:client_id).count).to eq(1)
         end
       end

data/spec/octopus/replicated_slave_grouped_spec.rb

@@ -61,4 +61,31 @@ describe 'when the database is replicated and has slave groups' do
       expect(Cat.count).to eq(2)
     end
   end
+
+  it 'should keep sending to slaves in a using block' do
+    OctopusHelper.using_environment :replicated_slave_grouped do
+      Cat.create!(:name => 'Thiago1')
+      Cat.create!(:name => 'Thiago2')
+
+      expect(Cat.count).to eq(2)
+      Octopus.using(:slave_group => :slaves1) do
+        expect(Cat.count).to eq(0)
+        expect(Cat.count).to eq(0)
+      end
+    end
+  end
+  
+  it 'should restore previous slave group after a using block' do
+    OctopusHelper.using_environment :replicated_slave_grouped do
+      Cat.create!(:name => 'Thiago1')
+      Cat.create!(:name => 'Thiago2')
+
+      Octopus.using(:slave_group => :slaves1) do
+        Octopus.using(:slave_group => :slaves2) do
+          expect(Cat.count).to eq(2)
+        end
+        expect(Cat.count).to eq(0)
+      end
+    end
+  end
 end

data/spec/octopus/replication_spec.rb

@@ -123,4 +123,15 @@ describe 'when the database is replicated and the entire application is replicat
       expect(Cat.connection.current_shard).to eql(:master)
     end
   end
+
+  it 'should reset current shard if slave throws an exception with custom master' do
+    OctopusHelper.using_environment :production_fully_replicated do
+      Octopus.config[:master_shard] = :slave2
+      Cat.create!(:name => 'Slave Cat')
+      expect(Cat.connection.current_shard).to eql(:slave2)
+      Cat.where(:rubbish => true)
+      expect(Cat.connection.current_shard).to eql(:slave2)
+      Octopus.config[:master_shard] = nil
+    end
+  end
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: ar-octopus
 version: !ruby/object:Gem::Version
-  version: 0.8.5
+  version: 0.8.6
 platform: ruby
 authors:
 - Thiago Pradi
@@ -10,7 +10,7 @@ authors:
 autorequire: 
 bindir: bin
 cert_chain: []
-date: 2015-03-12 00:00:00.000000000 Z
+date: 2016-03-05 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: activerecord
@@ -58,16 +58,16 @@ dependencies:
   name: mysql2
   requirement: !ruby/object:Gem::Requirement
     requirements:
-    - - ">"
+    - - "~>"
       - !ruby/object:Gem::Version
-        version: '0.3'
+        version: 0.3.18
   type: :development
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
-    - - ">"
+    - - "~>"
       - !ruby/object:Gem::Version
-        version: '0.3'
+        version: 0.3.18
 - !ruby/object:Gem::Dependency
   name: pg
   requirement: !ruby/object:Gem::Requirement
@@ -187,6 +187,7 @@ files:
 - lib/octopus/association_shard_tracking.rb
 - lib/octopus/collection_association.rb
 - lib/octopus/collection_proxy.rb
+- lib/octopus/exception.rb
 - lib/octopus/has_and_belongs_to_many_association.rb
 - lib/octopus/load_balancing.rb
 - lib/octopus/load_balancing/round_robin.rb
@@ -342,7 +343,7 @@ required_rubygems_version: !ruby/object:Gem::Requirement
       version: '0'
 requirements: []
 rubyforge_project: 
-rubygems_version: 2.4.4
+rubygems_version: 2.4.3
 signing_key: 
 specification_version: 4
 summary: Easy Database Sharding for ActiveRecord