active_record_shards 2.5.0 → 2.5.1

data/lib/active_record_shards/connection_switcher.rb

@@ -133,7 +133,7 @@ module ActiveRecordShards
     end
 
     def shard_env
-      @shard_env = defined?(Rails.env) ? Rails.env : RAILS_ENV
+      ActiveRecordShards.rails_env
     end
 
     def establish_shard_connection

data/lib/active_record_shards/shard_selection.rb

@@ -30,7 +30,7 @@ module ActiveRecordShards
     end
 
     def shard_name(klass = nil, try_slave = true)
-      s = "#{RAILS_ENV}"
+      s = ActiveRecordShards.rails_env.dup
       if the_shard = shard(klass)
         s << '_shard_'
         s << the_shard
@@ -45,4 +45,4 @@ module ActiveRecordShards
       {:shard => @shard, :slave => @on_slave}
     end
   end
-end
+end

data/lib/active_record_shards.rb

@@ -12,3 +12,9 @@ ActiveRecord::Base.extend(ActiveRecordShards::ConfigurationParser)
 ActiveRecord::Base.extend(ActiveRecordShards::Model)
 ActiveRecord::Base.extend(ActiveRecordShards::ConnectionSwitcher)
 ActiveRecord::Associations::AssociationCollection.send(:include, ActiveRecordShards::AssociationCollectionConnectionSelection)
+
+module ActiveRecordShards
+  def self.rails_env
+    Kernel.const_defined?(:Rails) ? Rails.env : RAILS_ENV
+  end
+end