Subclass switching

README.md

@@ -44,15 +44,13 @@ basically connections inherit configuration from the parent configuration file.
 ## Usage
 
 Normally you have some models that live on a shared database, and you might need to query this data in order to know what shard to switch to.
-All the models that live on the shared database must be marked as not\_sharded:
+All the models that live on the sharded database must inherit from ActiveRecordShards::ShardedModel:
 
     class Account < ActiveRecord::Base
-      not_sharded
-
       has_many :projects
     end
 
-    class Project < ActiveRecord::Base
+    class Project < ActiveRecordShards::ShardedModel
       belongs_to :account
     end
 

lib/active_record_shards.rb

@@ -3,6 +3,7 @@
 require 'active_record/base'
 require 'active_record_shards/configuration_parser'
 require 'active_record_shards/model'
+require 'active_record_shards/sharded_model'
 require 'active_record_shards/shard_selection'
 require 'active_record_shards/connection_switcher'
 require 'active_record_shards/association_collection_connection_selection'
@@ -24,5 +25,3 @@ def self.rails_env
 ActiveRecord::Relation.include(ActiveRecordShards::DefaultSlavePatches::ActiveRelationPatches)
 ActiveRecord::Associations::CollectionProxy.include(ActiveRecordShards::AssociationCollectionConnectionSelection)
 ActiveRecord::Associations::Builder::HasAndBelongsToMany.include(ActiveRecordShards::DefaultSlavePatches::HasAndBelongsToManyBuilderExtension)
-
-ActiveRecord::InternalMetadata.not_sharded

lib/active_record_shards/connection_switcher.rb

@@ -13,11 +13,6 @@ def self.extended(base)
       base.singleton_class.send(:alias_method, :table_exists?, :table_exists_with_default_shard?)
     end
 
-    def default_shard=(new_default_shard)
-      ActiveRecordShards::ShardSelection.default_shard = new_default_shard
-      switch_connection(shard: new_default_shard)
-    end
-
     def on_shard(shard)
       old_options = current_shard_selection.options
       switch_connection(shard: shard) if supports_sharding?
@@ -135,7 +130,7 @@ def on_slave?
     end
 
     def current_shard_selection
-      Thread.current[:shard_selection] ||= ShardSelection.new
+      Thread.current[:shard_selection] ||= ShardSelection.new(shard_names.first)
     end
 
     def current_shard_id
@@ -146,27 +141,33 @@ def shard_names
       unless config = configurations[shard_env]
         raise "Did not find #{shard_env} in configurations, did you forget to add it to your database config? (configurations: #{configurations.inspect})"
       end
-      unless config.fetch(SHARD_NAMES_CONFIG_KEY, []).all? { |shard_name| shard_name.is_a?(Integer) }
+      unless config[SHARD_NAMES_CONFIG_KEY]
+        raise "No shards configured for #{shard_env}"
+      end
+      unless config[SHARD_NAMES_CONFIG_KEY].all? { |shard_name| shard_name.is_a?(Integer) }
         raise "All shard names must be integers: #{config[SHARD_NAMES_CONFIG_KEY].inspect}."
       end
-      config[SHARD_NAMES_CONFIG_KEY] || []
+      config[SHARD_NAMES_CONFIG_KEY]
     end
 
     private
 
     def switch_connection(options)
       if options.any?
-        if options.key?(:slave)
-          current_shard_selection.on_slave = options[:slave]
-        end
-
         if options.key?(:shard)
-          unless configurations[shard_env]
+          unless config = configurations[shard_env]
             raise "Did not find #{shard_env} in configurations, did you forget to add it to your database config? (configurations: #{configurations.inspect})"
           end
+          unless config['shard_names'].include?(options[:shard])
+            raise "Did not find shard #{options[:shard]} in configurations"
+          end
           current_shard_selection.shard = options[:shard]
         end
 
+        if options.key?(:slave)
+          current_shard_selection.on_slave = options[:slave]
+        end
+
         ensure_shard_connection
       end
     end

lib/active_record_shards/default_slave_patches.rb

@@ -127,8 +127,8 @@ def model.on_slave_by_default?
           left_reflection.klass.on_slave_by_default?
         end
 
-        # also transfer the sharded-ness of the left table to the join model
-        model.not_sharded unless model.left_reflection.klass.is_sharded?
+        model.extend(ActiveRecordShards::Ext::ShardedModel) if model.left_reflection.klass.is_sharded?
+
         model
       end
     end

lib/active_record_shards/model.rb

@@ -3,24 +3,14 @@
 module ActiveRecordShards
   module Model
     def not_sharded
-      if self != ActiveRecord::Base && self != base_class
-        raise "You should only call not_sharded on direct descendants of ActiveRecord::Base"
-      end
-      self.sharded = false
+      ActiveSupport::Deprecation.warn("Calling not_sharded is deprecated. "\
+                                      "Please ensure to still read from the "\
+                                      "account db slave after removing the "\
+                                      "call.")
     end
 
     def is_sharded? # rubocop:disable Naming/PredicateName
-      if self == ActiveRecord::Base
-        sharded != false && supports_sharding?
-      elsif self == base_class
-        if sharded.nil?
-          ActiveRecord::Base.is_sharded?
-        else
-          sharded != false
-        end
-      else
-        base_class.is_sharded?
-      end
+      false
     end
 
     def on_slave_by_default?
@@ -61,9 +51,5 @@ def self.extended(base)
       base.send(:include, InstanceMethods)
       base.after_initialize :initialize_shard_and_slave
     end
-
-    private
-
-    attr_accessor :sharded
   end
 end

lib/active_record_shards/shard_selection.rb

@@ -1,20 +1,14 @@
 # frozen_string_literal: true
 module ActiveRecordShards
   class ShardSelection
-    NO_SHARD = :_no_shard
-    cattr_accessor :default_shard
-
-    def initialize
+    def initialize(shard)
       @on_slave = false
-      @shard = nil
+      self.shard = shard
     end
 
     def shard
-      if @shard.nil? || @shard == NO_SHARD
-        nil
-      else
-        @shard || self.class.default_shard
-      end
+      raise "Missing shard information on connection" unless @shard
+      @shard
     end
 
     PRIMARY = "primary".freeze
@@ -40,7 +34,7 @@ def resolve_connection_name(sharded:, configurations:)
     end
 
     def shard=(new_shard)
-      @shard = (new_shard || NO_SHARD)
+      @shard = Integer(new_shard)
     end
 
     def on_slave?

lib/active_record_shards/sharded_model.rb

@@ -0,0 +1,15 @@
+module ActiveRecordShards
+  module Ext
+    module ShardedModel
+      def is_sharded? # rubocop:disable Naming/PredicateName
+        true
+      end
+    end
+  end
+
+  class ShardedModel < ActiveRecord::Base
+    self.abstract_class = true
+
+    extend ActiveRecordShards::Ext::ShardedModel
+  end
+end

lib/active_record_shards/sql_comments.rb

@@ -10,7 +10,7 @@ def execute(query, name = nil)
     end
 
     def self.enable
-      ActiveRecord::Base.on_shard(nil) { ActiveRecord::Base.connection.class.prepend(Methods) }
+      ActiveRecord::Base.connection.class.prepend(Methods)
     end
   end
 end