ar-octopus-ruby-3 0.11.2

data/lib/ar-octopus.rb

@@ -0,0 +1 @@
+require 'octopus'

data/lib/octopus/abstract_adapter.rb

@@ -0,0 +1,33 @@
+# Implementation courtesy of db-charmer.
+module Octopus
+  module AbstractAdapter
+    module OctopusShard
+      class InstrumenterDecorator < ActiveSupport::ProxyObject
+        def initialize(adapter, instrumenter)
+          @adapter = adapter
+          @instrumenter = instrumenter
+        end
+
+        def instrument(name, payload = {}, &block)
+          payload[:octopus_shard] ||= @adapter.octopus_shard
+          @instrumenter.instrument(name, payload, &block)
+        end
+
+        def method_missing(meth, *args, &block)
+          @instrumenter.send(meth, *args, &block)
+        end
+      end
+
+      def octopus_shard
+        @config[:octopus_shard]
+      end
+
+      def initialize(*args)
+        super
+        @instrumenter = InstrumenterDecorator.new(self, @instrumenter)
+      end
+    end
+  end
+end
+
+ActiveRecord::ConnectionAdapters::AbstractAdapter.send(:prepend, Octopus::AbstractAdapter::OctopusShard)

data/lib/octopus/association.rb

@@ -0,0 +1,14 @@
+module Octopus
+  module Association
+    def self.included(base)
+      base.send(:include, Octopus::ShardTracking::Dynamic)
+      base.sharded_methods :target_scope
+    end
+
+    def current_shard
+      owner.current_shard
+    end
+  end
+end
+
+ActiveRecord::Associations::Association.send(:include, Octopus::Association)

data/lib/octopus/association_shard_tracking.rb

@@ -0,0 +1,74 @@
+module Octopus
+  module AssociationShardTracking
+    class MismatchedShards < StandardError
+      attr_reader :record, :current_shard
+
+      def initialize(record, current_shard)
+        @record = record
+        @current_shard = current_shard
+      end
+
+      def message
+        [
+          "Association Error: Records are from different shards",
+          "Record: #{record.inspect}",
+          "Current Shard: #{current_shard.inspect}",
+          "Current Record Shard: #{record.current_shard.inspect}",
+        ].join(" ")
+      end
+    end
+
+    def self.extended(base)
+      base.send(:include, InstanceMethods)
+    end
+
+    module InstanceMethods
+      def connection_on_association=(record)
+        return unless ::Octopus.enabled?
+        return if !self.class.connection.respond_to?(:current_shard) || !self.respond_to?(:current_shard)
+
+        if !record.current_shard.nil? && !current_shard.nil? && record.current_shard.to_s != current_shard.to_s
+          raise MismatchedShards.new(record, current_shard)
+        end
+
+        record.current_shard = self.class.connection.current_shard = current_shard if should_set_current_shard?
+      end
+    end
+
+    def has_many(name, scope = nil, **options, &extension)
+      if options == {} && scope.is_a?(Hash)
+        default_octopus_opts(scope)
+      else
+        default_octopus_opts(options)
+      end
+      super
+    end
+
+    def has_and_belongs_to_many(association_id, scope = nil, options = {}, &extension)
+      assign_octopus_opts(scope, options)
+      super
+    end
+
+    def default_octopus_opts(options)
+      options[:before_add] = [ :connection_on_association=, options[:before_add] ].compact.flatten
+      options[:before_remove] = [ :connection_on_association=, options[:before_remove] ].compact.flatten
+    end
+
+    def assign_octopus_opts(scope, options)
+      if options == {} && scope.is_a?(Hash)
+        default_octopus_opts(scope)
+      else
+        default_octopus_opts(options)
+      end
+    end
+
+    if Octopus.atleast_rails51?
+      def has_and_belongs_to_many(association_id, scope = nil, **options, &extension)
+        assign_octopus_opts(scope, options)
+        super
+      end
+    end
+  end
+end
+
+ActiveRecord::Base.extend(Octopus::AssociationShardTracking)

data/lib/octopus/collection_association.rb

@@ -0,0 +1,17 @@
+module Octopus
+  module CollectionAssociation
+    def self.included(base)
+      if Octopus.atleast_rails51?
+        base.sharded_methods :reader, :writer, :ids_reader, :ids_writer, :create, :create!,
+                             :build, :include?,
+                             :load_target, :reload, :size, :select
+      else
+        base.sharded_methods :reader, :writer, :ids_reader, :ids_writer, :create, :create!,
+                             :build, :any?, :count, :empty?, :first, :include?, :last, :length,
+                             :load_target, :many?, :reload, :size, :select, :uniq
+      end
+    end
+  end
+end
+
+ActiveRecord::Associations::CollectionAssociation.send(:include, Octopus::CollectionAssociation)

data/lib/octopus/collection_proxy.rb

@@ -0,0 +1,16 @@
+module Octopus
+  module CollectionProxy
+    def self.included(base)
+      base.send(:include, Octopus::ShardTracking::Dynamic)
+      base.sharded_methods :any?, :build, :count, :create, :create!, :concat, :delete, :delete_all,
+                           :destroy, :destroy_all, :empty?, :find, :first, :include?, :last, :length,
+                           :many?, :pluck, :replace, :select, :size, :sum, :to_a, :uniq
+    end
+
+    def current_shard
+      @association.owner.current_shard
+    end
+  end
+end
+
+ActiveRecord::Associations::CollectionProxy.send(:include, Octopus::CollectionProxy)

data/lib/octopus/exception.rb

@@ -0,0 +1,4 @@
+module Octopus
+  class Exception < ::Exception
+  end
+end

data/lib/octopus/finder_methods.rb

@@ -0,0 +1,8 @@
+# find_nth / find_nth! must be public here to allow Octopus to call
+# them on the scope proxy object
+if Octopus.rails42? || Octopus.rails50?
+  module ActiveRecord::FinderMethods
+    public :find_nth
+    public :find_nth!
+  end
+end

data/lib/octopus/load_balancing/round_robin.rb

@@ -0,0 +1,20 @@
+require 'octopus/load_balancing'
+
+# The round-robin load balancing of slaves belonging to the same shard.
+# It is a pool that contains slaves which queries are distributed to.
+module Octopus
+  module LoadBalancing
+    class RoundRobin
+      def initialize(slaves_list)
+        raise Octopus::Exception.new("No slaves available") if slaves_list.empty?
+        @slaves_list = slaves_list
+        @slave_index = 0
+      end
+
+      # Returns the next available slave in the pool
+      def next(options)
+        @slaves_list[@slave_index = (@slave_index + 1) % @slaves_list.length]
+      end
+    end
+  end
+end

data/lib/octopus/load_balancing.rb

@@ -0,0 +1,4 @@
+module Octopus
+  module LoadBalancing
+  end
+end

data/lib/octopus/log_subscriber.rb

@@ -0,0 +1,26 @@
+# Implementation courtesy of db-charmer.
+module Octopus
+  module LogSubscriber
+    def self.included(base)
+      base.send(:attr_accessor, :octopus_shard)
+
+      base.send :alias_method, :sql_without_octopus_shard, :sql
+      base.send :alias_method, :sql, :sql_with_octopus_shard
+
+      base.send :alias_method, :debug_without_octopus_shard, :debug
+      base.send :alias_method, :debug, :debug_with_octopus_shard
+    end
+
+    def sql_with_octopus_shard(event)
+      self.octopus_shard = event.payload[:octopus_shard]
+      sql_without_octopus_shard(event)
+    end
+
+    def debug_with_octopus_shard(msg)
+      conn = octopus_shard ? color("[Shard: #{octopus_shard}]", ActiveSupport::LogSubscriber::GREEN, true) : ''
+      debug_without_octopus_shard(conn + msg)
+    end
+  end
+end
+
+ActiveRecord::LogSubscriber.send(:include, Octopus::LogSubscriber)

data/lib/octopus/migration.rb

@@ -0,0 +1,236 @@
+require 'set'
+require 'active_support/core_ext/module/aliasing'
+require 'active_support/core_ext/array/wrap'
+
+module Octopus
+  module Migration
+    module InstanceOrClassMethods
+      def announce_with_octopus(message)
+        announce_without_octopus("#{message} - #{current_shard}")
+      end
+
+      def current_shard
+        "Shard: #{connection.current_shard}" if connection.respond_to?(:current_shard)
+      end
+    end
+
+    include InstanceOrClassMethods
+
+    def self.included(base)
+      base.extend(ClassMethods)
+
+      base.send :alias_method, :announce_without_octopus, :announce
+      base.send :alias_method, :announce, :announce_with_octopus
+
+      base.class_attribute :current_shard, :current_group, :current_group_specified, :instance_reader => false, :instance_writer => false
+    end
+
+    module ClassMethods
+      def using(*args)
+        return self unless connection.is_a?(Octopus::Proxy)
+
+        self.current_shard = args
+        self
+      end
+
+      def using_group(*groups)
+        return self unless connection.is_a?(Octopus::Proxy)
+
+        self.current_group = groups
+        self.current_group_specified = true
+        self
+      end
+
+      def shards
+        shards = Set.new
+
+        if (groups = (current_group_specified ? current_group : Octopus.config[:default_migration_group]))
+          Array.wrap(groups).each do |group|
+            group_shards = connection.shards_for_group(group)
+            shards.merge(group_shards) if group_shards
+          end
+        elsif (shard = current_shard)
+          shards.merge(Array.wrap(shard))
+        end
+
+        shards.to_a.presence || [Octopus.master_shard]
+      end
+    end
+  end
+end
+
+module Octopus
+  module Migrator
+    def self.included(base)
+      unless Octopus.atleast_rails52?
+        base.extend(ClassMethods)
+
+        base.class_eval do
+          class << self
+              alias_method :migrate_without_octopus, :migrate
+              alias_method :migrate, :migrate_with_octopus
+              
+              alias_method :up_without_octopus, :up
+              alias_method :up, :up_with_octopus
+
+              alias_method :down_without_octopus, :down
+              alias_method :down, :down_with_octopus
+
+              alias_method :run_without_octopus, :run
+              alias_method :run, :run_with_octopus
+
+              alias_method :rollback_without_octopus, :rollback
+              alias_method :rollback, :rollback_with_octopus
+          end
+        end
+      end
+
+      base.send :alias_method, :run_without_octopus, :run
+      base.send :alias_method, :run, :run_with_octopus
+
+      base.send :alias_method, :migrate_without_octopus, :migrate
+      base.send :alias_method, :migrate, :migrate_with_octopus
+
+      base.send :alias_method, :migrations_without_octopus, :migrations
+      base.send :alias_method, :migrations, :migrations_with_octopus
+    end
+    if Octopus.atleast_rails52?
+      ### Post RAILS 5.2 Migration methods
+      
+      def run_with_octopus(&block)
+        return run_without_octopus(&block) unless connection.is_a?(Octopus::Proxy)
+        shards = migrations.map(&:shards).flatten.map(&:to_s)
+        connection.send_queries_to_multiple_shards(shards) do
+          run_without_octopus(&block)
+        end
+        rescue ActiveRecord::UnknownMigrationVersionError => e
+          raise unless migrations(true).detect { |m| m.version == e.version }
+      end
+
+      def migrate_with_octopus(&block)
+        return migrate_without_octopus(&block) unless connection.is_a?(Octopus::Proxy)
+        shards = migrations.map(&:shards).flatten.map(&:to_s)
+        connection.send_queries_to_multiple_shards(shards) do
+          migrate_without_octopus(&block)
+        end
+        rescue ActiveRecord::UnknownMigrationVersionError => e
+          raise unless migrations(true).detect { |m| m.version == e.version }
+      end
+
+      def migrations_with_octopus(shard_agnostic = true)
+        migrations = migrations_without_octopus
+        return migrations if !connection.is_a?(Octopus::Proxy) || shard_agnostic
+
+        migrations.select { |m| m.shards.include?(connection.current_shard.to_sym) }
+      end
+      
+      private 
+      
+      def connection
+        ActiveRecord::Base.connection
+      end
+      
+    else 
+      ### Pre RAILS 5.2 Migration methods
+      
+      def run_with_octopus(&block)
+        run_without_octopus(&block)
+      rescue ActiveRecord::UnknownMigrationVersionError => e
+        raise unless migrations(true).detect { |m| m.version == e.version }
+      end
+
+      def migrate_with_octopus(&block)
+        migrate_without_octopus(&block)
+      rescue ActiveRecord::UnknownMigrationVersionError => e
+        raise unless migrations(true).detect { |m| m.version == e.version }
+      end
+
+      def migrations_with_octopus(shard_agnostic = false)
+        connection = ActiveRecord::Base.connection
+        migrations = migrations_without_octopus
+        return migrations if !connection.is_a?(Octopus::Proxy) || shard_agnostic
+
+        migrations.select { |m| m.shards.include?(connection.current_shard.to_sym) }
+      end
+
+      module ClassMethods
+        def migrate_with_octopus(migrations_paths, target_version = nil, &block)
+          return migrate_without_octopus(migrations_paths, target_version, &block) unless connection.is_a?(Octopus::Proxy)
+
+          connection.send_queries_to_multiple_shards(connection.shard_names) do
+            migrate_without_octopus(migrations_paths, target_version, &block)
+          end
+        end
+
+        def up_with_octopus(migrations_paths, target_version = nil, &block)
+          return up_without_octopus(migrations_paths, target_version, &block) unless connection.is_a?(Octopus::Proxy)
+          return up_without_octopus(migrations_paths, target_version, &block) unless connection.current_shard.to_s == Octopus.master_shard.to_s
+
+          connection.send_queries_to_multiple_shards(connection.shard_names) do
+            up_without_octopus(migrations_paths, target_version, &block)
+          end
+        end
+
+        def down_with_octopus(migrations_paths, target_version = nil, &block)
+          return down_without_octopus(migrations_paths, target_version, &block) unless connection.is_a?(Octopus::Proxy)
+          return down_without_octopus(migrations_paths, target_version, &block) unless connection.current_shard.to_s == Octopus.master_shard.to_s
+
+          connection.send_queries_to_multiple_shards(connection.shard_names) do
+            down_without_octopus(migrations_paths, target_version, &block)
+          end
+        end
+
+        def run_with_octopus(direction, migrations_paths, target_version)
+          return run_without_octopus(direction, migrations_paths, target_version) unless connection.is_a?(Octopus::Proxy)
+
+          connection.send_queries_to_multiple_shards(connection.shard_names) do
+            run_without_octopus(direction, migrations_paths, target_version)
+          end
+        end
+
+        def rollback_with_octopus(migrations_paths, steps = 1)
+          return rollback_without_octopus(migrations_paths, steps) unless connection.is_a?(Octopus::Proxy)
+
+          connection.send_queries_to_multiple_shards(connection.shard_names) do
+            rollback_without_octopus(migrations_paths, steps)
+          end
+        end
+
+        private
+
+        def connection
+          ActiveRecord::Base.connection
+        end
+      end
+    end
+  end
+end
+
+module Octopus
+  module MigrationProxy
+    def shards
+      migration.class.shards
+    end
+  end
+end
+
+module Octopus
+  module UnknownMigrationVersionError
+    def self.included(base)
+      base.send :alias_method, :initialize_without_octopus, :initialize
+      base.send :alias_method, :initialize, :initialize_with_octopus
+
+      base.send(:attr_accessor, :version)
+    end
+
+    def initialize_with_octopus(version)
+      @version = version
+      initialize_without_octopus(version)
+    end
+  end
+end
+
+ActiveRecord::Migration.send(:include, Octopus::Migration)
+ActiveRecord::Migrator.send(:include, Octopus::Migrator)
+ActiveRecord::MigrationProxy.send(:include, Octopus::MigrationProxy)
+ActiveRecord::UnknownMigrationVersionError.send(:include, Octopus::UnknownMigrationVersionError)

data/lib/octopus/model.rb

@@ -0,0 +1,216 @@
+require 'active_support/deprecation'
+
+module Octopus
+  module Model
+    def self.extended(base)
+      base.send(:include, Octopus::ShardTracking::Attribute)
+      base.send(:include, InstanceMethods)
+      base.extend(ClassMethods)
+    end
+
+    module SharedMethods
+      def using(shard)
+        if block_given?
+          raise Octopus::Exception, <<-EOF
+#{name}.using is not allowed to receive a block, it works just like a regular scope.
+
+If you are trying to scope everything to a specific shard, use Octopus.using instead.
+          EOF
+        end
+
+        if Octopus.enabled?
+          Octopus::ScopeProxy.new(shard, self)
+        else
+          self
+        end
+      end
+    end
+
+    module InstanceMethods
+      include SharedMethods
+
+      def self.included(base)
+        base.send(:alias_method, :equality_without_octopus, :==)
+        base.send(:alias_method, :==, :equality_with_octopus)
+        base.send(:alias_method, :eql?, :==)
+        base.send(:alias_method, :perform_validations_without_octopus, :perform_validations)
+        base.send(:alias_method, :perform_validations, :perform_validations_with_octopus)
+      end
+
+      def set_current_shard
+        return unless Octopus.enabled?
+        shard = self.class.connection_proxy.current_shard
+        self.current_shard = shard if self.class.allowed_shard?(shard)
+      end
+
+      def init_with(coder)
+        obj = super
+
+        return obj unless Octopus.enabled?
+        return obj if obj.class.connection_proxy.current_model_replicated?
+
+        current_shard_value = coder['attributes']['current_shard'].value if coder['attributes']['current_shard'].present? && coder['attributes']['current_shard'].value.present?
+
+        coder['attributes'].send(:attributes).send(:values).delete('current_shard')
+        coder['attributes'].send(:attributes).send(:delegate_hash).delete('current_shard')
+
+        obj.current_shard = current_shard_value if current_shard_value.present?
+        obj
+      end
+
+      def should_set_current_shard?
+        self.respond_to?(:current_shard) && !current_shard.nil?
+      end
+
+      def equality_with_octopus(comparison_object)
+        equality_without_octopus(comparison_object) && comparison_object.current_shard.to_s == current_shard.to_s
+      end
+
+      def perform_validations_with_octopus(*args)
+        if Octopus.enabled? && should_set_current_shard?
+          Octopus.using(current_shard) do
+            perform_validations_without_octopus(*args)
+          end
+        else
+          perform_validations_without_octopus(*args)
+        end
+      end
+    end
+
+    module ClassMethods
+      include SharedMethods
+
+      def self.extended(base)
+        base.class_attribute(:replicated)
+        base.class_attribute(:sharded)
+        base.class_attribute(:allowed_shards)
+        base.hijack_methods
+      end
+
+      def replicated_model
+        self.replicated = true
+      end
+
+      def sharded_model
+        self.sharded = true
+      end
+
+      def allow_shard(*shards)
+        self.allowed_shards ||= []
+        self.allowed_shards += shards
+      end
+
+      def hijack_methods
+        after_initialize :set_current_shard
+
+        around_save :run_on_shard, :unless => lambda { self.class.custom_octopus_connection }
+
+        class_attribute :custom_octopus_connection
+
+        class << self
+          attr_accessor :custom_octopus_table_name
+
+          alias_method :connection_without_octopus, :connection
+          alias_method :connection, :connection_with_octopus
+
+          alias_method :connection_pool_without_octopus, :connection_pool
+          alias_method :connection_pool, :connection_pool_with_octopus
+
+          alias_method :clear_all_connections_without_octopus!, :clear_all_connections!
+          alias_method :clear_all_connections!, :clear_all_connections_with_octopus!
+
+          alias_method :clear_active_connections_without_octopus!, :clear_active_connections!
+          alias_method :clear_active_connections!, :clear_active_connections_with_octopus!
+
+          alias_method :connected_without_octopus?, :connected?
+          alias_method :connected?, :connected_with_octopus?
+
+          def table_name=(value = nil)
+            self.custom_octopus_table_name = true
+            super
+          end
+        end
+      end
+
+      def connection_proxy
+        ActiveRecord::Base.class_variable_defined?(:@@connection_proxy) &&
+          ActiveRecord::Base.class_variable_get(:@@connection_proxy) ||
+          ActiveRecord::Base.class_variable_set(:@@connection_proxy, Octopus::Proxy.new)
+      end
+
+      def should_use_normal_connection?
+        if !Octopus.enabled?
+          true
+        elsif custom_octopus_connection
+          !connection_proxy.block || !allowed_shard?(connection_proxy.current_shard)
+        end
+      end
+
+      def allowed_shard?(shard)
+        if custom_octopus_connection
+          allowed_shards && shard && (allowed_shards.include?(shard.to_s) || allowed_shards.include?(shard.to_sym))
+        else
+          true
+        end
+      end
+
+      def connection_with_octopus
+        if should_use_normal_connection?
+          connection_without_octopus
+        else
+          connection_proxy.current_model = self
+          connection_proxy
+        end
+      end
+
+      def connection_pool_with_octopus
+        if should_use_normal_connection?
+          connection_pool_without_octopus
+        else
+          connection_proxy.connection_pool
+        end
+      end
+
+      def clear_active_connections_with_octopus!
+        if should_use_normal_connection?
+          clear_active_connections_without_octopus!
+        else
+          connection_proxy.clear_active_connections!
+        end
+      end
+
+      def clear_all_connections_with_octopus!
+        if should_use_normal_connection?
+          clear_all_connections_without_octopus!
+        else
+          connection_proxy.clear_all_connections!
+        end
+      end
+
+      def connected_with_octopus?
+        if should_use_normal_connection?
+          connected_without_octopus?
+        else
+          connection_proxy.connected?
+        end
+      end
+
+      def set_table_name_with_octopus(value = nil, &block)
+        self.custom_octopus_table_name = true
+        set_table_name_without_octopus(value, &block)
+      end
+
+      def octopus_establish_connection(spec = ENV['DATABASE_URL'])
+        self.custom_octopus_connection = true if spec
+        establish_connection(spec)
+      end
+
+      def octopus_set_table_name(value = nil)
+        ActiveSupport::Deprecation.warn 'Calling `octopus_set_table_name` is deprecated and will be removed in Octopus 1.0.', caller
+        set_table_name(value)
+      end
+    end
+  end
+end
+
+ActiveRecord::Base.extend(Octopus::Model)