ar-octopus 0.3.4 → 0.4.0

data/lib/octopus.rb

@@ -11,10 +11,16 @@ module Octopus
   end
 
   def self.config()
-    @config ||= HashWithIndifferentAccess.new(YAML.load(ERB.new(File.open(Octopus.directory() + "/config/shards.yml").read()).result))[Octopus.env()]
+    file_name = Octopus.directory() + "/config/shards.yml"
 
-    if @config && @config['environments']
-      self.environments = @config['environments']
+    if File.exists? file_name
+      @config ||= HashWithIndifferentAccess.new(YAML.load(ERB.new(File.open(file_name).read()).result))[Octopus.env()]
+
+      if @config && @config['environments']
+        self.environments = @config['environments']
+      end
+    else
+      @config ||= HashWithIndifferentAccess.new
     end
 
     @config
@@ -49,6 +55,12 @@ module Octopus
     defined?(Rails) 
   end
   
+  def self.shards=(shards)
+    @config ||= HashWithIndifferentAccess.new
+    @config[rails_env()] = HashWithIndifferentAccess.new(shards)
+    ActiveRecord::Base.connection.initialize_shards(@config)
+  end
+  
   def self.using(shard, &block)
     ActiveRecord::Base.hijack_initializer()
     conn = ActiveRecord::Base.connection
@@ -71,6 +83,7 @@ require "octopus/association"
 if Octopus.rails3?
   require "octopus/rails3/association"
   require "octopus/rails3/persistence"
+  require "octopus/rails3/arel"
 else
   require "octopus/rails2/association"
   require "octopus/rails2/persistence"
@@ -78,4 +91,4 @@ end
 
 require "octopus/proxy"
 require "octopus/scope_proxy"
-
+require "octopus/logger"

data/lib/octopus/logger.rb

@@ -0,0 +1,14 @@
+require "logger"
+
+class Octopus::Logger < Logger
+  def format_message(severity, timestamp, progname, msg)
+    str = super 
+    
+    if ActiveRecord::Base.connection.respond_to?(:current_shard)
+      str += "Shard: #{ActiveRecord::Base.connection.current_shard} -" 
+    end
+    
+    str
+  end
+end
+

data/lib/octopus/model.rb

@@ -7,14 +7,22 @@ module Octopus::Model
 
   module SharedMethods
     def clean_table_name
-      self.reset_table_name() if self != ActiveRecord::Base && self.respond_to?(:reset_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() 
+      end
+
+      if Octopus.rails3?
+        self.reset_column_information
+        self.instance_variable_set(:@quoted_table_name, nil)
+      end
     end
 
     def using(shard)
       return self if defined?(::Rails) && !Octopus.environments.include?(Rails.env.to_s)
       
-      clean_table_name()
       hijack_initializer() if !respond_to?(:set_current_shard)
+      clean_table_name()
 
       self.connection_proxy.using_enabled = true
 
@@ -23,7 +31,6 @@ module Octopus::Model
 
     def hijack_initializer()
       attr_accessor :current_shard
-      after_initialize :set_current_shard
       before_save :reload_connection
 
       def set_current_shard
@@ -34,7 +41,9 @@ module Octopus::Model
         end
       end
 
-      if !Octopus.rails3?
+      if Octopus.rails3?
+        after_initialize :set_current_shard
+      else
         def after_initialize
           set_current_shard()
         end
@@ -90,7 +99,12 @@ module Octopus::Model
       write_inheritable_attribute(:establish_connection, true)      
       establish_connection(spec)
     end
+
+    def octopus_set_table_name(value = nil, &block)
+      write_inheritable_attribute(:set_table_name, true)      
+      set_table_name(value, &block)
+    end
   end
 end
 
-ActiveRecord::Base.extend(Octopus::Model)
+ActiveRecord::Base.extend(Octopus::Model)

data/lib/octopus/proxy.rb

@@ -1,3 +1,5 @@
+require "set"
+
 class Octopus::Proxy
   attr_accessor :current_model, :current_shard, :current_group, :block, :using_enabled, :last_current_shard, :config
 
@@ -7,8 +9,9 @@ class Octopus::Proxy
   end
 
   def initialize_shards(config)
-    @shards = {}
-    @groups = {}
+    @shards = HashWithIndifferentAccess.new
+    @groups = HashWithIndifferentAccess.new
+    @adapters = Set.new
     @shards[:master] = ActiveRecord::Base.connection_pool()
     @config = ActiveRecord::Base.connection_pool.connection.instance_variable_get(:@config)
     @current_shard = :master
@@ -87,6 +90,10 @@ class Octopus::Proxy
     current_shard.is_a?(Array) ? current_shard.first : current_shard
   end
   
+  def should_clean_table_name?
+    @adapters.size > 1
+  end
+  
   def run_queries_on_shard(shard, &block)
     older_shard = self.current_shard
     last_block = self.block
@@ -143,12 +150,17 @@ class Octopus::Proxy
     end
   end
 
+  def respond_to?(method, include_private = false)
+    super || select_connection.respond_to?(method, include_private)
+  end
+
   protected
   def connection_pool_for(adapter, config)
     ActiveRecord::ConnectionAdapters::ConnectionPool.new(ActiveRecord::Base::ConnectionSpecification.new(adapter, config))
   end
 
   def initialize_adapter(adapter)
+    @adapters << adapter
     begin
       require "active_record/connection_adapters/#{adapter}_adapter"
     rescue LoadError
@@ -167,16 +179,19 @@ class Octopus::Proxy
   def send_queries_to_selected_slave(method, *args, &block)        
     old_shard = self.current_shard
     
-    if current_model.read_inheritable_attribute(:replicated) || @fully_replicated
-      self.current_shard = @slaves_list.shift.to_sym
-      @slaves_list << self.current_shard
-    else
-      self.current_shard = :master
-    end
+    begin
+      if current_model.read_inheritable_attribute(:replicated) || @fully_replicated
+        self.current_shard = @slaves_list.shift.to_sym
+        @slaves_list << self.current_shard
+      else
+        self.current_shard = :master
+      end
     
-    sql = select_connection().send(method, *args, &block)     
-    self.current_shard = old_shard
-    @using_enabled = nil
-    return sql    
+      sql = select_connection().send(method, *args, &block)     
+      return sql    
+    ensure
+      self.current_shard = old_shard
+      @using_enabled = nil
+    end
   end
 end

data/lib/octopus/rails3/arel.rb

@@ -0,0 +1,13 @@
+class Arel::Visitors::ToSql
+  def quote value, column = nil
+    ActiveRecord::Base.connection.quote value, column
+  end
+  
+  def quote_table_name name
+    ActiveRecord::Base.connection.quote_table_name(name)
+  end
+
+  def quote_column_name name
+    Arel::Nodes::SqlLiteral === name ? name : ActiveRecord::Base.connection.quote_column_name(name)
+  end
+end

data/lib/octopus/scope_proxy.rb

@@ -24,11 +24,11 @@ class Octopus::ScopeProxy
     @klass.connection()
   end
   
-  def method_missing(method, *args, &block)
+  def method_missing(method, *args, &block)        
     @klass.connection.run_queries_on_shard(@shard) do
       @klass = @klass.send(method, *args, &block)
     end
-
+                
     return @klass if @klass.is_a?(ActiveRecord::Base) or @klass.is_a?(Array) or @klass.is_a?(Fixnum) or @klass.nil?
     return self
   end
@@ -37,4 +37,4 @@ class Octopus::ScopeProxy
     @shard == other.shard
     @klass == other.klass
   end
-end
+end

data/sample_app/Gemfile

@@ -1,6 +1,6 @@
 source :rubygems
 
-gem 'rails', '3.0.0'
+gem 'rails', '3.0.5'
 
 # Bundle edge Rails instead:
 # gem 'rails', :git => 'git://github.com/rails/rails.git'
@@ -18,4 +18,4 @@ group :test do
   gem "rspec-rails", ">= 2.0.0.beta.16"
   gem 'ruby-debug' if RUBY_VERSION < "1.9"
   gem "aruba"
-end
+end

data/sample_app/Gemfile.lock

@@ -1,130 +1,146 @@
 GIT
   remote: git://github.com/tchandy/octopus.git
-  revision: ff16c71
+  revision: 13320d08a1d722f1cd4d12468b328a59367a2c88
   specs:
-    ar-octopus (0.1.0)
+    ar-octopus (0.3.4)
       activerecord (>= 2.3)
 
 GEM
   remote: http://rubygems.org/
   specs:
     abstract (1.0.0)
-    actionmailer (3.0.0)
-      actionpack (= 3.0.0)
-      mail (~> 2.2.5)
-    actionpack (3.0.0)
-      activemodel (= 3.0.0)
-      activesupport (= 3.0.0)
+    actionmailer (3.0.5)
+      actionpack (= 3.0.5)
+      mail (~> 2.2.15)
+    actionpack (3.0.5)
+      activemodel (= 3.0.5)
+      activesupport (= 3.0.5)
       builder (~> 2.1.2)
       erubis (~> 2.6.6)
-      i18n (~> 0.4.1)
+      i18n (~> 0.4)
       rack (~> 1.2.1)
-      rack-mount (~> 0.6.12)
-      rack-test (~> 0.5.4)
+      rack-mount (~> 0.6.13)
+      rack-test (~> 0.5.7)
       tzinfo (~> 0.3.23)
-    activemodel (3.0.0)
-      activesupport (= 3.0.0)
+    activemodel (3.0.5)
+      activesupport (= 3.0.5)
       builder (~> 2.1.2)
-      i18n (~> 0.4.1)
-    activerecord (3.0.0)
-      activemodel (= 3.0.0)
-      activesupport (= 3.0.0)
-      arel (~> 1.0.0)
+      i18n (~> 0.4)
+    activerecord (3.0.5)
+      activemodel (= 3.0.5)
+      activesupport (= 3.0.5)
+      arel (~> 2.0.2)
       tzinfo (~> 0.3.23)
-    activeresource (3.0.0)
-      activemodel (= 3.0.0)
-      activesupport (= 3.0.0)
-    activesupport (3.0.0)
-    arel (1.0.1)
-      activesupport (~> 3.0.0)
-    aruba (0.2.1)
+    activeresource (3.0.5)
+      activemodel (= 3.0.5)
+      activesupport (= 3.0.5)
+    activesupport (3.0.5)
+    arel (2.0.9)
+    aruba (0.2.7)
+      background_process
+      cucumber (~> 0.10.0)
+    background_process (1.2)
     builder (2.1.2)
-    capybara (0.3.9)
+    capybara (0.4.1.2)
+      celerity (>= 0.7.9)
       culerity (>= 0.2.4)
       mime-types (>= 1.16)
       nokogiri (>= 1.3.3)
       rack (>= 1.0.0)
       rack-test (>= 0.5.4)
-      selenium-webdriver (>= 0.0.3)
-    columnize (0.3.1)
-    configuration (1.1.0)
-    cucumber (0.8.5)
-      builder (~> 2.1.2)
+      selenium-webdriver (>= 0.0.27)
+      xpath (~> 0.1.3)
+    celerity (0.8.8)
+    childprocess (0.1.7)
+      ffi (~> 0.6.3)
+    columnize (0.3.2)
+    configuration (1.2.0)
+    cucumber (0.10.0)
+      builder (>= 2.1.2)
       diff-lcs (~> 1.1.2)
-      gherkin (~> 2.1.4)
-      json_pure (~> 1.4.3)
-      term-ansicolor (~> 1.0.4)
+      gherkin (~> 2.3.2)
+      json (~> 1.4.6)
+      term-ansicolor (~> 1.0.5)
     cucumber-rails (0.3.2)
       cucumber (>= 0.8.0)
-    culerity (0.2.12)
-    database_cleaner (0.5.2)
+    culerity (0.2.15)
+    database_cleaner (0.6.4)
     diff-lcs (1.1.2)
     erubis (2.6.6)
       abstract (>= 1.0.0)
     ffi (0.6.3)
       rake (>= 0.8.7)
-    gherkin (2.1.5)
-      trollop (~> 1.16.2)
-    i18n (0.4.1)
-    json_pure (1.4.6)
+    gherkin (2.3.3)
+      json (~> 1.4.6)
+    i18n (0.5.0)
+    json (1.4.6)
+    json_pure (1.5.1)
     launchy (0.3.7)
       configuration (>= 0.0.5)
       rake (>= 0.8.1)
     linecache (0.43)
-    mail (2.2.5)
+    mail (2.2.15)
       activesupport (>= 2.3.6)
-      mime-types
-      treetop (>= 1.4.5)
+      i18n (>= 0.4.0)
+      mime-types (~> 1.16)
+      treetop (~> 1.4.8)
     mime-types (1.16)
-    nokogiri (1.4.3.1)
+    nokogiri (1.4.4)
     polyglot (0.3.1)
     rack (1.2.1)
     rack-mount (0.6.13)
       rack (>= 1.0.0)
-    rack-test (0.5.4)
+    rack-test (0.5.7)
       rack (>= 1.0)
-    rails (3.0.0)
-      actionmailer (= 3.0.0)
-      actionpack (= 3.0.0)
-      activerecord (= 3.0.0)
-      activeresource (= 3.0.0)
-      activesupport (= 3.0.0)
-      bundler (~> 1.0.0)
-      railties (= 3.0.0)
-    railties (3.0.0)
-      actionpack (= 3.0.0)
-      activesupport (= 3.0.0)
-      rake (>= 0.8.4)
-      thor (~> 0.14.0)
+    rails (3.0.5)
+      actionmailer (= 3.0.5)
+      actionpack (= 3.0.5)
+      activerecord (= 3.0.5)
+      activeresource (= 3.0.5)
+      activesupport (= 3.0.5)
+      bundler (~> 1.0)
+      railties (= 3.0.5)
+    railties (3.0.5)
+      actionpack (= 3.0.5)
+      activesupport (= 3.0.5)
+      rake (>= 0.8.7)
+      thor (~> 0.14.4)
     rake (0.8.7)
-    rspec (2.0.0.beta.20)
-      rspec-core (= 2.0.0.beta.20)
-      rspec-expectations (= 2.0.0.beta.20)
-      rspec-mocks (= 2.0.0.beta.20)
-    rspec-core (2.0.0.beta.20)
-    rspec-expectations (2.0.0.beta.20)
-      diff-lcs (>= 1.1.2)
-    rspec-mocks (2.0.0.beta.20)
-    rspec-rails (2.0.0.beta.20)
-      rspec (= 2.0.0.beta.20)
-    ruby-debug (0.10.3)
+    rspec (2.5.0)
+      rspec-core (~> 2.5.0)
+      rspec-expectations (~> 2.5.0)
+      rspec-mocks (~> 2.5.0)
+    rspec-core (2.5.1)
+    rspec-expectations (2.5.0)
+      diff-lcs (~> 1.1.2)
+    rspec-mocks (2.5.0)
+    rspec-rails (2.5.0)
+      actionpack (~> 3.0)
+      activesupport (~> 3.0)
+      railties (~> 3.0)
+      rspec (~> 2.5.0)
+    ruby-debug (0.10.4)
       columnize (>= 0.1)
-      ruby-debug-base (~> 0.10.3.0)
-    ruby-debug-base (0.10.3)
+      ruby-debug-base (~> 0.10.4.0)
+    ruby-debug-base (0.10.4)
       linecache (>= 0.3)
     rubyzip (0.9.4)
-    selenium-webdriver (0.0.28)
-      ffi (>= 0.6.1)
+    selenium-webdriver (0.1.3)
+      childprocess (~> 0.1.5)
+      ffi (~> 0.6.3)
       json_pure
       rubyzip
     spork (0.8.4)
-    sqlite3-ruby (1.3.1)
+    sqlite3 (1.3.3)
+    sqlite3-ruby (1.3.3)
+      sqlite3 (>= 1.3.3)
     term-ansicolor (1.0.5)
-    thor (0.14.0)
-    treetop (1.4.8)
+    thor (0.14.6)
+    treetop (1.4.9)
       polyglot (>= 0.3.1)
-    trollop (1.16.2)
-    tzinfo (0.3.23)
+    tzinfo (0.3.24)
+    xpath (0.1.3)
+      nokogiri (~> 1.3)
 
 PLATFORMS
   ruby
@@ -137,7 +153,7 @@ DEPENDENCIES
   cucumber-rails
   database_cleaner
   launchy
-  rails (= 3.0.0)
+  rails (= 3.0.5)
   rspec-rails (>= 2.0.0.beta.16)
   ruby-debug
   spork