ar-octopus 0.0.1

data/lib/octopus.rb

@@ -0,0 +1,30 @@
+require "yaml"
+
+module Octopus
+  def self.env()
+    if defined?(Rails)
+      Rails.env.to_s
+    else
+      "production"
+    end
+  end
+  
+  def self.config()
+    @@config ||= YAML.load_file(Octopus.directory() + "/config/shards.yml") 
+  end
+
+  def self.directory()
+    if defined?(Rails)
+      # Running in a normal Rails application
+      Rails.root.to_s
+    else
+      # Running in a generic Ruby process
+      Dir.pwd
+    end
+  end
+end
+
+require "octopus/proxy"
+require "octopus/migration"
+require "octopus/model"
+require "octopus/controller"

data/lib/octopus/controller.rb

@@ -0,0 +1,2 @@
+class Octopus::Controller
+end

data/lib/octopus/migration.rb

@@ -0,0 +1,41 @@
+module Octopus::Migration
+  def self.included(base)
+    base.extend(ClassMethods)
+    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
+    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
+end
+
+
+ActiveRecord::Migration.send(:include, Octopus::Migration)

data/lib/octopus/model.rb

@@ -0,0 +1,57 @@
+module Octopus::Model  
+  def self.included(base)
+    base.extend(ClassMethods)    
+    base.send(:include, InstanceMethods)
+  end
+
+  module InstanceMethods
+    def connection_proxy
+      self.class.connection_proxy
+    end
+    
+    def using(shard, &block)
+      class << self
+        def connection_proxy
+          @@connection_proxy ||= Octopus::Proxy.new(Octopus.config())
+        end
+
+        def connection 
+          if self.respond_to?(:replicated)
+            self.connection_proxy().set_replicated_model(self)
+          end
+
+          self.connection_proxy()
+        end
+
+        def connected?
+          self.connection_proxy().connected?
+        end
+      end
+      
+      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
+      else
+        self.connection_proxy.current_shard = shard
+        return self
+      end
+    end
+  end
+
+  module ClassMethods
+    include InstanceMethods
+
+    def replicated_model()
+      self.cattr_accessor :replicated
+    end
+  end  
+end
+
+ActiveRecord::Base.send(:include, Octopus::Model)

data/lib/octopus/proxy.rb

@@ -0,0 +1,195 @@
+require "set"
+
+class Octopus::Proxy
+  attr_accessor :shards, :current_shard, :block, :groups, :current_group,  :replicated, :slaves_list, :replicated_models
+
+  delegate :increment_open_transactions, :decrement_open_transactions,  :to => :select_connection
+
+  def initialize(config)
+    @shards = {}
+    @groups = {}
+    @replicated_models = Set.new
+    @block = false
+    @replicated = config[Octopus.env()]["replicated"] || false
+    @shards[:master] = ActiveRecord::Base.connection_pool()
+    @current_shard = :master
+
+    config[Octopus.env()]["shards"].each do |key, value|
+      if value.has_key?("adapter")
+        initialize_adapter(value['adapter'])
+        @shards[key.to_sym] = connection_pool_for(value, "#{value['adapter']}_connection")
+      else
+        @groups[key.to_sym] = []
+
+        value.each do |k, v|
+          raise "You have duplicated shard names!" if @shards.has_key?(k.to_sym)
+          initialize_adapter(v['adapter'])
+          @shards[k.to_sym] = connection_pool_for(v, "#{v['adapter']}_connection")
+          @groups[key.to_sym] << k.to_sym
+        end
+      end
+    end
+
+    if @replicated
+      @slaves_list = @shards.keys
+      @slaves_list.delete(:master)
+      @slaves_list = @slaves_list.map {|sym| sym.to_s}.sort
+    end
+  end
+
+  def current_shard=(shard_symbol)
+    if shard_symbol.is_a?(Array)
+      shard_symbol.each {|symbol| raise "Nonexistent Shard Name: #{symbol}" if @shards[symbol].nil? }
+    else
+      raise "Nonexistent Shard Name: #{shard_symbol}" if @shards[shard_symbol].nil?
+    end
+
+    @current_shard = shard_symbol
+  end
+
+  def current_group=(group_symbol)
+    if group_symbol.is_a?(Array)
+      group_symbol.each {|symbol| raise "Nonexistent Group Name: #{symbol}" if @groups[symbol].nil? }
+    else
+      raise "Nonexistent Group Name: #{group_symbol}" if @groups[group_symbol].nil? && !group_symbol.nil?
+    end
+
+    @current_group = group_symbol
+  end
+
+  def select_connection()
+    @shards[shard_name].connection()
+  end
+
+  def shard_name
+    if(current_shard.is_a?(Array))
+      current_shard.first
+    else
+      current_shard
+    end
+  end
+  
+  def set_replicated_model(model)
+    replicated_models << model.to_s
+  end
+
+  def transaction(start_db_transaction = true, &block)
+    if should_send_queries_to_multiple_shards?
+      self.send_transaction_to_multiple_shards(current_shard, start_db_transaction, &block)
+      self.current_shard = :master
+    elsif should_send_queries_to_multiple_groups?
+      self.send_transaction_to_multiple_groups(start_db_transaction, &block)
+      self.current_group = nil      
+    elsif should_send_queries_to_a_group_of_shards?
+      self.send_transaction_to_multiple_shards(@groups[current_group], start_db_transaction, &block)
+      self.current_group = nil
+    else
+      select_connection.transaction(start_db_transaction, &block) 
+    end
+  end
+
+  def method_missing(method, *args, &block)
+    if should_clean_connection?(method)
+      conn = select_connection()
+      self.current_shard = :master
+      conn.send(method, *args, &block)
+    elsif should_send_queries_to_replicated_databases?(method)
+      send_queries_to_selected_slave(method, *args, &block)      
+    elsif should_send_queries_to_multiple_groups?
+      send_queries_to_multiple_groups(method, *args, &block)
+    elsif should_send_queries_to_multiple_shards?
+      send_queries_to_shards(current_shard, method, *args, &block)
+    elsif should_send_queries_to_a_group_of_shards?
+      send_queries_to_shards(@groups[current_group], method, *args, &block)
+    else
+      select_connection().send(method, *args, &block)
+    end
+  end
+
+  protected
+  def connection_pool_for(adapter, config)
+    ActiveRecord::ConnectionAdapters::ConnectionPool.new(ActiveRecord::Base::ConnectionSpecification.new(adapter, config))
+  end
+
+  def initialize_adapter(adapter)
+    begin
+      require 'rubygems'
+      gem "activerecord-#{adapter}-adapter"
+      require "active_record/connection_adapters/#{adapter}_adapter"
+    rescue LoadError
+      begin
+        require "active_record/connection_adapters/#{adapter}_adapter"
+      rescue LoadError
+        raise "Please install the #{adapter} adapter: `gem install activerecord-#{adapter}-adapter` (#{$!})"
+      end
+    end
+  end
+
+  def should_clean_connection?(method)
+    method.to_s =~ /begin_db_transaction|insert|select/ && !should_send_queries_to_multiple_shards? && !self.current_group && !replicated
+  end
+
+  def should_send_queries_to_multiple_shards?
+    current_shard.is_a?(Array)
+  end
+
+  def should_send_queries_to_multiple_groups?
+    current_group.is_a?(Array)
+  end
+
+  def should_send_queries_to_a_group_of_shards?
+    !current_group.nil?
+  end
+
+  def should_send_queries_to_replicated_databases?(method)
+    @replicated && method.to_s == "select_all"
+  end
+
+  def send_queries_to_multiple_groups(method, *args, &block)
+    method_return = nil
+
+    current_group.each do |group_symbol|
+      method_return = self.send_queries_to_shards(@groups[group_symbol], method, *args, &block)
+    end
+
+    return method_return
+  end
+
+  def send_queries_to_shards(shard_array, method, *args, &block)
+    method_return = nil
+
+    shard_array.each do |shard_symbol|
+      method_return = @shards[shard_symbol].connection().send(method, *args, &block) 
+    end
+
+    return method_return
+  end
+
+  def send_queries_to_selected_slave(method, *args, &block)        
+    #TODO: ugly code, needs refactor
+    if args.last =~ /#{replicated_models.to_a.join('|')}/
+      old_shard = self.current_shard
+      self.current_shard = slaves_list.shift.to_sym
+      slaves_list << self.current_shard      
+    else
+      old_shard = self.current_shard
+      self.current_shard = :master
+    end
+
+    sql = select_connection().send(method, *args, &block)     
+    self.current_shard = old_shard
+    return sql    
+  end
+
+  def send_transaction_to_multiple_shards(shard_array, start_db_transaction, &block)
+    shard_array.each do |shard_symbol|
+      @shards[shard_symbol].connection().transaction(start_db_transaction, &block)
+    end
+  end
+
+  def send_transaction_to_multiple_groups(start_db_transaction, &block)
+    current_group.each do |group_symbol|
+      self.send_transaction_to_multiple_shards(@groups[group_symbol], start_db_transaction, &block)
+    end
+  end
+end

data/spec/config/shards.yml

@@ -0,0 +1,76 @@
+production:
+  shards:
+    alone_shard:
+      adapter: mysql
+      host: localhost
+      database: octopus_shard5
+    
+    postgresql_shard:
+      adapter: postgresql
+      username: postgres
+      password:
+      database: octopus_shard1
+      encoding: unicode
+    
+    history_shards:
+      aug2009:
+        adapter: mysql
+        host: localhost
+        database: octopus_shard2
+      aug2010:
+        adapter: mysql
+        host: localhost
+        database: octopus_shard3
+      aug2011:
+        adapter: mysql
+        host: localhost
+        database: octopus_shard4
+    
+    country_shards:
+      canada:
+        adapter: mysql
+        host: localhost
+        database: octopus_shard2
+      brazil:
+        adapter: mysql
+        host: localhost
+        database: octopus_shard3
+      russia:
+        adapter: mysql
+        host: localhost
+        database: octopus_shard4
+        
+        
+production_raise_error:
+  shards:
+    history_shards:
+      duplicated_shard_name:
+        adapter: mysql
+        host: localhost
+        database: octopus_shard5
+    
+    country_shards:
+      duplicated_shard_name:
+        adapter: mysql
+        host: localhost
+        database: octopus_shard4
+        
+production_replicated:
+  replicated: true
+  shards:
+    slave1:
+      adapter: mysql
+      host: localhost
+      database: octopus_shard2
+    slave2:
+      adapter: mysql
+      host: localhost
+      database: octopus_shard3
+    slave3:
+      adapter: mysql
+      host: localhost
+      database: octopus_shard4
+    slave4:
+      adapter: mysql
+      host: localhost
+      database: octopus_shard5

data/spec/database_connection.rb

@@ -0,0 +1,6 @@
+require 'rubygems'
+require 'active_record'
+require 'logger'
+require "pg"
+
+ActiveRecord::Base.establish_connection(:adapter => "mysql", :database => "octopus_shard1", :username => "root", :password => "")

data/spec/database_models.rb

@@ -0,0 +1,23 @@
+#The user class is just sharded, not replicated
+class User < ActiveRecord::Base
+  def awesome_queries
+    using(:canada) do
+      User.create(:name => "teste")
+    end
+  end
+end
+
+#The client class isn't replicated
+class Client < ActiveRecord::Base
+  has_many :items
+end
+
+#This class is replicated
+class Cat < ActiveRecord::Base
+  replicated_model()
+end
+
+#This items belongs to a client
+class Item < ActiveRecord::Base
+  belongs_to :client
+end

data/spec/migrations/1_create_users_on_master.rb

@@ -0,0 +1,9 @@
+class CreateUsersOnMaster < ActiveRecord::Migration  
+  def self.up
+    User.create!(:name => "Master")
+  end
+
+  def self.down
+    User.delete_all()
+  end
+end