ar-octopus-ruby-3 0.11.2

data/lib/octopus/query_cache_for_shards.rb

@@ -0,0 +1,24 @@
+# query cache methods are moved to ConnectionPool for Rails >= 5.0
+module Octopus
+  module ConnectionPool
+    module QueryCacheForShards
+      %i(enable_query_cache! disable_query_cache!).each do |method|
+        define_method(method) do
+          if(Octopus.enabled? && (shards = ActiveRecord::Base.connection.shards)['master'] == self)
+            shards.each do |shard_name, v|
+              if shard_name == 'master'
+                super()
+              else
+                v.public_send(method)
+              end
+            end
+          else
+            super()
+          end
+        end
+      end
+    end
+  end
+end
+
+ActiveRecord::ConnectionAdapters::ConnectionPool.send(:prepend, Octopus::ConnectionPool::QueryCacheForShards)

data/lib/octopus/railtie.rb

@@ -0,0 +1,11 @@
+begin
+  module Octopus
+    class Railtie < Rails::Railtie
+      rake_tasks do
+        Dir[File.join(File.dirname(__FILE__), '../tasks/*.rake')].each { |ext| load ext }
+      end
+    end
+  end
+rescue LoadError
+  # nop
+end

data/lib/octopus/relation_proxy.rb

@@ -0,0 +1,74 @@
+module Octopus
+  class RelationProxy < BasicObject
+    include ::Octopus::ShardTracking::Attribute
+
+    module CaseFixer
+      def ===(other)
+        other = other.ar_relation while ::Octopus::RelationProxy === other
+        super
+      end
+    end
+
+    attr_accessor :ar_relation
+
+    def initialize(shard, ar_relation)
+      @current_shard = shard
+      @ar_relation = ar_relation
+    end
+
+    def respond_to?(*args)
+      method_missing(:respond_to?, *args)
+    end
+
+    # methods redefined in ActiveRecord that should run_on_shard
+    ENUM_METHODS = (::Enumerable.instance_methods - ::Object.instance_methods).reject do |m|
+      ::ActiveRecord::Relation.instance_method(m).source_location rescue nil
+    end + [:each, :map, :index_by]
+    # `find { ... }` etc. should run_on_shard, `find(id)` should be sent to relation
+    ENUM_WITH_BLOCK_METHODS = [:find, :select, :none?, :any?, :one?, :many?, :sum]
+    BATCH_METHODS = [:find_each, :find_in_batches, :in_batches]
+    WHERE_CHAIN_METHODS = [:not]
+
+    def method_missing(method, *args, &block)
+      if !block && BATCH_METHODS.include?(method)
+        ::Enumerator.new do |yielder|
+          run_on_shard do
+            parsed_args = args.empty? ? {} : args.first
+
+            @ar_relation.public_send(method, **parsed_args) do |batch_item|
+              yielder << batch_item
+            end
+          end
+        end
+      elsif ENUM_METHODS.include?(method) || block && ENUM_WITH_BLOCK_METHODS.include?(method)
+        run_on_shard { @ar_relation.to_a }.public_send(method, *args, &block)
+      elsif WHERE_CHAIN_METHODS.include?(method)
+        ::Octopus::ScopeProxy.new(@current_shard, run_on_shard { @ar_relation.public_send(method, *args) } )
+      elsif block
+        parsed_args = args.empty? ? {} : args.first
+
+        @ar_relation.public_send(method, **parsed_args, &block)
+      else
+        run_on_shard do
+          if method == :load_records
+            @ar_relation.send(method, *args)
+          else
+            @ar_relation.public_send(method, *args)
+          end
+        end
+      end
+    end
+
+    def ==(other)
+      case other
+      when ::Octopus::RelationProxy
+        method_missing(:==, other.ar_relation)
+      else
+        method_missing(:==, other)
+      end
+    end
+    alias_method :eql?, :==
+  end
+end
+
+ActiveRecord::Relation.extend(Octopus::RelationProxy::CaseFixer)

data/lib/octopus/result_patch.rb

@@ -0,0 +1,19 @@
+module Octopus::ResultPatch
+  attr_accessor :current_shard
+
+  private
+
+  def hash_rows
+    if current_shard.blank?
+      super
+    else
+      foo = super
+      foo.each { |f| f.merge!('current_shard' => current_shard) }
+      foo
+    end
+  end
+end
+
+class ActiveRecord::Result
+  prepend Octopus::ResultPatch
+end

data/lib/octopus/scope_proxy.rb

@@ -0,0 +1,68 @@
+module Octopus
+  class ScopeProxy < BasicObject
+    include ::Octopus::ShardTracking::Attribute
+
+    module CaseFixer
+      def ===(other)
+        other = other.klass while ::Octopus::ScopeProxy === other
+        super
+      end
+    end
+
+    attr_accessor :klass
+
+    # Dup and clone should be delegated to the class.
+    # We want to dup the query, not the scope proxy.
+    delegate :dup, :clone, to: :klass
+
+    def initialize(shard, klass)
+      @current_shard = shard
+      @klass = klass
+    end
+
+    def using(shard)
+      fail "Nonexistent Shard Name: #{shard}" if @klass.connection.shards[shard].nil?
+      @current_shard = shard
+      self
+    end
+
+    # Transaction Method send all queries to a specified shard.
+    def transaction(options = {}, &block)
+      run_on_shard { klass.transaction(**options, &block) }
+    end
+
+    def connection
+      @klass.connection_proxy.current_shard = @current_shard
+
+      if @klass.custom_octopus_connection && @klass.allowed_shard?(@current_shard)
+        # Force use of proxy, given we called 'using' explicitly to get here
+        @klass.connection_proxy.current_model = @klass
+        @klass.connection_proxy
+      else
+        @klass.connection
+      end
+    end
+
+    def method_missing(method, *args, &block)
+      result = run_on_shard { @klass.__send__(method, *args, &block) }
+      if result.respond_to?(:all)
+        return ::Octopus::ScopeProxy.new(current_shard, result)
+      end
+
+      if result.respond_to?(:current_shard)
+        result.current_shard = current_shard
+      end
+
+      result
+    end
+
+    # Delegates to method_missing (instead of @klass) so that User.using(:blah).where(:name => "Mike")
+    # gets run in the correct shard context when #== is evaluated.
+    def ==(other)
+      method_missing(:==, other)
+    end
+    alias_method :eql?, :==
+  end
+end
+
+ActiveRecord::Relation.extend(Octopus::ScopeProxy::CaseFixer)

data/lib/octopus/shard_tracking/attribute.rb

@@ -0,0 +1,22 @@
+# Adds current_shard as an attribute; provide a default
+# implementation of set_current_shard which considers
+# only the current ActiveRecord::Base.connection_proxy
+module Octopus
+  module ShardTracking
+    module Attribute
+      def self.included(base)
+        base.send(:include, Octopus::ShardTracking)
+      end
+
+      attr_accessor :current_shard
+
+      def set_current_shard
+        return unless Octopus.enabled?
+
+        if ActiveRecord::Base.connection_proxy.block
+          self.current_shard = ActiveRecord::Base.connection_proxy.current_shard
+        end
+      end
+    end
+  end
+end

data/lib/octopus/shard_tracking/dynamic.rb

@@ -0,0 +1,11 @@
+require 'octopus/shard_tracking'
+
+module Octopus
+  module ShardTracking
+    module Dynamic
+      def self.included(base)
+        base.send(:include, Octopus::ShardTracking)
+      end
+    end
+  end
+end

data/lib/octopus/shard_tracking.rb

@@ -0,0 +1,46 @@
+module Octopus
+  module ShardTracking
+    def self.included(base)
+      base.extend(ClassMethods)
+    end
+
+    module ClassMethods
+      # If the class which includes this module responds to the class
+      # method sharded_methods, then automagically alias_method_chain
+      # a sharding-friendly version of each of those methods into existence
+      def sharded_methods(*methods)
+        methods.each { |m| create_sharded_method(m) }
+      end
+
+      def create_sharded_method(name)
+        name.to_s =~ /([^!?]+)([!?])?/
+        method, punctuation = [Regexp.last_match[1], Regexp.last_match[2]]
+        with = :"#{method}_with_octopus#{punctuation}"
+        without = :"#{method}_without_octopus#{punctuation}"
+        define_method with do |*args, &block|
+          run_on_shard { send(without, *args, &block) }
+        end
+        alias_method without.to_sym, name.to_sym
+        alias_method name.to_sym, with.to_sym
+      end
+    end
+
+    # Adds run_on_shard method, but does not implement current_shard method
+    def run_on_shard(&block)
+      if (cs = current_shard)
+        r = ActiveRecord::Base.connection_proxy.run_queries_on_shard(cs, &block)
+        # Use a case statement to avoid any path through ActiveRecord::Delegation's
+        # respond_to? code. We want to avoid the respond_to? code because it can have
+        # the side effect of causing a call to load_target
+
+        if (ActiveRecord::Relation === r || ActiveRecord::QueryMethods::WhereChain === r) && !(Octopus::RelationProxy === r)
+          Octopus::RelationProxy.new(cs, r)
+        else
+          r
+        end
+      else
+        yield
+      end
+    end
+  end
+end

data/lib/octopus/singular_association.rb

@@ -0,0 +1,9 @@
+module Octopus
+  module SingularAssociation
+    def self.included(base)
+      base.sharded_methods :reader, :writer, :create, :create!, :build
+    end
+  end
+end
+
+ActiveRecord::Associations::SingularAssociation.send(:include, Octopus::SingularAssociation)

data/lib/octopus/slave_group.rb

@@ -0,0 +1,13 @@
+module Octopus
+  class SlaveGroup
+    def initialize(slaves)
+      slaves = HashWithIndifferentAccess.new(slaves)
+      slaves_list = slaves.values
+      @load_balancer = Octopus.load_balancer.new(slaves_list)
+    end
+
+    def next(options)
+      @load_balancer.next options
+    end
+  end
+end

data/lib/octopus/version.rb

@@ -0,0 +1,3 @@
+module Octopus
+  VERSION = '0.11.2'
+end

data/lib/octopus.rb

@@ -0,0 +1,209 @@
+require 'active_record'
+require 'active_support/version'
+require 'active_support/core_ext/class'
+require 'byebug'
+
+require 'yaml'
+require 'erb'
+
+module Octopus
+  def self.env
+    @env ||= 'octopus'
+  end
+
+  def self.rails_env
+    @rails_env ||= defined?(::Rails.env) ? Rails.env.to_s : 'shards'
+  end
+
+  def self.config
+    @config ||= begin
+      file_name = File.join(Octopus.directory, 'config/shards.yml').to_s
+
+      if File.exist?(file_name) || File.symlink?(file_name)
+        config ||= HashWithIndifferentAccess.new(YAML.load(ERB.new(File.read(file_name)).result))[Octopus.env]
+      else
+        config ||= HashWithIndifferentAccess.new
+      end
+
+      config
+    end
+  end
+
+  def self.load_balancer=(balancer)
+    @load_balancer = balancer
+  end
+
+  def self.load_balancer
+    @load_balancer ||= Octopus::LoadBalancing::RoundRobin
+  end
+
+  def self.master_shard
+    ((config && config[:master_shard]) || :master).to_sym
+  end
+
+  # Public: Whether or not Octopus is configured and should hook into the
+  # current environment. Checks the environments config option for the Rails
+  # environment by default.
+  #
+  # Returns a boolean
+  def self.enabled?
+    if defined?(::Rails.env)
+      Octopus.environments.include?(Rails.env.to_s)
+    else
+      # TODO: This doens't feel right but !Octopus.config.blank? is breaking a
+      #       test. Also, Octopus.config is always returning a hash.
+      Octopus.config
+    end
+  end
+
+  # Returns the Rails.root_to_s when you are using rails
+  # Running the current directory in a generic Ruby process
+  def self.directory
+    @directory ||= defined?(::Rails.root) ? Rails.root.to_s : Dir.pwd
+  end
+
+  # This is the default way to do Octopus Setup
+  # Available variables:
+  # :enviroments => the enviroments that octopus will run. default: 'production'
+  def self.setup
+    yield self
+  end
+
+  def self.environments=(environments)
+    @environments = environments.map(&:to_s)
+  end
+
+  def self.environments
+    @environments ||= config['environments'] || ['production']
+  end
+
+  def self.robust_environments=(environments)
+    @robust_environments = environments.map(&:to_s)
+  end
+
+  # Environments in which to swallow failures from a single shard
+  # when iterating through all.
+  def self.robust_environments
+    @robust_environments ||= config['robust_environments'] || ['production']
+  end
+
+  def self.robust_environment?
+    robust_environments.include? rails_env
+  end
+
+  def self.rails4?
+    ActiveRecord::VERSION::MAJOR == 4
+  end
+
+  def self.rails42?
+    rails4? && ActiveRecord::VERSION::MINOR == 2
+  end
+
+  def self.rails50?
+    ActiveRecord::VERSION::MAJOR == 5 && ActiveRecord::VERSION::MINOR == 0
+  end
+
+  def self.atleast_rails50?
+    ActiveRecord::VERSION::MAJOR >= 5
+  end
+
+  def self.rails51?
+    ActiveRecord::VERSION::MAJOR == 5 && ActiveRecord::VERSION::MINOR == 1
+  end
+
+  def self.rails52?
+    ActiveRecord::VERSION::MAJOR == 5 && ActiveRecord::VERSION::MINOR == 2
+  end
+
+  def self.rails60?
+    ActiveRecord::VERSION::MAJOR > 6 || ActiveRecord::VERSION::MAJOR == 6
+  end
+
+  def self.atleast_rails51?
+    ActiveRecord::VERSION::MAJOR > 5 || (ActiveRecord::VERSION::MAJOR == 5 && ActiveRecord::VERSION::MINOR >= 1)
+  end
+
+  def self.atleast_rails52?
+    ActiveRecord::VERSION::MAJOR > 5 || (ActiveRecord::VERSION::MAJOR == 5 && ActiveRecord::VERSION::MINOR > 1)
+  end
+
+  attr_writer :logger
+
+  def self.logger
+    if defined?(Rails.logger)
+      @logger ||= Rails.logger
+    else
+      @logger ||= Logger.new($stderr)
+    end
+  end
+
+  def self.shards=(shards)
+    config[rails_env] = HashWithIndifferentAccess.new(shards)
+    ActiveRecord::Base.connection.initialize_shards(@config)
+  end
+
+  def self.using(shard, &block)
+    conn = ActiveRecord::Base.connection
+
+    if conn.is_a?(Octopus::Proxy)
+      conn.run_queries_on_shard(shard, &block)
+    else
+      yield
+    end
+  end
+
+  def self.using_group(group, &block)
+    conn = ActiveRecord::Base.connection
+
+    if conn.is_a?(Octopus::Proxy)
+      conn.send_queries_to_group(group, &block)
+    else
+      yield
+    end
+  end
+
+  def self.using_all(&block)
+    conn = ActiveRecord::Base.connection
+
+    if conn.is_a?(Octopus::Proxy)
+      conn.send_queries_to_all_shards(&block)
+    else
+      yield
+    end
+  end
+
+  def self.fully_replicated(&_block)
+    old_fully_replicated = Thread.current[Octopus::ProxyConfig::FULLY_REPLICATED_KEY]
+    Thread.current[Octopus::ProxyConfig::FULLY_REPLICATED_KEY] = true
+    yield
+  ensure
+    Thread.current[Octopus::ProxyConfig::FULLY_REPLICATED_KEY] = old_fully_replicated
+  end
+end
+
+require 'octopus/exception'
+
+require 'octopus/shard_tracking'
+require 'octopus/shard_tracking/attribute'
+require 'octopus/shard_tracking/dynamic'
+
+require 'octopus/model'
+require 'octopus/result_patch'
+require 'octopus/migration'
+require 'octopus/association'
+require 'octopus/collection_association'
+require 'octopus/association_shard_tracking'
+require 'octopus/persistence'
+require 'octopus/log_subscriber'
+require 'octopus/abstract_adapter'
+require 'octopus/singular_association'
+require 'octopus/finder_methods'
+require 'octopus/query_cache_for_shards' unless Octopus.rails4?
+
+require 'octopus/railtie' if defined?(::Rails::Railtie)
+
+require 'octopus/proxy_config'
+require 'octopus/proxy'
+require 'octopus/collection_proxy'
+require 'octopus/relation_proxy'
+require 'octopus/scope_proxy'

data/lib/tasks/octopus.rake

@@ -0,0 +1,16 @@
+namespace :octopus do
+  desc 'Copy schema version information from master to all shards'
+  task :copy_schema_versions => :environment do
+    abort('Octopus is not enabled for this environment') unless Octopus.enabled?
+
+    connection = ActiveRecord::Base.connection
+
+    current_version  = ActiveRecord::Migrator.current_version
+    migrations_paths = ActiveRecord::Migrator.migrations_paths
+
+    connection.send_queries_to_multiple_shards(connection.shard_names) do
+      ActiveRecord::Schema.initialize_schema_migrations_table
+      ActiveRecord::Schema.assume_migrated_upto_version(current_version, migrations_paths)
+    end
+  end
+end

data/sample_app/.gitignore

@@ -0,0 +1,4 @@
+.bundle
+db/*.sqlite3
+log/*.log
+tmp/**/*

data/sample_app/.rspec

@@ -0,0 +1 @@
+--colour

data/sample_app/Gemfile

@@ -0,0 +1,20 @@
+source 'https://rubygems.org'
+
+gem 'rails', '3.2.13'
+
+# Bundle edge Rails instead:
+# gem 'rails', :git => 'git://github.com/rails/rails.git'
+
+gem 'pg'
+gem 'sqlite3'
+gem 'ar-octopus', '0.6.0'
+gem 'pry'
+
+group :test do
+  gem 'capybara'
+  gem 'cucumber-rails', :require => false
+  gem 'launchy'
+  gem 'rspec-rails'
+  gem 'aruba'
+  gem 'database_cleaner'
+end