switchman 0.0.1 → 1.0.0

data/lib/switchman/database_server.rb

@@ -1,6 +1,6 @@
 module Switchman
   class DatabaseServer
-    attr_accessor :id, :config
+    attr_accessor :id
 
     class << self
       def all
@@ -9,12 +9,12 @@ module Switchman
 
       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]
+        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?
+        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))
@@ -42,12 +42,13 @@ module Switchman
     end
 
     def initialize(settings = {})
-      self.id = settings[:id]
-      self.config = (settings[:config] || {}).symbolize_keys
+      @id = settings[:id]
+      @config = (settings[:config] || {}).symbolize_keys
+      @configs = {}
     end
 
     def destroy
-      raise "database servers should be set up in database.yml" unless Rails.env.test?
+      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
 
@@ -55,33 +56,83 @@ module Switchman
       @fake
     end
 
+    def config(environment = :master)
+      @configs[environment] ||= begin
+        if @config[environment].is_a?(Array)
+          @config[environment].map do |config|
+            config = @config.merge((config || {}).symbolize_keys)
+            # make sure Shackles doesn't get any brilliant ideas about choosing the first possible server
+            config.delete(environment)
+            config
+          end
+        elsif @config[environment].is_a?(Hash)
+          @config.merge(@config[environment])
+        else
+          @config
+        end
+      end
+    end
+
+    def shackles_environment
+      @shackles_environment || ::Shackles.environment
+    end
+
+    # locks this db to a specific environment, except for
+    # when doing writes (then it falls back to the current
+    # value of Shackles.environment)
+    def shackle!(environment = :slave)
+      @shackles_environment = environment
+    end
+
+    def unshackle!
+      @shackles_environment = nil
+    end
+
+    def unshackle
+      old_env = @shackles_environment
+      unshackle!
+      yield
+    ensure
+      shackle!(old_env)
+    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
+      if explicit_user
+        username = explicit_user
+      else
+        config = self.config(environment)
+        config = config.first if config.is_a?(Array)
+        username = config[:username]
+      end
       @shareable = self.config[:adapter] != 'sqlite3' && username !~ /%?\{[a-zA-Z0-9_]+\}/
     end
 
     def shards
-      if self.id == Rails.env
+      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 pool_key
+      self.id == ::Rails.env ? nil : self.id
+    end
+
     def create_new_shard(options = {})
+      raise MethodNotImplemented unless Shard.default.is_a?(Shard)
+
       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
+      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'
@@ -113,67 +164,78 @@ module Switchman
         }
       end
 
-      shard = Shard.create!(:name => temp_db_name,
+      create_shard = lambda do
+        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
+          shard.id = options[:id] if options[:id]
         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!
+        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
+
+                reset_column_information
+                ::ActiveRecord::Migrator.migrate(::Rails.root + "db/migrate/") unless create_schema == false
+                reset_column_information
+              ensure
+                ::ActiveRecord::Migration.verbose = old_verbose
+                ::ActiveRecord::Base.connection.raw_connection.set_notice_processor(&old_proc) if old_proc
               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
+          shard.save!
+          shard
+        rescue
+          shard.destroy
+          shard.drop_database if shard.name == db_name rescue nil
+          raise
         end
-        shard.save!
-        shard
-      rescue
-        shard.destroy
-        raise
+      end
+
+      if Shard.connection.supports_ddl_transactions? && self.shareable? && other_shard
+        Shard.transaction do
+          other_shard.activate do
+            ::ActiveRecord::Base.connection.transaction(&create_shard)
+          end
+        end
+      else
+        create_shard.call
       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}" }
+      unless @cache_store
+        @cache_store = Switchman.config[:cache_map][self.id] || Switchman.config[:cache_map][::Rails.env]
       end
       @cache_store
     end
@@ -193,5 +255,11 @@ module Switchman
         config[:database]
       end
     end
+
+    private
+    def reset_column_information
+      ::ActiveRecord::Base.connection.schema_cache.clear!
+      ::ActiveRecord::Base.descendants.reject{ |m| m == Shard }.each(&:reset_column_information)
+    end
   end
 end

data/lib/switchman/default_shard.rb

@@ -1,28 +1,43 @@
 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)
+module Switchman
+  class 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; 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
+
+    def ==(rhs)
+      return true if rhs.is_a?(Shard) && rhs.default?
+      super
+    end
+
+    class << self
+      def instance
+        @instance ||= new
+      end
+
+      private :new
     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

@@ -2,10 +2,55 @@ 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'
+    initializer 'switchman.initialize_cache', :before => 'initialize_cache' do
+      # if we haven't already setup our cache map out-of-band, set it up from
+      # config.cache_store now. behaves similarly to Rails' default
+      # initialize_cache initializer, but for each value in the map, rather
+      # than just Rails.cache. if config.cache_store is a flat value, uses it
+      # to fill just the Rails.env entry in the cache map.
+      unless Switchman.config[:cache_map].present?
+        cache_store_config = ::Rails.configuration.cache_store
+        unless cache_store_config.is_a?(Hash)
+          cache_store_config = {::Rails.env => cache_store_config}
+        end
+
+        Switchman.config[:cache_map] = {}
+        cache_store_config.each do |key, value|
+          value = ::ActiveSupport::Cache.lookup_store(value)
+          Switchman.config[:cache_map][key] = value
+          if value.respond_to?(:middleware)
+            config.middleware.insert_before("Rack::Runtime", value.middleware)
+          end
+        end
+      end
+
+      # if the configured cache map (either from before, or as populated from
+      # config.cache_store) didn't have an entry for Rails.env, add one using
+      # lookup_store(nil); matches the behavior of Rails' default
+      # initialize_cache initializer when config.cache_store is nil.
+      unless Switchman.config[:cache_map].has_key?(::Rails.env)
+        value = ::ActiveSupport::Cache.lookup_store(nil)
+        Switchman.config[:cache_map][::Rails.env] = value
+        if value.respond_to?(:middleware)
+          config.middleware.insert_before("Rack::Runtime", value.middleware)
+        end
+      end
 
+      # prevent :initialize_cache from trying to (or needing to) set
+      # Rails.cache. once our switchman.extend_ar initializer (below) runs
+      # Rails.cache will be overridden to pull appropriate values from the
+      # cache map, but between now and then, Rails.cache should return the
+      # Rails.env entry in the cache map.
+      if ::Rails.version < '4'
+        silence_warnings { Object.const_set "RAILS_CACHE", Switchman.config[:cache_map][::Rails.env] }
+      else
+        ::Rails.cache = Switchman.config[:cache_map][::Rails.env]
+      end
+    end
+
+    initializer 'switchman.extend_ar', :before => "active_record.initialize_database" do
+      ::ActiveSupport.on_load(:active_record) do
+        require "switchman/active_support/cache"
         require "switchman/active_record/abstract_adapter"
         require "switchman/active_record/association"
         require "switchman/active_record/attribute_methods"
@@ -18,11 +63,14 @@ module Switchman
         require "switchman/active_record/query_cache"
         require "switchman/active_record/query_methods"
         require "switchman/active_record/relation"
-        require "switchman/cache_extensions"
+        require "switchman/active_record/spawn_methods"
+        require "switchman/rails"
 
+        ::ActiveSupport::Cache::Store.send(:include, ActiveSupport::Cache::Store)
         include ActiveRecord::Base
         include ActiveRecord::AttributeMethods
         ::ActiveRecord::Associations::Association.send(:include, ActiveRecord::Association)
+        ::ActiveRecord::Associations::BelongsToAssociation.send(:include, ActiveRecord::BelongsToAssociation)
         ::ActiveRecord::Associations::CollectionProxy.send(:include, ActiveRecord::CollectionProxy)
         ::ActiveRecord::Associations::Builder::Association.send(:include, ActiveRecord::Builder::Association)
 
@@ -31,12 +79,19 @@ module Switchman
         ::ActiveRecord::ConnectionAdapters::ConnectionHandler.send(:include, ActiveRecord::ConnectionHandler)
         ::ActiveRecord::ConnectionAdapters::ConnectionPool.send(:include, ActiveRecord::ConnectionPool)
         ::ActiveRecord::ConnectionAdapters::AbstractAdapter.send(:include, ActiveRecord::QueryCache)
+        # when we call super in Switchman::ActiveRecord::QueryCache#select_all,
+        # we want it to find the definition from
+        # ActiveRecord::ConnectionAdapters::DatabaseStatements, not
+        # ActiveRecord::ConnectionAdapters::QueryCache
+        ::ActiveRecord::ConnectionAdapters::QueryCache.send(:remove_method, :select_all)
+
         ::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)
+        ::ActiveRecord::Relation.send(:include, ActiveRecord::SpawnMethods)
+        ::Rails.send(:include, Rails)
       end
     end
 
@@ -47,7 +102,7 @@ module Switchman
     end
 
     initializer 'switchman.extend_connection_adapters', :after => "active_record.initialize_database" do
-      ActiveSupport.on_load(:active_record) 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 = {})
@@ -74,18 +129,26 @@ module Switchman
     end
 
     initializer 'switchman.eager_load' do
-      ActiveSupport.on_load(:before_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
+      ::ActiveSupport.on_load(:active_record) do
         require "switchman/shackles"
 
         ::Shackles.send(:include, Shackles)
       end
     end
+
+    initializer 'switchman.extend_controller', :after => "shackles.extend_ar" do
+      ::ActiveSupport.on_load(:action_controller) do
+        require "switchman/action_controller/caching"
+
+        ::ActionController::Base.send(:include, ActionController::Caching)
+      end
+    end
   end
 end

data/lib/switchman/r_spec_helper.rb

@@ -17,48 +17,28 @@ module Switchman
       # 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..."
+          Shard.delete_all
+
           @@shard1, @@shard2 = TestHelper.recreate_persistent_test_shards
+          @@default_shard = Shard.default
           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
+            begin
+              @@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
+            rescue
+              @@shard1 = @@shard2 = @@shard3 = nil
+              raise
             end
           end
           puts "Done!"
@@ -81,6 +61,25 @@ module Switchman
             @@shard2.database_server.destroy
             exit status if status
           end
+        else
+          dup = @@default_shard.dup
+          dup.id = @@default_shard.id
+          dup.save!
+          Shard.default(true)
+          dup = @@shard1.dup
+          dup.id = @@shard1.id
+          dup.save!
+          @@shard1.instance_variable_set(:@destroyed, false)
+          dup = @@shard2.dup
+          dup.id = @@shard2.id
+          dup.save!
+          @@shard2.instance_variable_set(:@destroyed, false)
+          if @@shard3
+            dup = @@shard3.dup
+            dup.id = @@shard3.id
+            dup.save!
+            @@shard3.instance_variable_set(:@destroyed, false)
+          end
         end
         @shard1, @shard2 = @@shard1, @@shard2
         @shard3 = @@shard3 ? @@shard3 : Shard.default
@@ -119,6 +118,14 @@ module Switchman
           end
         end
       end
+
+      klass.after(:all) do
+        Shard.default.destroy
+        @@shard1.destroy
+        @@shard2.destroy
+        @@shard3.destroy if @@shard3
+        Shard.default(true)
+      end
     end
   end
 end

data/lib/switchman/rails.rb

@@ -0,0 +1,21 @@
+module Switchman::Rails
+  module ClassMethods
+    def cache_with_sharding
+      Switchman::Shard.current.database_server.cache_store
+    end
+
+    # in Rails 4+, the Rails.cache= method was used during bootstrap to set
+    # Rails.cache(_without_sharding) to the value from the config file. but now
+    # that that's done (the bootstrap happened before this module is included
+    # into Rails), we want to make sure no one tries to assign to Rails.cache,
+    # because it would be wrong w.r.t. sharding.
+    def cache=(new_cache)
+      raise NoMethodError
+    end
+  end
+
+  def self.included(klass)
+    klass.extend(ClassMethods)
+    klass.singleton_class.alias_method_chain :cache, :sharding
+  end
+end

data/lib/switchman/shackles.rb

@@ -2,33 +2,33 @@ module Switchman
   module Shackles
     module ClassMethods
       def ensure_handler
-        Shard.default.activate(*Shard.categories) do
-          new_handler = @connection_handlers[self.environment]
-          if !new_handler
-            new_handler = @connection_handlers[self.environment] = ::ActiveRecord::ConnectionAdapters::ConnectionHandler.new
-            pools = ::ActiveRecord::Base.connection_handler.instance_variable_get(:@class_to_pool)
-            pools.each do |model, pool|
-              # don't call establish_connection for pools created just for different sharding categories
-              if model != ::ActiveRecord::Base.name
-                klass = model.constantize
-                next if ::ActiveRecord::Base.connection_handler.retrieve_connection_pool(klass.superclass).default_pool == pools[model].default_pool
-              end
+        raise "This should not be called with switchman installed"
+      end
+
+      # drops the save_handler and ensure_handler calls from the vanilla
+      # Shackles' implementation.
+      def activate!(environment)
+        environment ||= :master
+        old_environment = self.environment
+        @environment = environment
+        old_environment
+      end
 
-              new_handler.establish_connection(model, pool.spec)
-            end
-          end
-          # make sure it picks up the environment change
-          new_handler.connection_pools.each do |_, pool|
-            pool.spec.instance_variable_set(:@current_config, nil)
-          end
-          new_handler
-        end
+      # since activate! really is just a variable swap now, it's safe to use in
+      # the ensure block, simplifying the implementation
+      def activate(environment)
+        old_environment = activate!(environment)
+        yield
+      ensure
+        activate!(old_environment)
       end
     end
 
     def self.included(klass)
       klass.extend(ClassMethods)
       klass.singleton_class.send(:remove_method, :ensure_handler)
+      klass.singleton_class.send(:remove_method, :activate!)
+      klass.singleton_class.send(:remove_method, :activate)
     end
   end
 end

data/lib/switchman/sharded_instrumenter.rb

@@ -0,0 +1,19 @@
+module Switchman
+  class ShardedInstrumenter < ::SimpleDelegator
+    def initialize(instrumenter, shard_host)
+      super instrumenter
+      @shard_host = shard_host
+    end
+
+    def instrument(name, payload={})
+      shard = @shard_host.try(:shard)
+      if shard.is_a?(Shard)
+        payload[:shard] = {
+          id: shard.id,
+          env: shard.database_server.shackles_environment
+        }
+      end
+      super name, payload
+    end
+  end
+end