switchman 1.14.4 → 1.15.0

data/lib/switchman/active_record/association.rb

@@ -30,17 +30,14 @@ module Switchman
     end
 
     module CollectionAssociation
-      method = ::Rails.version < '5.1' ? :get_records : :find_target
-      module_eval <<-RUBY, __FILE__, __LINE__ + 1
-        def #{method}
-          shards = reflection.options[:multishard] && owner.respond_to?(:associated_shards) ? owner.associated_shards : [shard]
-          # activate both the owner and the target's shard category, so that Reflection#join_id_for,
-          # when called for the owner, will be returned relative to shard the query will execute on
-          Shard.with_each_shard(shards, [klass.shard_category, owner.class.shard_category].uniq) do
-            super
-          end
+      def find_target
+        shards = reflection.options[:multishard] && owner.respond_to?(:associated_shards) ? owner.associated_shards : [shard]
+        # activate both the owner and the target's shard category, so that Reflection#join_id_for,
+        # when called for the owner, will be returned relative to shard the query will execute on
+        Shard.with_each_shard(shards, [klass.shard_category, owner.class.shard_category].uniq) do
+          super
         end
-      RUBY
+      end
     end
 
     module BelongsToAssociation
@@ -88,7 +85,11 @@ module Switchman
           # Copypasta from Activerecord but with added global_id_for goodness.
           def records_for(ids)
             scope.where(association_key_name => ids).load do |record|
-              global_key = Shard.global_id_for(record[association_key_name], record.shard)
+              global_key = if record.class.shard_category == :unsharded
+                             convert_key(record[association_key_name])
+                           else
+                             Shard.global_id_for(record[association_key_name], record.shard)
+                           end
               owner = owners_by_key[global_key.to_s].first
               association = owner.association(reflection.name)
               association.set_inverse_instance(record)

data/lib/switchman/active_record/base.rb

@@ -68,6 +68,14 @@ module Switchman
             result
           end
         end
+
+        def clear_query_caches_for_current_thread
+          ::ActiveRecord::Base.connection_handlers.each_value do |handler|
+            handler.connection_pool_list.each do |pool|
+              pool.connection(switch_shard: false).clear_query_cache if pool.active_connection?
+            end
+          end
+        end
       end
 
       def self.included(klass)
@@ -99,12 +107,12 @@ module Switchman
 
       def save(*args)
         @shard_set_in_stone = true
-        self.class.shard(shard, :implicit).scoping { super }
+        (self.class.current_scope || self.class.default_scoped).shard(shard, :implicit).scoping { super }
       end
 
       def save!(*args)
         @shard_set_in_stone = true
-        self.class.shard(shard, :implicit).scoping { super }
+        (self.class.current_scope || self.class.default_scoped).shard(shard, :implicit).scoping { super }
       end
 
       def destroy

data/lib/switchman/active_record/calculations.rb

@@ -121,9 +121,8 @@ module Switchman
           group_fields = group_attrs
         end
 
-        # .clone below corrects for what I consider a Rails bug. column_alias_for modifies the string in place.
         # to_s is because Rails 5 returns a string but Rails 6 returns a symbol.
-        group_aliases = group_fields.map { |field| column_alias_for(field.clone).to_s }
+        group_aliases = group_fields.map { |field| column_alias_for(field.downcase.to_s).to_s }
         group_columns = group_aliases.zip(group_fields).map { |aliaz, field|
           [aliaz, type_for(field), column_name_for(field)]
         }

data/lib/switchman/active_record/connection_handler.rb

@@ -24,6 +24,10 @@ module Switchman
       end
 
       def establish_connection(spec)
+        # Just skip establishing a sharded connection if sharding isn't loaded; we'll do it again later
+        # This only can happen when loading ActiveRecord::Base; after everything is loaded Shard will
+        # be defined and this will actually establish a connection
+        return unless defined?(Shard)
         pool = super
 
         # this is the first place that the adapter would have been required; but now we
@@ -39,16 +43,24 @@ module Switchman
           # to sharding will recurse onto itself trying to access column information
           Shard.default
 
+          config = pool.spec.config
           # automatically change config to allow for sharing connections with simple config
-          config = ::Rails.version < '5.1' ? spec.config : pool.spec.config
           ConnectionHandler.make_sharing_automagic(config)
           ConnectionHandler.make_sharing_automagic(Shard.default.database_server.config)
 
-          if ::Rails.version < '5.1'
-            ::ActiveRecord::Base.configurations[::Rails.env] = spec.instance_variable_get(:@config).stringify_keys
-          else
+          if ::Rails.version < '6.0'
             ::ActiveRecord::Base.configurations[::Rails.env] = config.stringify_keys
+          else
+            # Adopted from the deprecated code that currently lives in rails proper
+            remaining_configs = ::ActiveRecord::Base.configurations.configurations.reject { |db_config| db_config.env_name == ::Rails.env }
+            new_config = ::ActiveRecord::DatabaseConfigurations.new(::Rails.env => config.stringify_keys).configurations
+            new_configs = remaining_configs + new_config
+
+            ::ActiveRecord::Base.configurations = new_configs
           end
+        else
+          # this is probably wrong now
+          Shard.default.remove_instance_variable(:@name) if Shard.default.instance_variable_defined?(:@name)
         end
 
         @shard_connection_pools ||= { [:master, Shard.default.database_server.shareable? ? ::Rails.env : Shard.default] => pool}
@@ -122,8 +134,7 @@ module Switchman
             else
               ancestor_pool.spec
             end
-            spec = spec.to_hash if ::Rails.version >= '5.1'
-            pool = establish_connection spec
+            pool = establish_connection(spec.to_hash)
             pool.instance_variable_set(:@schema_cache, ancestor_pool.schema_cache) if ancestor_pool.schema_cache
             pool
           elsif spec_name != "primary"

data/lib/switchman/active_record/connection_pool.rb

@@ -32,10 +32,10 @@ module Switchman
         conn
       end
 
-      def connection
-        conn = super
+      def connection(switch_shard: true)
+        conn = super()
         raise NonExistentShardError if shard.new_record?
-        switch_database(conn) if conn.shard != self.shard
+        switch_database(conn) if conn.shard != self.shard && switch_shard
         conn
       end
 
@@ -47,6 +47,20 @@ module Switchman
         end
       end
 
+      def remove_shard!(shard)
+        synchronize do
+          # The shard might be currently active, so we need to update our own shard
+          if self.shard == shard
+            self.shard = Shard.default
+          end
+          # Update out any connections that may be using this shard
+          @connections.each do |conn|
+            # This will also update the connection's shard to the default shard
+            switch_database(conn) if conn.shard == shard
+          end
+        end
+      end
+
       def clear_idle_connections!(since_when)
         synchronize do
           @connections.reject! do |conn|

data/lib/switchman/active_record/log_subscriber.rb

@@ -18,18 +18,14 @@ module Switchman
         shard = "  [#{shard[:database_server_id]}:#{shard[:id]} #{shard[:env]}]" if shard
 
         unless (payload[:binds] || []).empty?
-          if ::Rails.version < '5.0.3'
-            binds = "  " + payload[:binds].map { |attr| render_bind(attr) }.inspect
-          else
-            use_old_format = (::Rails.version < '5.1') ? (::Rails.version < '5.0.7') : (::Rails.version < '5.1.5')
-            args = use_old_format ?
-              [payload[:binds], payload[:type_casted_binds]] :
-              [payload[:type_casted_binds]]
-            casted_params = type_casted_binds(*args)
-            binds = "  " + payload[:binds].zip(casted_params).map { |attr, value|
-              render_bind(attr, value)
-            }.inspect
-          end
+          use_old_format = (::Rails.version < '5.1.5')
+          args = use_old_format ?
+            [payload[:binds], payload[:type_casted_binds]] :
+            [payload[:type_casted_binds]]
+          casted_params = type_casted_binds(*args)
+          binds = "  " + payload[:binds].zip(casted_params).map { |attr, value|
+            render_bind(attr, value)
+          }.inspect
         end
 
         name = colorize_payload_name(name, payload[:name])

data/lib/switchman/active_record/model_schema.rb

@@ -4,7 +4,7 @@ module Switchman
       module ClassMethods
         def quoted_table_name
           @quoted_table_name ||= {}
-          @quoted_table_name[Shard.current.id] ||= connection.quote_table_name(table_name)
+          @quoted_table_name[Shard.current(shard_category).id] ||= connection.quote_table_name(table_name)
         end
       end
     end

data/lib/switchman/active_record/postgresql_adapter.rb

@@ -52,31 +52,12 @@ module Switchman
         SQL
       end
 
-      if ::Rails.version >= '5.1'
-        def extract_schema_qualified_name(string)
-          name = ::ActiveRecord::ConnectionAdapters::PostgreSQL::Utils.extract_schema_qualified_name(string.to_s)
-          if string && !name.schema && use_qualified_names?
-            name.instance_variable_set(:@schema, shard.name)
-          end
-          [name.schema, name.identifier]
-        end
-      else
-        def data_source_exists?(name)
-          name = ::ActiveRecord::ConnectionAdapters::PostgreSQL::Utils.extract_schema_qualified_name(name.to_s)
-          return false unless name.identifier
-          if !name.schema && use_qualified_names?
-            name.instance_variable_set(:@schema, shard.name)
-          end
-
-          exec_query(<<-SQL, 'SCHEMA').rows.first[0].to_i > 0
-            SELECT COUNT(*)
-            FROM pg_class c
-            LEFT JOIN pg_namespace n ON n.oid = c.relnamespace
-            WHERE c.relkind IN ('r','v','m') -- (r)elation/table, (v)iew, (m)aterialized view
-            AND c.relname = '#{name.identifier}'
-            AND n.nspname = #{name.schema ? "'#{name.schema}'" : 'ANY (current_schemas(false))'}
-          SQL
+      def extract_schema_qualified_name(string)
+        name = ::ActiveRecord::ConnectionAdapters::PostgreSQL::Utils.extract_schema_qualified_name(string.to_s)
+        if string && !name.schema && use_qualified_names?
+          name.instance_variable_set(:@schema, shard.name)
         end
+        [name.schema, name.identifier]
       end
 
       def view_exists?(name)

data/lib/switchman/active_record/query_cache.rb

@@ -1,122 +1,32 @@
 module Switchman
   module ActiveRecord
     module QueryCache
-      if ::Rails.version < '5.0.1'
-        # thread local accessors to replace @query_cache_enabled
-        def query_cache
-          thread_cache = Thread.current[:query_cache] ||= {}
-          thread_cache[self.object_id] ||= Hash.new { |h,sql| h[sql] = {} }
-        end
-
-        def query_cache_enabled
-          Thread.current[:query_cache_enabled]
-        end
-
-        def query_cache_enabled=(value)
-          Thread.current[:query_cache_enabled] = value
-        end
-
-        # basically wholesale repeat of the methods from the original (see
-        # https://github.com/rails/rails/blob/master/activerecord/lib/active_record/connection_adapters/abstract/query_cache.rb),
-        # but with self.query_cache_enabled and self.query_cache_enabled= instead
-        # of @query_cache_enabled.
-
-        def enable_query_cache!
-          self.query_cache_enabled = true
-        end
-
-        def disable_query_cache!
-          self.query_cache_enabled = false
-        end
-
-        def cache
-          old, self.query_cache_enabled = query_cache_enabled, true
-          yield
-        ensure
-          self.query_cache_enabled = old
-          clear_query_cache unless self.query_cache_enabled
-        end
-
-        def uncached
-          old, self.query_cache_enabled = query_cache_enabled, false
-          yield
-        ensure
-          self.query_cache_enabled = old
-        end
-
-        def clear_query_cache
-          Thread.current[:query_cache]&.clear
-        end
-
-        def select_all(arel, name = nil, binds = [], preparable: nil)
-          if self.query_cache_enabled && !locked?(arel)
-            arel, binds = binds_from_relation(arel, binds)
-            sql = to_sql(arel, binds)
-            cache_sql(sql, binds) { super(sql, name, binds, preparable: preparable) }
-          else
-            super
-          end
-        end
-
-        # no reason to define these on the including class directly. the super
-        # works just as well from a method on the included module
-        [:insert, :update, :delete].each do |method_name|
-          class_eval <<-end_code, __FILE__, __LINE__ + 1
-            def #{method_name}(*args)
-              clear_query_cache if self.query_cache_enabled
-              super
-            end
-          end_code
-        end
-      end
 
       private
 
-      if ::Rails.version < '5.1'
-        def cache_sql(sql, binds)
-          # have to include the shard id in the cache key because of switching dbs on the same connection
-          sql = "#{self.shard.id}::#{sql}"
+      def cache_sql(sql, name, binds)
+        # have to include the shard id in the cache key because of switching dbs on the same connection
+        sql = "#{self.shard.id}::#{sql}"
+        @lock.synchronize do
           result =
               if query_cache[sql].key?(binds)
-                args = {:sql => sql, :binds => binds, :name => "CACHE", :connection_id => object_id}
-                args[:type_casted_binds] = -> { type_casted_binds(binds) } if ::Rails.version >= '5.0.7'
-                ::ActiveSupport::Notifications.instrument("sql.active_record", args)
+                args = {
+                  sql: sql,
+                  binds: binds,
+                  name: name,
+                  connection_id: object_id,
+                  cached: true
+                }
+                args[:type_casted_binds] = -> { type_casted_binds(binds) } if ::Rails.version >= '5.1.5'
+                ::ActiveSupport::Notifications.instrument(
+                    "sql.active_record",
+                    args
+                )
                 query_cache[sql][binds]
               else
                 query_cache[sql][binds] = yield
               end
-
-          if ::ActiveRecord::Result === result
-            result.dup
-          else
-            result.collect { |row| row.dup }
-          end
-        end
-      else
-        def cache_sql(sql, name, binds)
-          # have to include the shard id in the cache key because of switching dbs on the same connection
-          sql = "#{self.shard.id}::#{sql}"
-          @lock.synchronize do
-            result =
-                if query_cache[sql].key?(binds)
-                  args = {
-                    sql: sql,
-                    binds: binds,
-                    name: name,
-                    connection_id: object_id,
-                    cached: true
-                  }
-                  args[:type_casted_binds] = -> { type_casted_binds(binds) } if ::Rails.version >= '5.1.5'
-                  ::ActiveSupport::Notifications.instrument(
-                      "sql.active_record",
-                      args
-                  )
-                  query_cache[sql][binds]
-                else
-                  query_cache[sql][binds] = yield
-                end
-            result.dup
-          end
+          result.dup
         end
       end
     end

data/lib/switchman/active_record/statement_cache.rb

@@ -56,15 +56,7 @@ module Switchman
         end
       end
 
-      if ::Rails.version < '5.1'
-        def generic_query_builder(connection)
-          @query_builder ||= connection.cacheable_query(@arel)
-        end
-
-        def qualified_query_builder(shard, klass)
-          @qualified_query_builders[shard.id] ||= klass.connection.cacheable_query(@arel)
-        end
-      elsif ::Rails.version < '5.2'
+      if ::Rails.version < '5.2'
         def generic_query_builder(connection)
           @query_builder ||= connection.cacheable_query(self.class, @arel)
         end

data/lib/switchman/active_support/cache.rb

@@ -4,10 +4,26 @@ module Switchman
       module ClassMethods
         def lookup_store(*store_options)
           store = super
+          # can't use defined?, because it's a _ruby_ autoloaded constant,
+          # so just checking that will cause it to get required
+          if store.class.name == "ActiveSupport::Cache::RedisCacheStore" && !::ActiveSupport::Cache::RedisCacheStore.ancestors.include?(RedisCacheStore)
+            ::ActiveSupport::Cache::RedisCacheStore.prepend(RedisCacheStore)
+          end
           store.options[:namespace] ||= lambda { Shard.current.default? ? nil : "shard_#{Shard.current.id}" }
           store
         end
       end
+
+      module RedisCacheStore
+        def clear(options = {})
+          # RedisCacheStore tries to be smart and only clear the cache under your namespace, if you have one set
+          # unfortunately, it uses the keys command, which is extraordinarily inefficient in a large redis instance
+          # fortunately, we can assume we control the entire instance, because we set up the namespacing, so just
+          # always unset it temporarily for clear calls
+          options[:namespace] = nil
+          super
+        end
+      end
     end
   end
 end

data/lib/switchman/connection_pool_proxy.rb

@@ -23,7 +23,14 @@ module Switchman
       @category = category
       @default_pool = default_pool
       @connection_pools = shard_connection_pools
-      @schema_cache = SchemaCache.new(self)
+      @schema_cache = default_pool.get_schema_cache(nil) if ::Rails.version >= '6'
+      @schema_cache = SchemaCache.new(self) unless @schema_cache.is_a?(SchemaCache)
+      if ::Rails.version >= '6'
+        @default_pool.set_schema_cache(@schema_cache)
+        @connection_pools.each_value do |pool|
+          pool.set_schema_cache(@schema_cache)
+        end
+      end
     end
 
     def active_shard
@@ -46,11 +53,11 @@ module Switchman
       connection_pools.map(&:connections).inject([], &:+)
     end
 
-    def connection
+    def connection(switch_shard: true)
       pool = current_pool
       begin
-        connection = pool.connection
-        connection.instance_variable_set(:@schema_cache, @schema_cache)
+        connection = pool.connection(switch_shard: switch_shard)
+        connection.instance_variable_set(:@schema_cache, @schema_cache) unless ::Rails.version >= '6'
         connection
       rescue ConnectionError
         raise if active_shard.database_server == Shard.default.database_server && active_shackles_environment == :master
@@ -60,7 +67,7 @@ module Switchman
           pool = create_pool(config.dup)
           begin
             connection = pool.connection
-            connection.instance_variable_set(:@schema_cache, @schema_cache)
+            connection.instance_variable_set(:@schema_cache, @schema_cache) unless ::Rails.version >= '6'
           rescue ConnectionError
             raise if idx == configs.length - 1
             next
@@ -71,6 +78,10 @@ module Switchman
       end
     end
 
+    def get_schema_cache(_connection)
+      @schema_cache
+    end
+
     %w{release_connection
        disconnect!
        flush!
@@ -98,6 +109,10 @@ module Switchman
       connection_pools.each { |pool| pool.clear_idle_connections!(since_when) }
     end
 
+    def remove_shard!(shard)
+      connection_pools.each { |pool| pool.remove_shard!(shard) }
+    end
+
     protected
 
     def connection_pools
@@ -132,18 +147,19 @@ module Switchman
           end
         end
       end
-      args = [config, "#{config[:adapter]}_connection"]
-      args.unshift(pool_key.join("/"))
-      spec = ::ActiveRecord::ConnectionAdapters::ConnectionSpecification.new(*args)
+      spec = ::ActiveRecord::ConnectionAdapters::ConnectionSpecification.new(
+        category,
+        config,
+        "#{config[:adapter]}_connection"
+      )
       # unfortunately the AR code that does this require logic can't really be
       # called in isolation
       require "active_record/connection_adapters/#{config[:adapter]}_adapter"
 
       ::ActiveRecord::ConnectionAdapters::ConnectionPool.new(spec).tap do |pool|
         pool.shard = shard
-        if ::Rails.version >= '5.0.1'
-          pool.enable_query_cache! if !@connection_pools.empty? && @connection_pools.first.last.query_cache_enabled
-        end
+        pool.set_schema_cache(@schema_cache) if ::Rails.version >= '6'
+        pool.enable_query_cache! if !@connection_pools.empty? && @connection_pools.first.last.query_cache_enabled
       end
     end
   end

data/lib/switchman/database_server.rb

@@ -185,6 +185,7 @@ module Switchman
         shard = Shard.create!(:id => shard_id,
                               :name => name,
                               :database_server => self)
+        schema_already_existed = false
 
         begin
           self.class.creating_new_shard = true
@@ -192,6 +193,10 @@ module Switchman
             ::Shackles.activate(:deploy) do
               begin
                 if create_statement
+                  if (::ActiveRecord::Base.connection.select_value("SELECT 1 FROM pg_namespace WHERE nspname=#{::ActiveRecord::Base.connection.quote(name)}"))
+                    schema_already_existed = true
+                    raise "This schema already exists; cannot overwrite"
+                  end
                   Array(create_statement.call).each do |stmt|
                     ::ActiveRecord::Base.connection.execute(stmt)
                   end
@@ -229,7 +234,9 @@ module Switchman
           shard
         rescue
           shard.destroy
-          shard.drop_database rescue nil
+          unless schema_already_existed
+            shard.drop_database rescue nil
+          end
           reset_column_information unless create_schema == false rescue nil
           raise
         ensure

data/lib/switchman/default_shard.rb

@@ -3,6 +3,7 @@ require 'switchman/database_server'
 module Switchman
   class DefaultShard
     def id; 'default'; end
+    alias cache_key id
     def activate(*categories); yield; end
     def activate!(*categories); end
     def default?; true; end

data/lib/switchman/engine.rb

@@ -88,7 +88,7 @@ module Switchman
         require "switchman/call_super"
         require "switchman/rails"
         require "switchman/shackles/relation"
-        require_dependency "switchman/shard_internal"
+        require_dependency "switchman/shard"
         require "switchman/standard_error"
 
         ::StandardError.include(StandardError)
@@ -118,11 +118,6 @@ module Switchman
         ::ActiveRecord::ConnectionAdapters::ConnectionHandler.prepend(ActiveRecord::ConnectionHandler)
         ::ActiveRecord::ConnectionAdapters::ConnectionPool.prepend(ActiveRecord::ConnectionPool)
         ::ActiveRecord::ConnectionAdapters::AbstractAdapter.prepend(ActiveRecord::QueryCache)
-        # when we call super in Switchman::ActiveRecord::QueryCache#select_all,
-        # we want it to find the definition from
-        # ActiveRecord::ConnectionAdapters::DatabaseStatements, not
-        # ActiveRecord::ConnectionAdapters::QueryCache
-        ::ActiveRecord::ConnectionAdapters::QueryCache.send(:remove_method, :select_all) if ::Rails.version < '5.0.1'
 
         ::ActiveRecord::LogSubscriber.prepend(ActiveRecord::LogSubscriber)
         ::ActiveRecord::Migration.prepend(ActiveRecord::Migration)
@@ -173,6 +168,12 @@ module Switchman
           require "switchman/active_record/postgresql_adapter"
           ::ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.prepend(ActiveRecord::PostgreSQLAdapter)
         end
+
+        # If Switchman::Shard wasn't loaded as of when ActiveRecord::Base initialized
+        # establish a connection here instead
+        if !Shard.instance_variable_get(:@default)
+          ::ActiveRecord::Base.establish_connection
+        end
       end
     end
 

data/lib/switchman/r_spec_helper.rb

@@ -121,7 +121,7 @@ module Switchman
       klass.before do
         raise "Sharding did not set up correctly" if @@sharding_failed
         Shard.clear_cache
-        if ::Rails.version >= '5.1' ? use_transactional_tests : (use_transactional_tests || use_transactional_fixtures)
+        if use_transactional_tests
           Shard.default(true)
           @shard1 = Shard.find(@shard1.id)
           @shard2 = Shard.find(@shard2.id)
@@ -137,7 +137,7 @@ module Switchman
 
       klass.after do
         next if @@sharding_failed
-        if ::Rails.version >= '5.1' ? use_transactional_tests : (use_transactional_tests || use_transactional_fixtures)
+        if use_transactional_tests
           shards = [@shard2]
           shards << @shard1 unless @shard1.database_server == Shard.default.database_server
           shards.each do |shard|

data/lib/switchman/version.rb

@@ -1,3 +1,3 @@
 module Switchman
-  VERSION = "1.14.4"
+  VERSION = "1.15.0"
 end

data/lib/tasks/switchman.rake

@@ -20,7 +20,12 @@ module Switchman
         open = servers.delete('open')
 
         servers = servers.map { |server| DatabaseServer.find(server) }.compact
-        servers.concat(DatabaseServer.all.select { |server| server.config[:open] }) if open
+        if open
+          open_servers = DatabaseServer.all.select { |server| server.config[:open] }
+          servers.concat(open_servers)
+          servers << DatabaseServer.find(nil) if open_servers.empty?
+          servers.uniq!
+        end
         servers = DatabaseServer.all - servers if negative
       end
 
@@ -67,7 +72,17 @@ module Switchman
                 shard = Shard.current
                 puts "#{shard.id}: #{shard.description}"
                 ::ActiveRecord::Base.connection_pool.spec.config[:shard_name] = Shard.current.name
-                ::ActiveRecord::Base.configurations[::Rails.env] = ::ActiveRecord::Base.connection_pool.spec.config.stringify_keys
+                if ::Rails.version < '6.0'
+                  ::ActiveRecord::Base.configurations[::Rails.env] = ::ActiveRecord::Base.connection_pool.spec.config.stringify_keys
+                else
+                  # Adopted from the deprecated code that currently lives in rails proper
+                  remaining_configs = ::ActiveRecord::Base.configurations.configurations.reject { |db_config| db_config.env_name == ::Rails.env }
+                  new_config = ::ActiveRecord::DatabaseConfigurations.new(::Rails.env =>
+                    ::ActiveRecord::Base.connection_pool.spec.config.stringify_keys).configurations
+                  new_configs = remaining_configs + new_config
+      
+                  ::ActiveRecord::Base.configurations = new_configs
+                end
                 shard.database_server.unshackle do
                   old_actions.each { |action| action.call(*task_args) }
                 end
@@ -195,9 +210,7 @@ module Switchman
       @filter_database_servers_chain ||= ->(servers) { servers }
     end
   end
-end
 
-module Switchman
   module ActiveRecord
     module PostgreSQLDatabaseTasks
       def structure_dump(filename, extra_flags=nil)