data_fabric 1.0.0

data/init.rb

@@ -0,0 +1 @@
+DataFabric.init

data/lib/data_fabric/version.rb

@@ -0,0 +1,5 @@
+module DataFabric
+	module Version
+		STRING = "1.0.0"
+	end
+end

data/lib/data_fabric.rb

@@ -0,0 +1,231 @@
+require 'active_record'
+require 'active_record/version'
+
+# DataFabric adds a new level of flexibility to ActiveRecord connection handling.
+# You need to describe the topology for your database infrastructure in your model(s).  As with ActiveRecord normally, different models can use different topologies.
+# 
+# class MyHugeVolumeOfDataModel < ActiveRecord::Base
+#   connection_topology :replicated => true, :shard_by => :city
+# end
+# 
+# There are four supported modes of operation, depending on the options given to the connection_topology method.  The plugin will look for connections in your config/database.yml with the following convention:
+# 
+# No connection topology:
+# #{environment} - this is the default, as with ActiveRecord, e.g. "production"
+# 
+# connection_topology :replicated => true
+# #{environment}_#{role} - no sharding, just replication, where role is "master" or "slave", e.g. "production_master"
+# 
+# connection_topology :shard_by => :city
+# #{group}_#{shard}_#{environment} - sharding, no replication, e.g. "city_austin_production"
+# 
+# connection_topology :replicated => true, :shard_by => :city
+# #{group}_#{shard}_#{environment}_#{role} - sharding with replication, e.g. "city_austin_production_master"
+# 
+# 
+# When marked as replicated, all write and transactional operations for the model go to the master, whereas read operations go to the slave.
+# 
+# Since sharding is an application-level concern, your application must set the shard to use based on the current request or environment.  The current shard for a group is set on a thread local variable.  For example, you can set the shard in an ActionController around_filter based on the user as follows:
+# 
+# class ApplicationController < ActionController::Base
+#   around_filter :select_shard
+#   
+#   private
+#   def select_shard(&action_block)
+#     DataFabric.activate_shard(:city => @current_user.city, &action_block)
+#   end
+# end
+module DataFabric
+  VERSION = "1.0.0"
+  
+  def self.logger
+    ActiveRecord::Base.logger
+  end
+
+  def self.init
+    logger.info "Loading data_fabric #{VERSION} with ActiveRecord #{ActiveRecord::VERSION::STRING}"
+    ActiveRecord::Base.send(:include, self)
+  end
+  
+  def self.activate_shard(shards, &block)
+    ensure_setup
+    shards.each do |key, value|
+      Thread.current[:shards][key.to_s] = value.to_s
+    end
+    if block_given?
+      begin
+        yield
+      ensure
+        shards.each do |key, value|
+          Thread.current[:shards].delete(key.to_s)
+        end
+      end
+    end
+  end
+  
+  # For cases where you can't pass a block to activate_shards, you can
+  # clean up the thread local settings by calling this method at the
+  # end of processing
+  def self.deactivate_shard(shards)
+    ensure_setup
+    shards.each do |key, value|
+      Thread.current[:shards].delete(key.to_s)
+    end
+  end
+  
+  def self.active_shard(group)
+    raise ArgumentError, 'No shard has been activated' unless Thread.current[:shards]
+
+    returning(Thread.current[:shards][group.to_s]) do |shard|
+      raise ArgumentError, "No active shard for #{group}" unless shard
+    end
+  end
+  
+  def self.included(model)
+    # Wire up ActiveRecord::Base
+    model.extend ClassMethods
+  end
+
+  def self.ensure_setup
+    Thread.current[:shards] = {} unless Thread.current[:shards]
+  end
+  
+  # Class methods injected into ActiveRecord::Base
+  module ClassMethods
+    def connection_topology(options)
+      proxy = DataFabric::ConnectionProxy.new(self, options)
+      ActiveRecord::Base.active_connections[name] = proxy
+      
+      raise ArgumentError, "data_fabric does not support ActiveRecord's allow_concurrency = true" if allow_concurrency
+      DataFabric.logger.info "Creating data_fabric proxy for class #{name}"
+    end
+  end
+  
+  class StringProxy
+    def initialize(&block)
+      @proc = block
+    end
+    def to_s
+      @proc.call
+    end
+  end
+
+  class ConnectionProxy
+    def initialize(model_class, options)
+      @model_class = model_class      
+      @replicated = options[:replicated]
+      @shard_group = options[:shard_by]
+      @prefix = options[:prefix]
+      @current_role = 'slave' if @replicated
+      @current_connection_name_builder = connection_name_builder
+      @cached_connection = nil
+      @current_connection_name = nil
+      @role_changed = false
+
+      @model_class.send :include, ActiveRecordConnectionMethods if @replicated
+    end
+
+    delegate :insert, :update, :delete, :create_table, :rename_table, :drop_table, :add_column, :remove_column, 
+      :change_column, :change_column_default, :rename_column, :add_index, :remove_index, :initialize_schema_information,
+      :dump_schema_information, :to => :master
+    
+    def transaction(start_db_transaction = true, &block)
+      with_master { raw_connection.transaction(start_db_transaction, &block) }
+    end
+
+    def method_missing(method, *args, &block)
+      unless @cached_connection and !@role_changed
+        raw_connection
+        @role_changed = false
+      end
+      if logger.debug?
+        logger.debug("Calling #{method} on #{@cached_connection}")
+      end
+      @cached_connection.send(method, *args, &block) 
+    end
+    
+    def connection_name
+      @current_connection_name_builder.join('_')
+    end
+    
+    def disconnect!
+      @cached_connection.disconnect! if @cached_connection
+      @cached_connection = nil
+    end
+    
+    def verify!(arg)
+      @cached_connection.verify!(0) if @cached_connection
+    end
+    
+    def with_master
+      set_role('master')
+      yield
+    ensure
+      set_role('slave')
+    end
+
+    private
+    
+    def connection_name_builder
+      clauses = []
+      clauses << @prefix if @prefix
+      clauses << @shard_group if @shard_group
+      clauses << StringProxy.new { DataFabric.active_shard(@shard_group) } if @shard_group
+      clauses << RAILS_ENV
+      clauses << StringProxy.new { @current_role } if @replicated
+      clauses
+    end
+    
+    def raw_connection
+      conn_name = connection_name
+      unless already_connected_to? conn_name 
+        @cached_connection = begin 
+          config = ActiveRecord::Base.configurations[conn_name]
+          raise ArgumentError, "Unknown database config: #{conn_name}, have #{ActiveRecord::Base.configurations.inspect}" unless config
+          @model_class.establish_connection config
+          if logger.debug?
+            logger.debug "Switching from #{@current_connection_name} to #{conn_name}"
+          end
+          @current_connection_name = conn_name
+          conn = @model_class.connection
+          conn.verify! 0
+          conn
+        end
+        @model_class.active_connections[@model_class.name] = self
+      end
+      @cached_connection
+    end
+    
+    def already_connected_to?(conn_name)
+      conn_name == @current_connection_name and @cached_connection
+    end
+    
+    def set_role(role)
+      if @replicated and @current_role != role
+        @current_role = role
+        @role_changed = true
+      end
+    end
+    
+    def master
+      set_role('master')
+      return raw_connection
+    ensure
+      set_role('slave')
+    end
+    
+    def logger
+      DataFabric.logger
+    end
+  end
+
+  module ActiveRecordConnectionMethods
+    def self.included(base)
+      base.alias_method_chain :reload, :master
+    end
+
+    def reload_with_master(*args, &block)
+      connection.with_master { reload_without_master }
+    end
+  end
+end

data/test/connection_test.rb

@@ -0,0 +1,103 @@
+require File.join(File.dirname(__FILE__), 'test_helper')
+require 'flexmock/test_unit'
+
+class PrefixModel < ActiveRecord::Base
+  connection_topology :prefix => 'prefix'
+end
+
+class ShardModel < ActiveRecord::Base
+  connection_topology :shard_by => :city
+end
+
+class TheWholeEnchilada < ActiveRecord::Base
+  connection_topology :prefix => 'fiveruns', :replicated => true, :shard_by => :city
+end
+
+class AdapterMock < ActiveRecord::ConnectionAdapters::AbstractAdapter
+  # Minimum required to perform a find with no results
+   def columns(table_name, name=nil)
+     []
+   end
+   def select(sql, name=nil)
+     []
+   end
+   def execute(sql, name=nil)
+     0
+   end
+   
+   def name
+     'fake-db'
+   end
+  
+   def method_missing(name, *args)
+     raise ArgumentError, "#{self.class.name} missing '#{name}': #{args.inspect}"
+   end
+end
+
+class RawConnection
+  def method_missing(name, *args)
+      puts "#{self.class.name} missing '#{name}': #{args.inspect}"
+  end
+end
+
+class ConnectionTest < Test::Unit::TestCase
+
+  def test_should_install_into_arbase
+    assert PrefixModel.methods.include?('connection_topology')
+  end
+  
+  def test_prefix_connection_name
+    setup_configuration_for PrefixModel, 'prefix_test'
+    assert_equal 'prefix_test', PrefixModel.connection.connection_name
+  end
+  
+  def test_shard_connection_name
+    setup_configuration_for ShardModel, 'city_austin_test'
+    # ensure unset means error
+    assert_raises ArgumentError do
+      ShardModel.connection.connection_name
+    end
+    DataFabric.activate_shard(:city => 'austin', :category => 'art') do
+      assert_equal 'city_austin_test', ShardModel.connection.connection_name
+    end
+    # ensure it got unset
+    assert_raises ArgumentError do
+      ShardModel.connection.connection_name
+    end
+  end
+  
+  def test_enchilada
+    setup_configuration_for TheWholeEnchilada, 'fiveruns_city_dallas_test_slave'
+    setup_configuration_for TheWholeEnchilada, 'fiveruns_city_dallas_test_master'
+    DataFabric.activate_shard :city => :dallas do
+      assert_equal 'fiveruns_city_dallas_test_slave', TheWholeEnchilada.connection.connection_name
+
+      # Should use the slave
+      assert_raises ActiveRecord::RecordNotFound do
+        TheWholeEnchilada.find(1)
+      end
+      
+      # Should use the master
+      mmmm = TheWholeEnchilada.new
+      mmmm.instance_variable_set(:@attributes, { 'id' => 1 })
+      assert_raises ActiveRecord::RecordNotFound do
+        mmmm.reload
+      end
+      # ...but immediately set it back to default to the slave
+      assert_equal 'fiveruns_city_dallas_test_slave', TheWholeEnchilada.connection.connection_name
+      
+      # Should use the master
+      TheWholeEnchilada.transaction do
+        mmmm.save!
+      end
+    end
+  end
+
+  private
+  
+  def setup_configuration_for(clazz, name)
+    flexmock(clazz).should_receive(:mysql_connection).and_return(AdapterMock.new(RawConnection.new))
+    ActiveRecord::Base.configurations ||= HashWithIndifferentAccess.new
+    ActiveRecord::Base.configurations[name] = HashWithIndifferentAccess.new({ :adapter => 'mysql', :database => name, :host => 'localhost'})
+  end
+end

data/test/database.yml

@@ -0,0 +1,27 @@
+fiveruns_city_austin_test_master:
+  adapter: mysql
+  host: localhost
+  database: vr_austin_master
+  username: root
+  password:
+
+fiveruns_city_austin_test_slave:
+  adapter: mysql
+  host: localhost
+  database: vr_austin_slave
+  username: root
+  password:
+
+fiveruns_city_dallas_test_master:
+  adapter: mysql
+  host: localhost
+  database: vr_dallas_master
+  username: root
+  password:
+
+fiveruns_city_dallas_test_slave:
+  adapter: mysql
+  host: localhost
+  database: vr_dallas_slave
+  username: root
+  password:

data/test/database_test.rb

@@ -0,0 +1,39 @@
+require File.join(File.dirname(__FILE__), 'test_helper')
+require 'flexmock/test_unit'
+require 'erb'
+
+class TheWholeBurrito < ActiveRecord::Base
+  connection_topology :prefix => 'fiveruns', :replicated => true, :shard_by => :city
+end
+
+class DatabaseTest < Test::Unit::TestCase
+  
+  def setup
+    filename = File.join(File.dirname(__FILE__), "database.yml")
+    ActiveRecord::Base.configurations = YAML::load(ERB.new(IO.read(filename)).result)
+  end
+
+  def test_live_burrito
+    DataFabric.activate_shard :city => :dallas do
+      assert_equal 'fiveruns_city_dallas_test_slave', TheWholeBurrito.connection.connection_name
+
+      # Should use the slave
+      burrito = TheWholeBurrito.find(1)
+      assert_equal 'vr_dallas_slave', burrito.name
+      
+      # Should use the master
+      burrito.reload
+      assert_equal 'vr_dallas_master', burrito.name
+
+      # ...but immediately set it back to default to the slave
+      assert_equal 'fiveruns_city_dallas_test_slave', TheWholeBurrito.connection.connection_name
+      
+      # Should use the master
+      TheWholeBurrito.transaction do
+        burrito = TheWholeBurrito.find(1)
+        assert_equal 'vr_dallas_master', burrito.name
+        burrito.save!
+      end
+    end
+  end
+end

data/test/shard_test.rb

@@ -0,0 +1,24 @@
+require File.join(File.dirname(__FILE__), 'test_helper')
+
+class ShardTest < Test::Unit::TestCase
+
+  def test_activation_should_persist_in_thread
+    DataFabric.activate_shard(:city => 'austin')
+    assert_equal 'austin', DataFabric.active_shard(:city)
+  end
+  
+  def test_activation_in_one_thread_does_not_change_another
+    assert_raises ArgumentError do
+       DataFabric.active_shard(:city)
+     end
+    DataFabric.activate_shard(:city => 'austin')
+
+    Thread.new do
+      assert_raises ArgumentError do
+         DataFabric.active_shard(:city)
+       end
+      DataFabric.activate_shard(:city => 'dallas')
+      assert_equal 'dallas', DataFabric.active_shard(:city)
+    end.join
+  end
+end

data/test/test_helper.rb

@@ -0,0 +1,17 @@
+ENV['RAILS_ENV']='test'
+RAILS_ENV='test'
+
+require 'rubygems'
+require 'test/unit'
+
+# Bootstrap AR
+gem 'activerecord', '=2.0.2'
+require 'active_record'
+require 'active_record/version'
+ActiveRecord::Base.logger = Logger.new(STDOUT)
+ActiveRecord::Base.logger.level = Logger::WARN
+ActiveRecord::Base.allow_concurrency = false
+
+# Bootstrap DF
+Dependencies.load_paths << File.join(File.dirname(__FILE__), '../lib')
+require 'init'