xbar 0.0.1 → 0.4.0

data/lib/xbar/mapper.rb

@@ -0,0 +1,304 @@
+module XBar
+  #
+  # This module holds the current configuration. It is read from a JSON document
+  # and always represents exactly the state of that document. The configuration
+  # should not be 'tweaked' by changing the state of in-memory structures.  The
+  # approved way to change the configuration is to call:
+  #
+  #     XBar::Mapper.reset(:xbar_env => "<config file>",
+  #       :app_env => "<application environment>")
+  # 
+  # This loads a new configuration file from the XBar 'config' directory.
+  # This file may contain multiple 'environments'.  Only the specified
+  # environment is used.  Both arguments are required.  Changes to the
+  # configuration via this API cause all Proxies and their Shards to be reset.
+  # This is, all Proxies are found and 'cleaned', and the Shard references held
+  # by each Proxy are dropped and new Shards are allocated.
+  #
+  # No thread-specific state is kept in the Mapper structure. In fact, the state
+  # it encapsulates is shared by all threads.  In contrast, thread-specific
+  # state is kept in instances of the Proxy class, or instances of the Shard
+  # class.  Thread specific state is handled as follows:
+  #
+  # Thread.current[:connection_proxy] references an instance of Proxy.
+  # In turn, an an instance of XBar::Proxy references a list of instances of
+  # XBar::Shard.
+  #
+  # Each Shard is considered to be an array of replicas, even if the
+  # configuration JSON specifies a single Connection (as a String literal).  In
+  # this case, the Shard is an array of one replica.  At any point in time, each
+  # Shard has one replica designated as the master.  Thus the complete Shard
+  # description always an array of Connections and the first Connection in the
+  # array is considered to be the master.
+  #
+  # This module is included in the XBar::Proxy class.  This adds instance
+  # methods to instances of XBar::Proxy that allow the configuration state to
+  # be read.
+  #
+  # The following data structures are maintained:
+  #
+  #     @@connections -- a Hash, the key is the connection name and the value
+  #       is a connection pool.
+  #     @@shards -- an array of Hashes.  For each hash the key is the shard name
+  #       and the value is a connection pool.
+  #
+  module Mapper
+    
+    def self.exports
+      # omit master_config
+      %w(connections shards adapters options app_env proxies).map(&:to_sym)
+    end
+    
+    module ClassMethods
+      
+      Mapper.exports.each do |var|
+        mattr_reader var
+      end
+      
+      @@cached_config = nil
+      @@shards = HashWithIndifferentAccess.new
+      @@connections = HashWithIndifferentAccess.new
+      @@proxies = []
+      @@adapters = Set.new
+      @@config = nil
+      @@app_env = nil
+      @@xbar_env = nil
+        
+      def config_file_name
+        file = "#{xbar_env}.json"
+        "#{XBar.directory}/config/#{file}"
+      end
+      
+      def connection_file_name
+        file = "connection.rb"
+        "#{XBar.directory}/config/#{file}"
+      end
+
+      def config_from_file
+        file_name = config_file_name
+        puts "XBar::Mapper, reading configuration from file #{file_name}"
+        if File.exists? file_name
+          config = JSON.parse(ERB.new(File.read(file_name)).result)
+        else
+          config = {}
+        end
+        HashWithIndifferentAccess.new(config)
+      end
+      
+      def config
+       @@cached_config ||= config_from_file
+      end 
+      
+      # Alter the configuration in-memory for the current XBar envirnoment.
+      def shards=(shards)
+        cached_config["environments"][app_env] = shards
+      end
+       
+      # This needs to be reconciled with the 'environments' method in the
+      # XBar module.  That method specifies the environments that XBar should
+      # be enabled for.  The present method returns the environments that
+      # the current config file contains. XXX
+      def environments
+        config['environments'].keys
+      end
+       
+      #
+      # When we switch the XBar env or the Rails env (either of which
+      # changes the set of available shards, we have to find all the
+      # connection proxies and reset their current shard to :master.)
+      #
+      # Q1.  Are all the connection proxies pointed to by model classes
+      # findable through Thread.current[:connection_proxy]?  We'll have
+      # to loop over all threads. XXX
+      #
+      # Alternatively, we can register each XBar::Proxy.new call to
+      # a hash in the XBar module.
+      #
+      def reset(options = {})
+        new_xbar_env = options[:xbar_env] || xbar_env
+        if (new_xbar_env != xbar_env) || (options[:clear_cache]) ||
+          (!@@cached_config.nil? && @@cached_config.empty?)
+          @@cached_config = nil
+        end
+        self.xbar_env = new_xbar_env
+        self.app_env = options[:app_env] if options[:app_env]
+            
+        # puts "XBar::Mapper#reset, xbar_env=#{xbar_env}, app_env=#{app_env}"
+        initialize_shards(config)
+        initialize_options(config)
+        
+        # If Rails or some other entity has not assigned a native connection
+        # for ActiveRecord, we will try to do something sensible.  This is only
+        # needed if you have some enviroments for which XBar is not enabled. 
+        # However, it's not likely you'll want to enable XBar for only some
+        # environments.  (What would be a use case?)  The first 
+        # choice is that if we have a shard called 'master', we will use its
+        # connection specification.  The second choice is to include a Ruby
+        # file that contains a call to 'establish connection'.  In this case,
+        # we will create a shard called master with the same connection 
+        # specification.  Thus there will always be a 'master' shard.
+        #
+        # Also, there is the case where there is a connection, but the config
+        # document didn't specify a master shard.  
+        
+        begin
+          connection_pool = ActiveRecord::Base.connection_pool_without_xbar
+        rescue
+          if @@shards.keys.include? "master"
+             ActiveRecord::Base.establish_connection(
+               XBar::Mapper.shards[:master][0].spec.config)
+          else
+            # The config file didn't exist or didn't specify a master shard. Or
+            # app_env wasn't specified (as an argument option).
+            require connection_file_name
+            connection_pool = ActiveRecord::Base.connection_pool_without_xbar
+          end  
+        end
+        if !@@shards.keys.include?("master") && connection_pool
+          @@shards[:master] = Array(connection_pool)
+          @@adapters << connection_pool.spec.config
+        end
+        
+        @@proxies.each do |proxy|
+          proxy.reset_proxy
+        end
+        self
+      end
+      
+      def initialize_options(aconfig)
+        @@options = aconfig["environments"][app_env].dup
+        @@options.delete("shards")
+      rescue
+        @@options = {}
+      ensure
+        @@options[:verify_connection] ||= false
+      end
+      
+      def register(proxy)
+        @@proxies << proxy
+        
+        # If we hang on to a reference to proxies here, the proxy will
+        # never be garbage collected, even when the thread that it was
+        # assigned to goes away.  Find a way to fix this. XXX  
+      end
+      
+      def app_env
+        @@app_env = XBar.rails_env || @@app_env
+      end
+
+      def xbar_env
+        @env ||= 'default'
+      end
+      
+      private
+      
+      def app_env=(env)
+        if XBar.rails_env && XBar.rails_env != env
+          raise XBar::ConfigError, "Can't change app_env when you have a Rails environment."
+        end
+        @@app_env = env
+      end 
+
+      # When XBar::Mapper is processing a reset, it will call this method.  No other
+      # method should call this.
+      def xbar_env=(xbar_env)
+        @env = xbar_env
+      end
+      
+      def initialize_shards(aconfig)
+       
+        @@connections.clear
+        @@adapters.clear
+        @@shards.clear
+        
+        if aconfig
+          begin
+            shards_config = aconfig["environments"][app_env]["shards"]
+          rescue
+            shards_config = nil
+          end
+        end
+        shards_config ||= []
+        shards_config.delete_if {|k| k == "__COMMENT"}
+        
+        shards_config.each do |shard_key, connection_key|
+          if @@shards.include? shard_key 
+            raise ConfigError, "You have duplicate shard names!"
+          end
+          if connection_key.kind_of? String
+            spec = aconfig["connections"][connection_key]
+            pool = install_connection(connection_key, spec)
+            @@shards[shard_key] = [pool]
+          else # an array of connection keys
+            @@shards[shard_key] = []
+            connection_key.each do |conn_key|
+              spec = aconfig["connections"][conn_key]
+              pool = install_connection(conn_key, spec)
+              @@shards[shard_key] << pool
+            end
+          end
+        end
+        # @@shards[:master] = [@@connections[:master]]
+      end
+      
+      # Should return a ConnectionPool.
+      def install_connection(conn_key, spec)
+        unless spec
+          raise XBar::ConfigError, "No connection for key #{conn_key}"
+        end
+        if defined? ActiveRecord::Base::ConnectionSpecification::Resolver
+          resolver = ActiveRecord::Base::ConnectionSpecification::Resolver.new(spec, {})
+          spec = resolver.spec
+          @@adapters << spec.config[:adapter]
+          @@connections[conn_key.to_sym] = 
+            ActiveRecord::ConnectionAdapters::ConnectionPool.new(spec)
+        else
+          old_install_connection(conn_key, spec)
+        end
+      end
+      
+      def old_install_connection(conn_key, spec)
+        unless spec
+           raise XBar::ConfigError, "No connection for key #{conn_key}"
+         end
+        install_adapter(spec['adapter'])
+        @@connections[conn_key.to_sym] =
+          connection_pool_for(spec, "#{spec['adapter']}_connection")  
+      end
+      
+      # Called only from 'old_install_connection'.  If you get here, you should not
+      # be using connection URI's.
+      def install_adapter(adapter)
+        @@adapters << adapter
+        begin
+          require "active_record/connection_adapters/#{adapter}_adapter"
+        rescue LoadError
+          raise "Please install the #{adapter} adapter: " +
+            "`gem install activerecord-#{adapter}-adapter` (#{$!})"
+         end
+      end
+      
+      def connection_pool_for(adapter, spec)
+        ActiveRecord::ConnectionAdapters::ConnectionPool.new(
+          ActiveRecord::Base::ConnectionSpecification.new(adapter, spec))
+      end
+    end
+    
+    # Give module XBar::Mapper the above class methods.
+     self.extend(ClassMethods)
+    
+    Mapper.exports.each do |meth|
+      define_method(meth) {Mapper.send(meth)}
+    end
+    
+    def reset_config(options = {})
+      Mapper.reset(options)
+    end
+    
+    def register
+      Mapper.register(self)
+    end
+  
+  end
+  
+end

data/lib/xbar/migration.rb

@@ -0,0 +1,76 @@
+module XBar::Migration
+  
+  def self.extended(base)
+    class << base
+      def announce_with_xbar(message)
+        announce_without_xbar("#{message} - #{get_current_shard}")
+      end
+      alias_method_chain :migrate, :xbar
+      alias_method_chain :announce, :xbar
+      attr_accessor :current_shard
+    end
+  end
+
+  def self.included(base)
+    base.class_eval do
+      def announce_with_xbar(message)
+        announce_without_xbar("#{message} - #{get_current_shard}")
+      end
+      alias_method_chain :migrate, :xbar
+      alias_method_chain :announce, :xbar
+      attr_accessor :current_shard
+    end
+    base.extend(ClassMethods)
+  end
+  
+  module ClassMethods
+    
+    def using(*args)
+      if self.connection.is_a?(XBar::Proxy)
+        # Doesn't it make sense to only keep the schema_migrations table on the
+        # master shard?  If we create these other tables, they are unused anyway.
+        #args.each do |shard|
+        #  self.connection.check_schema_migrations(shard)
+        #end
+        @current_shard = *args
+        self.connection.enter_block_scope 
+        self.current_shard = args
+        self.connection.current_shard = args
+      end
+      return self
+    end
+    
+  end
+  
+  def migrate_with_xbar(direction) 
+    conn = ActiveRecord::Base.connection
+    raise "XBar::Migration#mismatched connections" unless conn == self.connection
+    if conn.kind_of?(XBar::Proxy)
+      u2 = self.class.instance_variable_get(:@current_shard)
+      conn.current_shard = u2 if u2
+      conn.send_queries_to_multiple_shards(conn.current_shard) do
+        migrate_without_xbar(direction)
+      end
+    else
+      migrate_without_xbar(direction)
+    end
+  ensure
+    if conn.kind_of?(XBar::Proxy)
+      conn.clean_proxy
+    end
+  end
+
+  # Used by migration when printing out results.
+  def get_current_shard
+    if ActiveRecord::Base.connection.respond_to?(:current_shard)
+      "Shard: #{ActiveRecord::Base.connection.current_shard}" 
+    end
+  end
+
+end
+
+if XBar.rails31?
+  ActiveRecord::Migration.send :include, XBar::Migration
+else
+  ActiveRecord::Migration.extend(XBar::Migration)
+end

data/lib/xbar/model.rb

@@ -0,0 +1,165 @@
+require 'active_support/concern'
+
+module XBar::Model
+  
+  extend ActiveSupport::Concern
+  
+  included do   
+    attr_accessor :current_shard
+    before_save :reload_connection
+    if XBar.rails3?
+      after_initialize :set_current_shard
+    else
+      def after_initialize
+        set_current_shard
+      end
+    end
+    class << self
+      alias_method_chain :connection, :xbar
+      alias_method_chain :connection_pool, :xbar
+    end
+  end
+  
+  def should_set_current_shard?
+    current_shard
+  end
+  
+  def connection_proxy
+    self.class.connection_proxy
+  end
+  
+  def reload_connection_safe
+    return yield unless should_set_current_shard?
+    original = connection_proxy.current_shard
+    connection_proxy.current_shard = current_shard
+    result = yield
+    connection_proxy.current_shard = original
+    result
+  end
+
+  def reload_connection
+    return unless should_set_current_shard?
+    connection_proxy.current_shard = current_shard
+  end
+  
+  private
+  
+  # After initialize callback.  
+  def set_current_shard
+    if new_record? || connection_proxy.in_block_scope?
+      if XBar.debug
+        type = new_record? ? "New" : "Existing"
+        puts "#{type} model callback, current_shard=#{connection_proxy.current_shard}, " +
+          "block_scope=#{connection_proxy.in_block_scope?}"
+      end
+      self.current_shard = connection_proxy.current_shard
+    else
+      if XBar.debug
+        type = new_record? ? "New" : "Existing"
+        puts "#{type} model callback, current_shard=#{connection_proxy.current_shard} " +
+          "last_current_shard=#{connection_proxy.last_current_shard}, " +
+          "block_scope=#{connection_proxy.in_block_scope?}"
+      end
+      self.current_shard = connection_proxy.last_current_shard
+    end
+  end
+  
+  module ClassMethods
+  
+    def should_use_normal_connection?
+      (defined?(Rails) && XBar.config &&
+        !XBar.environments.include?(Rails.env.to_s)) || 
+        (if XBar.rails32?
+           _establish_connection
+         else
+           self.read_inheritable_attribute(:_establish_connection)
+         end
+        )
+    end
+
+    def connection_proxy
+      puts "Model allocating new connection proxy"  unless Thread.current[:connection_proxy]
+      Thread.current[:connection_proxy] ||= XBar::Proxy.new
+    end
+
+    def connection_with_xbar
+      if should_use_normal_connection?
+        connection_without_xbar
+      else
+        #puts "Model connection with octopus" if XBar.debug
+        #if (connection_proxy.current_model.nil?) || (self !=  ActiveRecord::Base)
+          connection_proxy.current_model = self  
+        #end
+        connection_proxy
+      end
+    end
+
+    def connection_pool_with_xbar
+      if should_use_normal_connection?
+        connection_pool_without_xbar
+      else 
+        connection_proxy.connection_pool
+      end
+    end
+    
+    def clean_table_name
+      return unless connection_proxy.should_clean_table_name?
+      if self != ActiveRecord::Base && self.respond_to?(:reset_table_name) &&
+        (if XBar.rails32?
+           !self._reset_table_name
+         else
+           !self.read_inheritable_attribute(:_reset_table_name)
+         end
+        )
+        self.reset_table_name
+      end
+
+      if XBar.rails3?
+        self.reset_column_information
+        self.instance_variable_set(:@quoted_table_name, nil)
+      end
+    end
+    
+    def using(shard_name)
+      return self if defined?(::Rails) && !XBar.environments.include?(Rails.env.to_s)
+      clean_table_name
+      return XBar::ScopeProxy.new(shard_name, self)
+    end
+
+    def unreplicated_model
+      if XBar.rails32?
+        self._unreplicated = true
+      else
+        write_inheritable_attribute(:_unreplicated, true)
+      end
+    end
+    
+    def unreplicated_model?
+      if XBar.rails32?
+         _unreplicated
+       else
+         read_inheritable_attribute(:_unreplicated)
+      end
+    end
+        
+    def xbar_establish_connection(spec = nil)
+      if XBar.rails32?
+        self._establish_connection = true
+      else
+        write_inheritable_attribute(:_establish_connection, true)
+      end
+      establish_connection(spec)
+    end
+
+    def xbar_set_table_name(value = nil)
+      if XBar.rails32?
+        self._reset_table_name = true
+        self.table_name = value
+      else
+        write_inheritable_attribute(:_reset_table_name, true)
+        set_table_name(value)
+      end
+    end
+  end
+end
+