ar-octopus 0.0.2 → 0.0.3

data/.gitignore

@@ -1 +1,2 @@
 pkg/*.*
+tmp/*

data/Rakefile

@@ -1,6 +1,22 @@
 $LOAD_PATH.unshift(File.join(File.dirname(__FILE__), 'lib'))
 require 'rubygems'
 require 'rake'
+require "yaml"
+require "active_support"
+require 'active_support/json'
+begin
+  require 'metric_fu'
+  MetricFu::Configuration.run do |config|
+    config.metrics  = [:churn,:flay, :flog, :reek, :roodi, :saikuro]
+    config.graphs   = [:flog, :flay, :reek, :roodi]
+    config.flay     = { :dirs_to_flay => ['spec', 'lib']  }
+    config.flog     = { :dirs_to_flog => ['spec', 'lib']  }
+    config.reek     = { :dirs_to_reek => ['spec', 'lib']  }
+    config.roodi    = { :dirs_to_roodi => ['spec', 'lib'] }
+    config.churn    = { :start_date => "1 year ago", :minimum_churn_count => 10}
+  end
+rescue LoadError
+end
 
 begin
   require 'jeweler'
@@ -12,7 +28,7 @@ begin
     gem.homepage = "http://github.com/tchandy/octopus"
     gem.authors = ["Thiago Pradi", "Mike Perham", "Amit Agarwal"]
     gem.add_development_dependency "rspec", ">= 1.2.9"
-    gem.version = "0.0.2"
+    gem.version = "0.0.3"
   end
   Jeweler::GemcutterTasks.new
 rescue LoadError
@@ -76,6 +92,8 @@ namespace :db do
     [:master, :brazil, :canada, :russia, :alone_shard, :postgresql_shard].each do |shard_symbol|
       ActiveRecord::Base.using(shard_symbol).connection.create_table(:users) do |u|
         u.string :name
+        u.integer :number
+        u.boolean :admin
       end
       
       ActiveRecord::Base.using(shard_symbol).connection.create_table(:clients) do |u|
@@ -95,6 +113,41 @@ namespace :db do
       ActiveRecord::Base.using(shard_symbol).connection.create_table(:schema_migrations) do |u|
         u.string :version, :unique => true, :null => false
       end
+      
+      ActiveRecord::Base.using(shard_symbol).connection.create_table(:computers) do |u|
+        u.string :name
+      end
+      
+      ActiveRecord::Base.using(shard_symbol).connection.create_table(:keyboards) do |u|
+        u.string :name
+        u.integer :computer_id
+      end
+      
+      ActiveRecord::Base.using(shard_symbol).connection.create_table(:roles) do |u|
+        u.string :name
+      end
+      
+      ActiveRecord::Base.using(shard_symbol).connection.create_table(:permissions) do |u|
+        u.string :name
+      end
+      
+      ActiveRecord::Base.using(shard_symbol).connection.create_table(:permissions_roles, :id => false) do |u|
+        u.integer :role_id
+        u.integer :permission_id
+      end
+      
+      ActiveRecord::Base.using(shard_symbol).connection.create_table(:assignments) do |u|
+        u.integer :programmer_id
+        u.integer :project_id
+      end
+      
+      ActiveRecord::Base.using(shard_symbol).connection.create_table(:programmers) do |u|
+        u.string :name
+      end
+      
+      ActiveRecord::Base.using(shard_symbol).connection.create_table(:projects) do |u|
+        u.string :name
+      end
     end
   end
   

data/ar-octopus.gemspec

@@ -5,11 +5,11 @@
 
 Gem::Specification.new do |s|
   s.name = %q{ar-octopus}
-  s.version = "0.0.2"
+  s.version = "0.0.3"
 
   s.required_rubygems_version = Gem::Requirement.new(">= 0") if s.respond_to? :required_rubygems_version=
   s.authors = ["Thiago Pradi", "Mike Perham", "Amit Agarwal"]
-  s.date = %q{2010-06-14}
+  s.date = %q{2010-06-22}
   s.description = %q{This gem allows you to use sharded databases with ActiveRecord. this also provides a interface for replication and for running migrations with multiples shards.}
   s.email = %q{tchandy@gmail.com}
   s.extra_rdoc_files = [
@@ -26,7 +26,10 @@ Gem::Specification.new do |s|
      "doc/libraries.textile",
      "doc/shards.yml",
      "lib/octopus.rb",
+     "lib/octopus/association.rb",
+     "lib/octopus/association_collection.rb",
      "lib/octopus/controller.rb",
+     "lib/octopus/has_and_belongs_to_many_association.rb",
      "lib/octopus/migration.rb",
      "lib/octopus/model.rb",
      "lib/octopus/proxy.rb",
@@ -44,6 +47,7 @@ Gem::Specification.new do |s|
      "spec/migrations/7_raise_exception_with_invalid_multiple_shard_names.rb",
      "spec/migrations/8_raise_exception_with_invalid_group_name.rb",
      "spec/migrations/9_raise_exception_with_multiple_invalid_group_names.rb",
+     "spec/octopus/association_spec.rb",
      "spec/octopus/controller_spec.rb",
      "spec/octopus/migration_spec.rb",
      "spec/octopus/model_spec.rb",
@@ -55,7 +59,7 @@ Gem::Specification.new do |s|
   s.homepage = %q{http://github.com/tchandy/octopus}
   s.rdoc_options = ["--charset=UTF-8"]
   s.require_paths = ["lib"]
-  s.rubygems_version = %q{1.3.6}
+  s.rubygems_version = %q{1.3.7}
   s.summary = %q{Easy Database Sharding for ActiveRecord}
   s.test_files = [
     "spec/database_connection.rb",
@@ -71,6 +75,7 @@ Gem::Specification.new do |s|
      "spec/migrations/7_raise_exception_with_invalid_multiple_shard_names.rb",
      "spec/migrations/8_raise_exception_with_invalid_group_name.rb",
      "spec/migrations/9_raise_exception_with_multiple_invalid_group_names.rb",
+     "spec/octopus/association_spec.rb",
      "spec/octopus/controller_spec.rb",
      "spec/octopus/migration_spec.rb",
      "spec/octopus/model_spec.rb",
@@ -83,7 +88,7 @@ Gem::Specification.new do |s|
     current_version = Gem::Specification::CURRENT_SPECIFICATION_VERSION
     s.specification_version = 3
 
-    if Gem::Version.new(Gem::RubyGemsVersion) >= Gem::Version.new('1.2.0') then
+    if Gem::Version.new(Gem::VERSION) >= Gem::Version.new('1.2.0') then
       s.add_development_dependency(%q<rspec>, [">= 1.2.9"])
     else
       s.add_dependency(%q<rspec>, [">= 1.2.9"])

data/lib/octopus.rb

@@ -2,17 +2,17 @@ require "yaml"
 
 module Octopus
   def self.env()
-    @@env ||= defined?(Rails) ? Rails.env.to_s : "production"      
+    @env ||= defined?(Rails) ? Rails.env.to_s : "production"      
   end
   
   def self.config()
-    @@config ||= YAML.load_file(Octopus.directory() + "/config/shards.yml") 
+    @config ||= YAML.load_file(Octopus.directory() + "/config/shards.yml") 
   end
 
   # 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) ?  Rails.root.to_s : Dir.pwd     
   end
 end
 
@@ -20,3 +20,6 @@ require "octopus/proxy"
 require "octopus/migration"
 require "octopus/model"
 require "octopus/controller"
+require "octopus/association"
+require "octopus/association_collection"
+require "octopus/has_and_belongs_to_many_association"

data/lib/octopus/association.rb

@@ -0,0 +1,109 @@
+module Octopus::Association
+  def collection_reader_method(reflection, association_proxy_class)
+    define_method(reflection.name) do |*params|
+      force_reload = params.first unless params.empty?
+      reload_connection() 
+
+      association = association_instance_get(reflection.name)
+
+      unless association
+        association  = association_proxy_class.new(self, reflection)
+        association_instance_set(reflection.name, association)
+      end
+
+      reflection.klass.uncached { association.reload } if force_reload
+
+      association
+    end
+
+    define_method("#{reflection.name.to_s.singularize}_ids") do
+      reload_connection()        
+      if send(reflection.name).loaded? || reflection.options[:finder_sql]
+        send(reflection.name).map(&:id)
+      else
+        if reflection.through_reflection && reflection.source_reflection.belongs_to?
+          through = reflection.through_reflection
+          primary_key = reflection.source_reflection.primary_key_name
+          send(through.name).select("DISTINCT #{through.quoted_table_name}.#{primary_key}").map!(&:"#{primary_key}")
+        else
+          send(reflection.name).select("#{reflection.quoted_table_name}.#{reflection.klass.primary_key}").except(:includes).map!(&:id)
+        end
+      end
+    end
+  end
+
+  def association_constructor_method(constructor, reflection, association_proxy_class)
+    define_method("#{constructor}_#{reflection.name}") do |*params|
+      reload_connection() 
+      attributees      = params.first unless params.empty?
+      replace_existing = params[1].nil? ? true : params[1]
+      association      = association_instance_get(reflection.name)
+
+      unless association
+        association = association_proxy_class.new(self, reflection)
+        association_instance_set(reflection.name, association)
+      end
+
+      if association_proxy_class == ActiveRecord::Associations::HasOneAssociation
+        return_val = association.send(constructor, attributees, replace_existing)
+      else
+        return_val = association.send(constructor, attributees)
+      end
+
+      if have_a_valid_shard?
+        return_val.current_shard = self.current_shard
+      end
+
+      return_val
+    end
+  end
+
+  def association_accessor_methods(reflection, association_proxy_class)
+    define_method(reflection.name) do |*params|
+      force_reload = true
+      reload_connection() 
+      association = association_instance_get(reflection.name)
+
+      if association.nil? || force_reload
+        association = association_proxy_class.new(self, reflection)
+        retval = have_a_valid_shard? ? reflection.klass.uncached { self.class.connection_proxy.run_query_on_shard(self.current_shard) { association.reload } } : association.reload
+        if retval.nil? and association_proxy_class == ActiveRecord::Associations::BelongsToAssociation
+          association_instance_set(reflection.name, nil)
+          return nil
+        end
+        association_instance_set(reflection.name, association)
+      end
+
+      association
+      #association.target.nil? ? nil : association
+    end
+
+    define_method("loaded_#{reflection.name}?") do
+      reload_connection() 
+      association = association_instance_get(reflection.name)
+      association && association.loaded?
+    end
+
+    define_method("#{reflection.name}=") do |new_value|
+      reload_connection() 
+      association = association_instance_get(reflection.name)
+
+      if association.nil? || association.target != new_value
+        association = association_proxy_class.new(self, reflection)
+      end
+
+      association.replace(new_value)
+      association_instance_set(reflection.name, new_value.nil? ? nil : association)
+    end
+
+    define_method("set_#{reflection.name}_target") do |target|
+      return if target.nil? and association_proxy_class == ActiveRecord::Associations::BelongsToAssociation
+      reload_connection() 
+      association = association_proxy_class.new(self, reflection)
+      association.target = target
+      association_instance_set(reflection.name, association)
+    end
+  end
+end
+
+ActiveRecord::Base.extend(Octopus::Association)

data/lib/octopus/association_collection.rb

@@ -0,0 +1,40 @@
+class ActiveRecord::Associations::AssociationCollection
+  def should_wrap_the_connection?
+    @owner.respond_to?(:current_shard) && @owner.current_shard != nil
+  end
+
+  def count(column_name = nil, options = {})
+    if @reflection.options[:counter_sql]
+      if should_wrap_the_connection?
+        @owner.using(@owner.current_shard) { @reflection.klass.count_by_sql(@counter_sql) } 
+      else        
+        @reflection.klass.count_by_sql(@counter_sql)
+      end
+    else
+      column_name, options = nil, column_name if column_name.is_a?(Hash)
+
+      if @reflection.options[:uniq]
+        # This is needed because 'SELECT count(DISTINCT *)..' is not valid SQL.
+        column_name = "#{@reflection.quoted_table_name}.#{@reflection.klass.primary_key}" unless column_name
+        options.merge!(:distinct => true)
+      end
+
+      value = @reflection.klass.send(:with_scope, construct_scope) do 
+        if should_wrap_the_connection?
+          @owner.using(@owner.current_shard) { @reflection.klass.count(column_name, options) } 
+        else        
+          @reflection.klass.count(column_name, options) 
+        end
+      end
+
+      limit  = @reflection.options[:limit]
+      offset = @reflection.options[:offset]
+
+      if limit || offset
+        [ [value - offset.to_i, 0].max, limit.to_i ].min
+      else
+        value
+      end
+    end
+  end
+end

data/lib/octopus/controller.rb

@@ -1,14 +1,6 @@
 module Octopus::Controller
   def using(shard, &block)
-    older_shard = ActiveRecord::Base.connection_proxy.current_shard
-    ActiveRecord::Base.connection_proxy.block = true
-    ActiveRecord::Base.connection_proxy.current_shard = shard
-    begin
-      yield
-    ensure
-      ActiveRecord::Base.connection_proxy.block = false
-      ActiveRecord::Base.connection_proxy.current_shard = older_shard
-    end
+    ActiveRecord::Base.connection_proxy.run_query_on_shard(shard, &block)
   end
 end
 

data/lib/octopus/has_and_belongs_to_many_association.rb

@@ -0,0 +1,43 @@
+class ActiveRecord::Associations::HasAndBelongsToManyAssociation
+  def should_wrap_the_connection?
+    @owner.respond_to?(:current_shard) && @owner.current_shard != nil
+  end
+  
+  def insert_record(record, force = true, validate = true)
+    if record.new_record?
+      if force
+        record.save!
+      else
+        return false unless record.save(:validate => validate)
+      end
+    end
+
+    if @reflection.options[:insert_sql]
+      @owner.connection.insert(interpolate_sql(@reflection.options[:insert_sql], record))
+    else
+      relation = Arel::Table.new(@reflection.options[:join_table])
+      attributes = columns.inject({}) do |attrs, column|
+        case column.name.to_s
+        when @reflection.primary_key_name.to_s
+          attrs[relation[column.name]] = owner_quoted_id
+        when @reflection.association_foreign_key.to_s
+          attrs[relation[column.name]] = record.quoted_id
+        else
+          if record.has_attribute?(column.name)
+            value = @owner.send(:quote_value, record[column.name], column)
+            attrs[relation[column.name]] = value unless value.nil?
+          end
+        end
+        attrs
+      end
+      
+      if should_wrap_the_connection?
+        @owner.using(@owner.current_shard) { relation.insert(attributes) } 
+      else
+        relation.insert(attributes)
+      end
+    end
+
+    return true
+  end
+end

data/lib/octopus/migration.rb

@@ -1,41 +1,34 @@
 module Octopus::Migration
-  def self.included(base)
-    base.extend(ClassMethods)
+  def self.extended(base)
     class << base
       def connection
         ActiveRecord::Base.connection_proxy()
       end
-
-      def connected?
-        ActiveRecord::Base.connection_proxy().connected?
-      end
     end
   end
 
-  module ClassMethods
-    def using(*args)
-      if args.size == 1
-        self.connection().block = true
-        self.connection().current_shard = args.first
-      else
-        self.connection().current_shard = args        
-      end
-
-      return self
+  def using(*args)
+    if args.size == 1
+      self.connection().block = true
+      self.connection().current_shard = args.first
+    else
+      self.connection().current_shard = args        
     end
 
-    def using_group(*args)
-      if args.size == 1
-        self.connection().block = true
-        self.connection().current_group = args.first
-      else
-        self.connection().current_group = args
-      end
-      
-      return self
+    return self
+  end
+
+  def using_group(*args)
+    if args.size == 1
+      self.connection().block = true
+      self.connection().current_group = args.first
+    else
+      self.connection().current_group = args
     end
+
+    return self
   end
 end
 
 
-ActiveRecord::Migration.send(:include, Octopus::Migration)
+ActiveRecord::Migration.extend(Octopus::Migration)

data/lib/octopus/model.rb

@@ -1,18 +1,54 @@
 module Octopus::Model  
-  def self.included(base)
-    base.extend(ClassMethods)    
+  def self.extended(base) 
     base.send(:include, InstanceMethods)
   end
 
   module InstanceMethods
-    def connection_proxy
-      self.class.connection_proxy
+    def reload_connection()
+      set_connection() if have_a_valid_shard?
     end
-    
-    def using(shard, &block)
+
+    def update_attribute(name, value)
+      reload_connection()
+      super(name, value)
+    end
+
+    def update_attributes(attributes)
+      reload_connection()
+      super(attributes)
+    end
+
+    def update_attributes!(attributes)
+      reload_connection()
+      super(attributes)
+    end
+
+    def reload
+      reload_connection()
+      super
+    end
+
+    def hijack_initializer()
+      attr_accessor :current_shard
+      after_initialize :set_current_shard
+      before_save :set_connection
+      before_destroy :set_connection
+
+      def set_current_shard
+        if self.class.respond_to?(:connection_proxy) && self.respond_to?(:current_shard)
+          if self.new_record?
+            self.current_shard = self.class.connection_proxy.current_shard    
+          else
+            self.current_shard = self.class.connection_proxy.last_current_shard  
+          end
+        end
+      end    
+    end
+
+    def hijack_connection()
       class << self
         def connection_proxy
-          @@connection_proxy ||= Octopus::Proxy.new(Octopus.config())
+          Thread.current[:connection_proxy] ||= Octopus::Proxy.new(Octopus.config())
         end
 
         def connection 
@@ -22,39 +58,71 @@ module Octopus::Model
 
           self.connection_proxy()
         end
-
-        def connected?
-          self.connection_proxy().connected?
-        end
       end
-      
-      self.reset_table_name() if self != ActiveRecord::Base && self.respond_to?(:reset_table_name)
-      
+    end
+
+    def using(shard, &block)
+      hijack_connection()  
+      clean_table_name()
+
       if block_given?
-        older_shard = self.connection_proxy.current_shard
-        self.connection_proxy.block = true
-        self.connection_proxy.current_shard = shard
-        begin
-          yield
-        ensure
-          self.connection_proxy.block = false
-          self.connection_proxy.current_shard = older_shard
-        end
+        self.connection_proxy.run_query_on_shard(shard, &block)
       else
+        hijack_initializer()  
         self.connection_proxy.current_shard = shard
         self.connection_proxy.using_enabled = true
+
         return self
       end
     end
+
+    def have_a_valid_shard?
+      self.respond_to?(:current_shard) && self.current_shard != nil
+    end
+
+    def set_connection(*args)
+      if(args.size == 1)
+        arg = args.first
+        arg.current_shard = self.current_shard if arg.respond_to?(:current_shard) && have_a_valid_shard?
+      end
+
+      self.class.connection_proxy.current_shard = self.current_shard if have_a_valid_shard?
+    end
+
+    def clean_table_name
+      self.reset_table_name() if self != ActiveRecord::Base && self.respond_to?(:reset_table_name)
+    end
+  end
+
+  include InstanceMethods
+
+  def replicated_model()
+    self.cattr_accessor :replicated
+  end
+
+  def has_many(association_id, options = {}, &extension)
+    default_octopus_opts(options)
+    super(association_id, options, &extension)
   end
 
-  module ClassMethods
-    include InstanceMethods
+  def has_and_belongs_to_many(association_id, options = {}, &extension)
+    default_octopus_opts(options)
+    super(association_id, options, &extension)    
+  end
 
-    def replicated_model()
-      self.cattr_accessor :replicated
+  def default_octopus_opts(options)
+    if options[:before_add].is_a?(Array)
+      options[:before_add] << :set_connection
+    else
+      options[:before_add] = :set_connection
     end
-  end  
+
+    if options[:before_remove].is_a?(Array)
+      options[:before_remove] << :set_connection
+    else
+      options[:before_remove] = :set_connection
+    end
+  end
 end
 
-ActiveRecord::Base.send(:include, Octopus::Model)
+ActiveRecord::Base.extend(Octopus::Model)