dynashard 0.1.0

data/.document

@@ -0,0 +1,5 @@
+lib/**/*.rb
+bin/*
+- 
+features/**/*.feature
+LICENSE.txt

data/Gemfile

@@ -0,0 +1,16 @@
+source "http://rubygems.org"
+
+gem 'activerecord', '>= 3.0'
+
+group :development do
+  gem "shoulda", ">= 0"
+  gem "bundler", "~> 1.0.0"
+  gem "jeweler", "~> 1.5.2"
+  gem "rcov", ">= 0"
+end
+
+group :test do
+  gem 'rspec', '>= 2.0'
+  gem 'sqlite3-ruby', :require => 'sqlite3'
+  gem 'factory_girl_rails'
+end

data/Gemfile.lock

@@ -0,0 +1,100 @@
+GEM
+  remote: http://rubygems.org/
+  specs:
+    abstract (1.0.0)
+    actionmailer (3.0.1)
+      actionpack (= 3.0.1)
+      mail (~> 2.2.5)
+    actionpack (3.0.1)
+      activemodel (= 3.0.1)
+      activesupport (= 3.0.1)
+      builder (~> 2.1.2)
+      erubis (~> 2.6.6)
+      i18n (~> 0.4.1)
+      rack (~> 1.2.1)
+      rack-mount (~> 0.6.12)
+      rack-test (~> 0.5.4)
+      tzinfo (~> 0.3.23)
+    activemodel (3.0.1)
+      activesupport (= 3.0.1)
+      builder (~> 2.1.2)
+      i18n (~> 0.4.1)
+    activerecord (3.0.1)
+      activemodel (= 3.0.1)
+      activesupport (= 3.0.1)
+      arel (~> 1.0.0)
+      tzinfo (~> 0.3.23)
+    activeresource (3.0.1)
+      activemodel (= 3.0.1)
+      activesupport (= 3.0.1)
+    activesupport (3.0.1)
+    arel (1.0.1)
+      activesupport (~> 3.0.0)
+    builder (2.1.2)
+    diff-lcs (1.1.2)
+    erubis (2.6.6)
+      abstract (>= 1.0.0)
+    factory_girl (1.3.2)
+    factory_girl_rails (1.0)
+      factory_girl (~> 1.3)
+      rails (>= 3.0.0.beta4)
+    git (1.2.5)
+    i18n (0.4.2)
+    jeweler (1.5.2)
+      bundler (~> 1.0.0)
+      git (>= 1.2.5)
+      rake
+    mail (2.2.9)
+      activesupport (>= 2.3.6)
+      i18n (~> 0.4.1)
+      mime-types (~> 1.16)
+      treetop (~> 1.4.8)
+    mime-types (1.16)
+    polyglot (0.3.1)
+    rack (1.2.1)
+    rack-mount (0.6.13)
+      rack (>= 1.0.0)
+    rack-test (0.5.6)
+      rack (>= 1.0)
+    rails (3.0.1)
+      actionmailer (= 3.0.1)
+      actionpack (= 3.0.1)
+      activerecord (= 3.0.1)
+      activeresource (= 3.0.1)
+      activesupport (= 3.0.1)
+      bundler (~> 1.0.0)
+      railties (= 3.0.1)
+    railties (3.0.1)
+      actionpack (= 3.0.1)
+      activesupport (= 3.0.1)
+      rake (>= 0.8.4)
+      thor (~> 0.14.0)
+    rake (0.8.7)
+    rcov (0.9.9)
+    rspec (2.1.0)
+      rspec-core (~> 2.1.0)
+      rspec-expectations (~> 2.1.0)
+      rspec-mocks (~> 2.1.0)
+    rspec-core (2.1.0)
+    rspec-expectations (2.1.0)
+      diff-lcs (~> 1.1.2)
+    rspec-mocks (2.1.0)
+    shoulda (2.11.3)
+    sqlite3-ruby (1.3.2)
+    thor (0.14.4)
+    treetop (1.4.8)
+      polyglot (>= 0.3.1)
+    tzinfo (0.3.23)
+
+PLATFORMS
+  ruby
+
+DEPENDENCIES
+  activerecord (>= 3.0)
+  bundler (~> 1.0.0)
+  factory_girl_rails
+  jeweler (~> 1.5.2)
+  rcov
+  rspec (>= 2.0)
+  shoulda
+  sqlite3-ruby

data/LICENSE.txt

@@ -0,0 +1,20 @@
+Copyright (c) 2011 Nick Hengeveld
+
+Permission is hereby granted, free of charge, to any person obtaining
+a copy of this software and associated documentation files (the
+"Software"), to deal in the Software without restriction, including
+without limitation the rights to use, copy, modify, merge, publish,
+distribute, sublicense, and/or sell copies of the Software, and to
+permit persons to whom the Software is furnished to do so, subject to
+the following conditions:
+
+The above copyright notice and this permission notice shall be
+included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
+EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
+MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
+NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
+LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
+OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
+WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

data/README.md

@@ -0,0 +1,72 @@
+# Dynashard - Dynamic sharding for ActiveRecord
+
+This package provides database sharding functionality for ActiveRecord models.
+
+Sharding is disabled by default and is enabled with +Dynashard.enable+.  This allows
+sharding behavior to be enabled globally or only for specific environments so for
+example production environments could be sharded while development environments could
+use a single database.
+
+Models may be configured to determine the appropriate shard (database connection) to
+use based on context defined prior to performing queries.
+
+    class Widget < ActiveRecord::Base
+      shard :by => :user
+    end
+  
+    class WidgetController < ApplicationController
+      around_filter :set_shard_context
+  
+      def index
+        # Widgets will be loaded using the connection for the current user's shard
+        @widgets = Widget.find(:all)
+      end
+  
+      private
+  
+        def set_shard_context
+          Dynashard.with_context(:user => current_user.shard) do
+            yield
+          end
+        end
+    end
+
+Associated models may be configured to use different shards determined by the
+association's owner.
+
+    class Company < ActiveRecord::Base
+      shard :associated, :using => :shard
+  
+      has_many :customers
+  
+      def shard
+        # logic to find the company's shard
+      end
+    end
+  
+    class Customer < ActiveRecord::Base
+      belongs_to :company
+      shard :by => :company
+    end
+  
+    > c = Company.find(:first)
+    => #<Company id:1>
+
+  Company is loaded using the default ActiveRecord connection.
+
+    > c.customers
+    => [#<Dynashard::Shard0::Customer id: 1>, #<Dynashard::Shard0::Customer id: 2>]
+
+  Customers are loaded using the connection for the Company's shard.  Associated models
+  are returned as shard-specific subclasses of the association class.
+
+    > c.customers.create(:name => 'Always right')
+    => #<Dynashard::Shard0::Customer id: 3>
+
+  New associations are saved on the Company's shard.
+
+## TODO: add gotcha section, eg:
+
+ - uniqueness validations should be scoped by whatever is sharding
+ - ways to shoot yourself in the foot with non-sharding association
+   owners of sharded models

data/Rakefile

@@ -0,0 +1,45 @@
+$LOAD_PATH.unshift File.expand_path("../lib", __FILE__)
+
+require 'rubygems'
+require 'bundler'
+require "rspec/core/rake_task"
+begin
+  Bundler.setup(:default, :development)
+rescue Bundler::BundlerError => e
+  $stderr.puts e.message
+  $stderr.puts "Run `bundle install` to install missing gems"
+  exit e.status_code
+end
+require 'rake'
+
+RSpec::Core::RakeTask.new(:spec)
+
+task :default => :spec
+
+require 'jeweler'
+Jeweler::Tasks.new do |gem|
+  # gem is a Gem::Specification... see http://docs.rubygems.org/read/chapter/20 for more options
+  gem.name = "dynashard"
+  gem.homepage = "http://github.com/nickh/dynashard"
+  gem.license = "MIT"
+  gem.summary = 'Dynamic sharding for ActiveRecord'
+  gem.description = 'Dynashard allows you to shard your ActiveRecord models.  Models can be configured to shard based on context that can be defined dynamically.'
+  gem.email = "nickh@verticalresponse.com"
+  gem.authors = ["Nick Hengeveld"]
+  # Include your dependencies below. Runtime dependencies are required when using your gem,
+  # and development dependencies are only needed for development (ie running rake tasks, tests, etc)
+  #  gem.add_runtime_dependency 'jabber4r', '> 0.1'
+  #  gem.add_development_dependency 'rspec', '> 1.2.3'
+  gem.add_runtime_dependency 'activerecord', '>= 3.0'
+end
+Jeweler::RubygemsDotOrgTasks.new
+
+require 'rake/rdoctask'
+Rake::RDocTask.new do |rdoc|
+  version = File.exist?('VERSION') ? File.read('VERSION') : ""
+
+  rdoc.rdoc_dir = 'rdoc'
+  rdoc.title = "dynashard #{version}"
+  rdoc.rdoc_files.include('README*')
+  rdoc.rdoc_files.include('lib/**/*.rb')
+end

data/VERSION

@@ -0,0 +1 @@
+0.1.0

data/lib/dynashard/arel_sql_engine.rb

@@ -0,0 +1,22 @@
+module Dynashard
+  module ArelEngineExtensions
+    def self.included(base)
+      base.alias_method_chain :connection, :dynashard
+    end
+    
+    def connection_with_dynashard
+      if @ar && @ar.respond_to?(:dynashard_klass)
+        @ar.dynashard_klass.connection
+      elsif @ar && @ar.sharding_enabled?
+        spec = Dynashard.shard_context[@ar.dynashard_context]
+        raise "Missing #{@ar.dynashard_context} shard context" if spec.nil?
+        spec = spec.call if spec.respond_to?(:call)
+        Dynashard.class_for(spec).connection
+      else
+        connection_without_dynashard
+      end
+    end
+  end
+end
+
+Arel::Sql::Engine.send(:include, Dynashard::ArelEngineExtensions)

data/lib/dynashard/associations.rb

@@ -0,0 +1,19 @@
+module Dynashard
+  module ProxyExtensions
+    def self.included(base)
+      base.alias_method_chain :initialize, :dynashard
+    end
+
+    # Initialize an association proxy.  If the proxy owner class is configured to
+    # shard its associations and the reflection klass is sharded, use a custom
+    # reflection with a sharded class.
+    def initialize_with_dynashard(owner, reflection)
+      if owner.class.shards_associated? && reflection.klass.sharding_enabled?
+        reflection = Dynashard.reflection_for(owner, reflection)
+      end
+      initialize_without_dynashard(owner, reflection)
+    end
+  end
+end
+
+ActiveRecord::Associations::AssociationProxy.send(:include, Dynashard::ProxyExtensions)

data/lib/dynashard/connection_handler.rb

@@ -0,0 +1,45 @@
+module Dynashard
+  module ConnectionHandlerExtensions
+    def self.included(base)
+      base.alias_method_chain :retrieve_connection_pool, :dynashard
+    end
+    
+    # Set an connection pool entry for the model class pointing at
+    # the shard class's pool.
+    # :nodoc:
+    # This is required for places that examine the connection pool to
+    # determine whether to use a given class's connection or a parent
+    # class's connection (eg. ActiveRecord::Base.arel_engine)
+    def dynashard_pool_alias(model_class, shard_class)
+      @connection_pools[model_class] = @connection_pools[shard_class]
+    end
+
+    # Return a connection pool for the specified class.  If sharding
+    # is enabled, the correct pool for the configured sharding context
+    # will be used.
+    #
+    # :nodoc:
+    # Reflection classes generated for an association proxy will
+    # respond to :dynashard_klass and return a class with an
+    # established connection to the correct shard.
+    #
+    # Sharded models will have a dynashard context defined, which
+    # can be used as a key to the Dynashard.shard_context hash to
+    # access the shard's connection specification.
+    def retrieve_connection_pool_with_dynashard(klass)
+      if klass.respond_to?(:dynashard_klass)
+        retrieve_connection_pool_without_dynashard(klass.dynashard_klass)
+      elsif klass.sharding_enabled?
+        spec = Dynashard.shard_context[klass.dynashard_context]
+        raise "Missing #{klass.dynashard_context} shard context" if spec.nil?
+        spec = spec.call if spec.respond_to?(:call)
+        shard_klass = Dynashard.class_for(spec)
+        retrieve_connection_pool_without_dynashard(shard_klass)
+      else
+        retrieve_connection_pool_without_dynashard(klass)
+      end
+    end
+  end
+end
+
+ActiveRecord::ConnectionAdapters::ConnectionHandler.send(:include, Dynashard::ConnectionHandlerExtensions)

data/lib/dynashard/model.rb

@@ -0,0 +1,92 @@
+module Dynashard
+  module ActiveRecordExtensions
+    def self.extended(base)
+      base.extend(ClassMethods)
+
+      # Change ActiveRecord::Base.arel_engine to create an engine for sharded models
+      # rather than using the ActiveRecord::Base class.
+      base.module_eval do
+        def self.arel_engine
+          if sharding_enabled?
+            Arel::Sql::Engine.new(self)
+          elsif self == ActiveRecord::Base
+            Arel::Table.engine
+          else
+            connection_handler.connection_pools[name] ? Arel::Sql::Engine.new(self) : superclass.arel_engine
+          end
+        end
+      end
+    end
+
+    module ClassMethods
+      # Configure sharding for the current model.  The following options are supported:
+      # * +:by => :context+ - the database connection for this model will be determined using a connection
+      #   proxy with the specified context.
+      # * +:associated, :using => :method+ - the database connection for associated models with sharding
+      #   enabled will be determined using the association's defined context after first setting that
+      #   context using the specified method on the current instance
+      #
+      #   class MyModel < ActiveRecord::Base
+      #     shard :by => :context
+      #     shard :assocated, :using => :sharding_method
+      #   end
+      #
+      #   def sharding_method
+      #     # return a valid shard descriptor:
+      #     # - a string key into the configurations in databases.yml
+      #     # - a hash with connection parameters
+      #     # - an object that responds to :call with one of the above
+      #   end
+      def shard(*args)
+        if args.first == :associated
+          using = args.last[:using] if args.last.respond_to?(:[])
+          raise ArgumentError.new(":associated specified without :using") if using.nil? 
+          @dynashard_association_using = using
+        else
+          shard_by = args.first[:by] if args.first.respond_to?(:[])
+          raise ArgumentError.new("Invalid options") if shard_by.nil?
+          @dynashard_context = shard_by
+        end
+      end
+
+      # Returns true if sharding has been globally enabled and has been configured for this model
+      def sharding_enabled?
+        Dynashard.enabled? && !@dynashard_context.nil?
+      end
+
+      # Returns true if sharding has been globally enabled and is configured to be used for
+      # sharded models associated with this model
+      def shards_associated?
+        Dynashard.enabled? && !@dynashard_association_using.nil?
+      end
+
+      # Returns true if the class was generated by Dynashard
+      def dynashard_model?
+        false
+      end
+
+      # Return a subclass configured to connect to the appropriate shard
+      def dynashard_sharded_subclass
+        if sharding_enabled?
+          spec = Dynashard.shard_context[dynashard_context]
+          raise "Missing #{dynashard_context} shard context" if spec.nil?
+          spec = spec.call if spec.respond_to?(:call)
+          shard_klass = Dynashard.class_for(spec)
+          Dynashard.sharded_model_class(shard_klass, self)
+        end
+      end
+
+      # Returns the shard context for this model
+      def dynashard_context
+        @dynashard_context
+      end
+
+      # Returns the method used to set the context used for associated models
+      def dynashard_association_using
+        @dynashard_association_using
+      end
+    end
+  end
+end
+
+ActiveRecord::Base.extend(Dynashard::ActiveRecordExtensions)

data/lib/dynashard/validations.rb

@@ -0,0 +1,23 @@
+module Dynashard
+  module ValidationExtensions
+    def self.included(base)
+      base.alias_method_chain :find_finder_class_for, :dynashard
+    end
+    
+    # Return the class that should be used to find other instances of
+    # the specified record's model class.  If the record is an instance
+    # of a sharded model, it should be used; otherwise the default
+    # behavior should be used.
+    def find_finder_class_for_with_dynashard(record)
+      if record.class.dynashard_model?
+        record.class
+      elsif record.class.sharding_enabled?
+        record.class.dynashard_sharded_subclass
+      else
+        find_finder_class_for_without_dynashard(record)
+      end
+    end
+  end
+end
+
+ActiveRecord::Validations::UniquenessValidator.send(:include, Dynashard::ValidationExtensions)

data/lib/dynashard.rb

@@ -0,0 +1,167 @@
+# = Dynashard - Dynamic sharding for ActiveRecord
+#
+# This package provides database sharding functionality for ActiveRecord models.
+#
+# Sharding is disabled by default and is enabled with +Dynashard.enable+.  This allows
+# sharding behavior to be enabled globally or only for specific environments so for
+# example production environments could be sharded while development environments could
+# use a single database.
+#
+# Models may be configured to determine the appropriate shard (database connection) to
+# use based on context defined prior to performing queries.
+#
+#   class Widget < ActiveRecord::Base
+#     shard :by => :user
+#   end
+#
+#   class WidgetController < ApplicationController
+#     around_filter :set_shard_context
+#
+#     def index
+#       # Widgets will be loaded using the connection for the current user's shard
+#       @widgets = Widget.find(:all)
+#     end
+#
+#     private
+#
+#       def set_shard_context
+#         Dynashard.with_context(:user => current_user.shard) do
+#           yield
+#         end
+#       end
+#   end
+#
+# Associated models may be configured to use different shards determined by the
+# association's owner.
+#
+#   class Company < ActiveRecord::Base
+#     shard :associated, :using => :shard
+#
+#     has_many :customers
+#
+#     def shard
+#       # logic to find the company's shard
+#     end
+#   end
+#
+#   class Customer < ActiveRecord::Base
+#     belongs_to :company
+#     shard :by => :company
+#   end
+#
+#   > c = Company.find(:first)
+#   => #<Company id:1>
+#
+#   Company is loaded using the default ActiveRecord connection.
+#
+#   > c.customers
+#   => [#<Dynashard::Shard0::Customer id: 1>, #<Dynashard::Shard0::Customer id: 2>]
+#
+#   Customers are loaded using the connection for the Company's shard.  Associated models
+#   are returned as shard-specific subclasses of the association class.
+#
+#   > c.customers.create(:name => 'Always right')
+#   => #<Dynashard::Shard0::Customer id: 3>
+#
+#   New associations are saved on the Company's shard.
+#
+# TODO: add gotcha section, eg:
+#  - uniqueness validations should be scoped by whatever is sharding
+
+module Dynashard
+  # Enable sharding for all models configured to do so
+  def self.enable
+    @enabled = true
+  end
+
+  # Disable sharding for all models configured to do so
+  def self.disable
+    @enabled = false
+  end
+
+  # Return true if sharding is globally enabled
+  def self.enabled?
+    @enabled == true
+  end
+
+  # Execute a block within a given sharding context
+  def self.with_context(new_context, &block)
+    orig_context = shard_context.dup
+    shard_context.merge! new_context
+    result = nil
+    begin
+      result = yield
+    ensure
+      shard_context.replace orig_context
+    end
+    result
+  end
+
+  # Return a threadsafe(?) current mapping of shard context to connection spec
+  def self.shard_context
+    Thread.current[:shard_context] ||= {}
+  end
+
+  # Return a class with an established connection to a database shard
+  def self.class_for(spec)
+    @class_cache ||= {}
+    @class_cache[spec] ||= new_shard_class(spec)
+    @class_cache[spec]
+  end
+
+  # Return a reflection with a sharded class
+  def self.reflection_for(owner, reflection)
+    reflection_copy = reflection.dup
+    shard_klass = Dynashard.class_for(owner.send(owner.class.dynashard_association_using))
+    klass = sharded_model_class(shard_klass, reflection.klass)
+    reflection_copy.instance_variable_set('@klass', klass)
+    reflection_copy.instance_variable_set('@class_name', klass.name)
+
+    reflection_copy
+  end
+
+  # Return a model subclass configured to use a specific shard
+  def self.sharded_model_class(shard_klass, base_klass)
+    class_name = "#{shard_klass.name}::#{base_klass.name}"
+    unless shard_klass.constants.include?(base_klass.name.to_sym)
+      class_eval <<EOE
+        class #{class_name} < #{base_klass.name}
+          @@dynashard_klass = #{shard_klass.name}
+
+          def self.dynashard_model?()
+            true
+          end
+
+          def self.dynashard_klass()
+            @@dynashard_klass
+          end
+
+          def self.connection
+            dynashard_klass.connection
+          end
+        end
+EOE
+      klass = class_name.constantize
+      klass.connection_handler.dynashard_pool_alias(klass.name, shard_klass.name)
+    end
+    class_name.constantize
+  end
+
+  private
+
+    def self.new_shard_class(spec)
+      shard_number = @class_cache.size
+      klass_name = "Shard#{shard_number}"
+      unless const_defined?(klass_name)
+        module_eval("class #{klass_name} < ActiveRecord::Base ; end")
+        "Dynashard::#{klass_name}".constantize.establish_connection(spec)
+      end
+      "Dynashard::#{klass_name}".constantize
+    end
+end
+
+require 'dynashard/model'
+require 'dynashard/associations'
+require 'dynashard/connection_handler'
+require 'dynashard/validations'
+require 'dynashard/arel_sql_engine'

data/spec/arel_sql_engine_spec.rb

@@ -0,0 +1,51 @@
+require 'spec_helper'
+
+describe 'Dynashard::ArelEngineExtensions' do
+  before(:each) do
+    @mock_record = mock()
+    @engine = Arel::Sql::Engine.new(@mock_record)
+  end
+
+  describe '#connection_with_dynashard' do
+    context 'with a sharding association owner' do
+      it "returns the generated model subclass connection" do
+        @mock_record.should_receive(:dynashard_klass).once.and_return(mock(:connection => :subclass_connection))
+        @engine.should_receive(:connection_without_dynashard).never
+        @engine.connection.should == :subclass_connection
+      end
+    end
+
+    context 'with a sharding model' do
+      before(:each) do
+        @mock_record.should_receive(:sharding_enabled?).and_return(true)
+        @mock_record.should_receive(:dynashard_context).at_least(:once).and_return(:dynashard_context)
+      end
+
+      context 'and no defined sharding context' do
+        it 'raises an exception' do
+          Dynashard.shard_context[:dynashard_context] = nil
+          lambda do
+            @engine.connection
+          end.should raise_error
+        end
+      end
+    
+      context 'and a defined sharding context' do
+        it 'returns the generated shard class connection' do
+          Dynashard.should_receive(:class_for).with(:shard_spec).and_return(mock(:connection => :shard_connection))
+          Dynashard.shard_context[:dynashard_context] = :shard_spec
+          @engine.should_receive(:connection_without_dynashard).never
+          @engine.connection.should == :shard_connection
+        end
+      end
+    end
+
+    context 'with a non-sharding model' do
+      it 'returns the default connection' do
+        @mock_record.should_receive(:sharding_enabled?).and_return(false)
+        @engine.should_receive(:connection_without_dynashard).and_return(:model_connection)
+        @engine.connection.should == :model_connection
+      end
+    end
+  end
+end