switchman 0.0.1

data/lib/switchman/cache_extensions.rb

@@ -0,0 +1,12 @@
+module Switchman::CacheExtensions
+  module ClassMethods
+    def cache_with_sharding
+      Switchman::Shard.current.database_server.cache_store
+    end
+  end
+
+  def self.included(klass)
+    klass.extend(ClassMethods)
+    klass.singleton_class.alias_method_chain :cache, :sharding
+  end
+end

data/lib/switchman/connection_pool_proxy.rb

@@ -0,0 +1,62 @@
+module Switchman
+  class ConnectionPoolProxy
+    delegate :spec, :connected?, :default_schema, :with_connection,
+             :to => :current_pool
+
+    attr_reader :category
+
+    def default_pool
+      @default_pool
+    end
+
+    def initialize(category, default_pool, shard_connection_pools)
+      @category = category
+      @default_pool = default_pool
+      @connection_pools = shard_connection_pools
+    end
+
+    def active_shard
+      Shard.current(@category)
+    end
+
+    def current_pool
+      pool = self.default_pool if active_shard.default?
+      pool = @connection_pools[pool_key] ||= create_pool unless pool
+      pool.shard = active_shard
+      pool
+    end
+
+    def connection
+      pool = current_pool
+      pool.connection
+    end
+
+    %w{release_connection disconnect! clear_reloadable_connections! verify_active_connections! clear_stale_cached_connections!}.each do |method|
+      class_eval(<<-EOS)
+          def #{method}
+            @connection_pools.values.each(&:#{method})
+          end
+      EOS
+    end
+
+    protected
+
+    def pool_key
+      active_shard.database_server.shareable? ? active_shard.database_server.id : active_shard
+    end
+
+    def create_pool
+      shard = active_shard
+      config = shard.database_server.config.dup
+      spec = ::ActiveRecord::Base::ConnectionSpecification.new(config, "#{config[:adapter]}_connection")
+      # unfortunately the AR code that does this require logic can't really be
+      # called in isolation
+      require "active_record/connection_adapters/#{config[:adapter]}_adapter"
+
+      ::ActiveRecord::ConnectionAdapters::ConnectionPool.new(spec).tap do |pool|
+        pool.shard = shard
+      end
+    end
+  end
+end
+

data/lib/switchman/database_server.rb

@@ -0,0 +1,197 @@
+module Switchman
+  class DatabaseServer
+    attr_accessor :id, :config
+
+    class << self
+      def all
+        database_servers.values
+      end
+
+      def find(id_or_all)
+        return self.all if id_or_all == :all
+        return id_or_all.map { |id| self.database_servers[id || Rails.env] }.compact.uniq if id_or_all.is_a?(Array)
+        database_servers[id_or_all || Rails.env]
+      end
+
+      def create(settings = {})
+        raise "database servers should be set up in database.yml" unless Rails.env.test?
+        id = 1
+        while database_servers[id.to_s]; id += 1; end
+        server = DatabaseServer.new({ :id => id.to_s }.merge(settings))
+        server.instance_variable_set(:@fake, true)
+        raise "database server #{server.id} already exists" if database_servers[server.id]
+        database_servers[server.id] = server
+      end
+
+      def server_for_new_shard
+        servers = all.select { |s| s.config[:open] }
+        return find(nil) if servers.empty?
+        servers[rand(servers.length)]
+      end
+
+      private
+      def database_servers
+        unless @database_servers
+          @database_servers = {}.with_indifferent_access
+          ::ActiveRecord::Base.configurations.each do |(id, config)|
+            @database_servers[id] = DatabaseServer.new(:id => id, :config => config)
+          end
+        end
+        @database_servers
+      end
+    end
+
+    def initialize(settings = {})
+      self.id = settings[:id]
+      self.config = (settings[:config] || {}).symbolize_keys
+    end
+
+    def destroy
+      raise "database servers should be set up in database.yml" unless Rails.env.test?
+      self.class.send(:database_servers).delete(self.id) if self.id
+    end
+
+    def fake?
+      @fake
+    end
+
+    def shareable?
+      @shareable_environment_key ||= []
+      environment = ::Shackles.environment
+      explicit_user = ::Shackles.global_config[:username]
+      return @shareable if @shareable_environment_key == [environment, explicit_user]
+      @shareable_environment_key = [environment, explicit_user]
+      username = self.config[:username]
+      username = self.config[environment][:username] if environment && self.config[environment] && self.config[environment][:username]
+      username = explicit_user if explicit_user
+      @shareable = self.config[:adapter] != 'sqlite3' && username !~ /%?\{[a-zA-Z0-9_]+\}/
+    end
+
+    def shards
+      if self.id == Rails.env
+        Shard.where("database_server_id IS NULL OR database_server_id=?", self.id)
+      else
+        Shard.where(:database_server_id => self.id)
+      end
+    end
+
+    def create_new_shard(options = {})
+      db_name = options[:db_name]
+      create_schema = options[:schema]
+      # look for another shard associated with this db
+      other_shard = self.shards.where("name<>':memory:' OR name IS NULL").order(:id).first
+      temp_db_name = other_shard.try(:name) unless id == Rails.env
+      temp_db_name = Shard.default.name if id == Rails.env
+
+      case config[:adapter]
+        when 'postgresql'
+          temp_db_name ||= 'public'
+          create_statement = lambda { "CREATE SCHEMA #{db_name}" }
+          password = " PASSWORD #{::ActiveRecord::Base.connection.quote(config[:password])}" if config[:password]
+        when 'sqlite3'
+          if db_name
+            # Try to create a db on-disk even if the only shards for sqlite are in-memory
+            temp_db_name = nil if temp_db_name == ':memory:'
+            # Put it in the db directory if there are no other sqlite shards
+            temp_db_name ||= 'db/dummy'
+            temp_db_name = File.join(File.dirname(temp_db_name), "#{db_name}.sqlite3")
+            # If they really asked for :memory:, give them :memory:
+            temp_db_name = db_name if db_name == ':memory:'
+            db_name = temp_db_name
+          end
+        else
+          temp_db_name ||= self.config[:database] % self.config
+          create_statement = lambda { "CREATE DATABASE #{db_name}" }
+      end
+      sharding_config = Switchman.config
+      config_create_statement = sharding_config[config[:adapter]].try(:[], :create_statement)
+      config_create_statement ||= sharding_config[:create_statement]
+      if config_create_statement
+        create_commands = Array(config_create_statement).dup
+        create_statement = lambda {
+          create_commands.map { |statement| statement.gsub('%{db_name}', db_name).gsub('%{password}', password || '') }
+        }
+      end
+
+      shard = Shard.create!(:name => temp_db_name,
+                            :database_server => self) do |shard|
+        shard.id = options[:id] if options[:id]
+      end
+      begin
+        if db_name.nil?
+          base_db_name = self.config[:database] % self.config
+          base_db_name = $1 if base_db_name =~ /(?:.*\/)(.+)_shard_\d+(?:\.sqlite3)?$/
+          base_db_name = nil if base_db_name == ':memory:'
+          base_db_name << '_' if base_db_name
+          db_name = "#{base_db_name}shard_#{shard.id}"
+          if config[:adapter] == 'sqlite3'
+            # Try to create a db on-disk even if the only shards for sqlite are in-memory
+            temp_db_name = nil if temp_db_name == ':memory:'
+            # Put it in the db directory if there are no other sqlite shards
+            temp_db_name ||= 'db/dummy'
+            db_name = File.join(File.dirname(temp_db_name), "#{db_name}.sqlite3")
+            shard.name = db_name
+          end
+        end
+        shard.activate(Shard.categories) do
+          ::Shackles.activate(:deploy) do
+            begin
+              if create_statement
+                Array(create_statement.call).each do |stmt|
+                  ::ActiveRecord::Base.connection.execute(stmt)
+                end
+                # have to disconnect and reconnect to the correct db
+                if self.shareable? && other_shard
+                  other_shard.activate { ::ActiveRecord::Base.connection }
+                else
+                  ::ActiveRecord::Base.connection_pool.current_pool.disconnect!
+                end
+              end
+              shard.name = db_name
+              old_proc = ::ActiveRecord::Base.connection.raw_connection.set_notice_processor {} if config[:adapter] == 'postgresql'
+              old_verbose = ::ActiveRecord::Migration.verbose
+              ::ActiveRecord::Migration.verbose = false
+              ::ActiveRecord::Migrator.migrate(Rails.root + "db/migrate/") unless create_schema == false
+            ensure
+              ::ActiveRecord::Migration.verbose = old_verbose
+              ::ActiveRecord::Base.connection.raw_connection.set_notice_processor(&old_proc) if old_proc
+            end
+          end
+        end
+        shard.save!
+        shard
+      rescue
+        shard.destroy
+        raise
+      end
+    end
+
+    def cache_store
+      if !@cache_store
+        @cache_store = Rails.cache_without_sharding if self.id == Rails.env
+        # TODO: get the config from somewhere
+        config = nil
+        @cache_store ||= ActiveSupport::Cache.lookup_store(config) if config
+        @cache_store ||= Rails.cache_without_sharding
+        @cache_store.options[:namespace] = lambda { Shard.current.default? ? nil : "shard_#{Shard.current.id}" }
+      end
+      @cache_store
+    end
+
+    def shard_name(shard)
+      if config[:shard_name]
+        config[:shard_name]
+      elsif config[:adapter] == 'postgresql'
+        if shard == :bootstrap
+          # rescue nil because the database may not exist yet; if it doesn't,
+          # it will shortly, and this will be re-invoked
+          ::ActiveRecord::Base.connection.schemas.first rescue nil
+        else
+          shard.activate { ::ActiveRecord::Base.connection_pool.default_schema }
+        end
+      else
+        config[:database]
+      end
+    end
+  end
+end

data/lib/switchman/default_shard.rb

@@ -0,0 +1,28 @@
+require_dependency 'switchman/database_server'
+
+class Switchman::DefaultShard
+  def id; 'default'; end
+  def activate(*categories); yield; end
+  def activate!(*categories); end
+  def default?; true; end
+  def relative_id_for(local_id, target = nil); local_id; end
+  def global_id_for(local_id); local_id; end
+  def database_server_id; nil; end
+  def database_server; ::Switchman::DatabaseServer.find(nil); end
+  def name
+    unless instance_variable_defined?(:@name)
+      @name = nil # prevent taking this branch on recursion
+      @name = database_server.shard_name(:bootstrap)
+    end
+    @name
+  end
+  def description; Rails.env; end
+  # The default's shard is always the default shard
+  def shard; self; end
+  def _dump(depth)
+    ''
+  end
+  def self._load(str)
+    Shard.default
+  end
+end

data/lib/switchman/engine.rb

@@ -0,0 +1,91 @@
+module Switchman
+  class Engine < ::Rails::Engine
+    isolate_namespace Switchman
+
+    initializer 'switchman.extend_ar', :before => "active_record.initialize_database" do
+      ActiveSupport.on_load(:active_record) do
+        #require 'active_record/associations/preloader/belongs_to'
+
+        require "switchman/active_record/abstract_adapter"
+        require "switchman/active_record/association"
+        require "switchman/active_record/attribute_methods"
+        require "switchman/active_record/base"
+        require "switchman/active_record/calculations"
+        require "switchman/active_record/connection_handler"
+        require "switchman/active_record/connection_pool"
+        require "switchman/active_record/finder_methods"
+        require "switchman/active_record/log_subscriber"
+        require "switchman/active_record/query_cache"
+        require "switchman/active_record/query_methods"
+        require "switchman/active_record/relation"
+        require "switchman/cache_extensions"
+
+        include ActiveRecord::Base
+        include ActiveRecord::AttributeMethods
+        ::ActiveRecord::Associations::Association.send(:include, ActiveRecord::Association)
+        ::ActiveRecord::Associations::CollectionProxy.send(:include, ActiveRecord::CollectionProxy)
+        ::ActiveRecord::Associations::Builder::Association.send(:include, ActiveRecord::Builder::Association)
+
+        ::ActiveRecord::Associations::Preloader::Association.send(:include, ActiveRecord::Preloader::Association)
+        ::ActiveRecord::ConnectionAdapters::AbstractAdapter.send(:include, ActiveRecord::AbstractAdapter)
+        ::ActiveRecord::ConnectionAdapters::ConnectionHandler.send(:include, ActiveRecord::ConnectionHandler)
+        ::ActiveRecord::ConnectionAdapters::ConnectionPool.send(:include, ActiveRecord::ConnectionPool)
+        ::ActiveRecord::ConnectionAdapters::AbstractAdapter.send(:include, ActiveRecord::QueryCache)
+        ::ActiveRecord::LogSubscriber.send(:include, ActiveRecord::LogSubscriber)
+        ::ActiveRecord::Relation.send(:include, ActiveRecord::Calculations)
+        ::ActiveRecord::Relation.send(:include, ActiveRecord::FinderMethods)
+        ::ActiveRecord::Relation.send(:include, ActiveRecord::QueryMethods)
+        ::ActiveRecord::Relation.send(:include, ActiveRecord::Relation)
+        Rails.send(:include, CacheExtensions)
+      end
+    end
+
+    def self.foreign_key_check(name, type, options)
+      if name.to_s =~ /_id\z/ && type.to_s == 'integer' && options[:limit].to_i < 8
+        puts "WARNING: All foreign keys need to be 8-byte integers. #{name} looks like a foreign key. If so, please add the option: `:limit => 8`"
+      end
+    end
+
+    initializer 'switchman.extend_connection_adapters', :after => "active_record.initialize_database" do
+      ActiveSupport.on_load(:active_record) do
+        ::ActiveRecord::ConnectionAdapters::AbstractAdapter.descendants.each do |klass|
+          klass.class_eval do
+            def add_column_with_foreign_key_check(table, name, type, options = {})
+              Switchman::Engine.foreign_key_check(name, type, options)
+              add_column_without_foreign_key_check(table, name, type, options)
+            end
+            alias_method_chain(:add_column, :foreign_key_check)
+          end
+        end
+
+        ::ActiveRecord::ConnectionAdapters::TableDefinition.class_eval do
+          def column_with_foreign_key_check(name, type, options = {})
+            Switchman::Engine.foreign_key_check(name, type, options)
+            column_without_foreign_key_check(name, type, options)
+          end
+          alias_method_chain(:column, :foreign_key_check)
+        end
+
+        if defined?(::ActiveRecord::ConnectionAdapters::PostgreSQLAdapter)
+          require "switchman/active_record/postgresql_adapter"
+          ::ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.send(:include, ActiveRecord::PostgreSQLAdapter)
+        end
+      end
+    end
+
+    initializer 'switchman.eager_load' do
+      ActiveSupport.on_load(:before_eager_load) do
+        # This needs to be loaded before Switchman::Shard, otherwise it won't autoload it correctly
+        require_dependency('active_record/base')
+      end
+    end
+
+    initializer 'switchman.extend_shackles', :after => "shackles.extend_ar" do
+      ActiveSupport.on_load(:active_record) do
+        require "switchman/shackles"
+
+        ::Shackles.send(:include, Shackles)
+      end
+    end
+  end
+end

data/lib/switchman/r_spec_helper.rb

@@ -0,0 +1,124 @@
+require_dependency "switchman/test_helper"
+
+module Switchman
+  # including this module in your specs will give you several shards to
+  # work with during specs:
+  #  * Shard.default - the test database itself
+  #  * @shard1 - a shard possibly using the same connection as Shard.default
+  #  * @shard2 - a shard using a dedicated connection
+  #  * @shard3 - a shard using the same connection as @shard1 (this might
+  #              be Shard.default if they already share a connection, or
+  #              a separate shard)
+  module RSpecHelper
+    @@keep_the_shards = false
+    @@shard1 = nil
+
+    def self.included(klass)
+      # our before handlers have already been configured from a parent group; don't add them again
+      return if klass.parent_groups[1..-1].any? { |group| group.included_modules.include?(self) }
+
+      ::RSpec.configure do |config|
+        block = proc do
+          if @@shard1
+            # some specs are mean, and blow away our shards
+            begin
+              @@shard1.reload
+              @@shard2.reload
+              @@shard3.reload if @@shard3
+            rescue ::ActiveRecord::RecordNotFound
+              Shard.default.clone.save!
+              Shard.default(true)
+              @@shard1 = @@shard1.clone
+              @@shard1.save!
+              @@shard2 = @@shard2.clone
+              @@shard2.save!
+              if @@shard3
+                @@shard3 = @@shard3.clone
+                @@shard3.save!
+              end
+            end
+          end
+        end
+        # this module will be included multiple times, but we don't want to run this global hook multiple times
+        if config.hooks[:before][:all].all? { |hook| hook.block.source_location != block.source_location }
+          config.before(:all, &block)
+        end
+      end
+
+      klass.before(:all) do
+        unless @@shard1
+          puts "Setting up sharding for all specs..."
+          @@shard1, @@shard2 = TestHelper.recreate_persistent_test_shards
+          if @@shard1.is_a?(Shard)
+            @@keep_the_shards = true
+            @@shard3 = nil
+          else # @@shard1.is_a?(DatabaseServer)
+            @@shard1 = @@shard1.create_new_shard
+            @@shard2 = @@shard2.create_new_shard
+            if @@shard1.database_server == Shard.default.database_server
+              @@shard3 = nil
+            else
+              @@shard3 = @@shard1.database_server.create_new_shard
+            end
+          end
+          puts "Done!"
+
+          at_exit do
+            # preserve rspec's exit status
+            status= $!.is_a?(::SystemExit) ? $!.status : nil
+            puts "Tearing down sharding for all specs"
+            @@shard1.database_server.destroy unless @@shard1.database_server == Shard.default.database_server
+            unless @@keep_the_shards
+              @@shard1.drop_database
+              @@shard1.destroy
+              @@shard2.drop_database
+              @@shard2.destroy
+              if @@shard3
+                @@shard3.drop_database
+                @@shard3.destroy
+              end
+            end
+            @@shard2.database_server.destroy
+            exit status if status
+          end
+        end
+        @shard1, @shard2 = @@shard1, @@shard2
+        @shard3 = @@shard3 ? @@shard3 : Shard.default
+      end
+
+      klass.before do
+        Shard.clear_cache
+        if use_transactional_fixtures
+          Shard.default(true)
+          @shard1 = Shard.find(@shard1)
+          @shard2 = Shard.find(@shard2)
+          shards = [@shard2]
+          shards << @shard1 unless @shard1.database_server == Shard.default.database_server
+          shards.each do |shard|
+            shard.activate do
+              # this is how AR does it in fixtures.rb
+              ::ActiveRecord::Base.connection.increment_open_transactions
+              ::ActiveRecord::Base.connection.transaction_joinable = false
+              ::ActiveRecord::Base.connection.begin_db_transaction
+            end
+          end
+        end
+      end
+
+      klass.after do
+        if use_transactional_fixtures
+          shards = [@shard2]
+          shards << @shard1 unless @shard1.database_server == Shard.default.database_server
+          shards.each do |shard|
+            shard.activate do
+              if ::ActiveRecord::Base.connection.open_transactions != 0
+                ::ActiveRecord::Base.connection.rollback_db_transaction
+                ::ActiveRecord::Base.connection.decrement_open_transactions
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+end