switchman 1.13.3 → 2.0.0

data/db/migrate/20130328224244_create_default_shard.rb

@@ -3,7 +3,7 @@ class CreateDefaultShard < ActiveRecord::Migration[4.2]
     unless Switchman::Shard.default.is_a?(Switchman::Shard)
       Switchman::Shard.reset_column_information
       Switchman::Shard.create!(:default => true)
-      Switchman::Shard.default(true)
+      Switchman::Shard.default(reload: true)
     end
   end
 end

data/lib/switchman/active_record/abstract_adapter.rb

@@ -4,8 +4,8 @@ module Switchman
   module ActiveRecord
     module AbstractAdapter
       module ForeignKeyCheck
-        def add_column(table, name, type, options = {})
-          Engine.foreign_key_check(name, type, options)
+        def add_column(table, name, type, limit: nil, **)
+          Engine.foreign_key_check(name, type, limit: limit)
           super
         end
       end
@@ -27,10 +27,6 @@ module Switchman
         quote_table_name(name)
       end
 
-      def use_qualified_names?
-        false
-      end
-
       protected
 
       def log(*args, &block)

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
@@ -76,15 +73,34 @@ module Switchman
 
     module Preloader
       module Association
-        if ::Rails.version >= "5.2"
+        if ::Rails.version >= "5.2" and ::Rails.version < "6.0"
           def run(preloader)
-            # TODO - can move associated_records_by_owner into this after 5.1 is gonzo
             associated_records_by_owner(preloader).each do |owner, records|
               associate_records_to_owner(owner, records)
             end
           end
         end
 
+        if ::Rails.version >= "6.0"
+          # 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 = 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)
+            end
+          end
+
+          def records_by_owner
+            associated_records_by_owner
+          end
+        end
+
         def associated_records_by_owner(preloader = nil)
           owners_map = owners_by_key
 
@@ -183,6 +199,13 @@ module Switchman
           (record.has_attribute?(reflection.foreign_key) && record.send(reflection.foreign_key) != key) || # have to use send instead of [] because sharding
           record.attribute_changed?(reflection.foreign_key)
       end
+
+      def save_belongs_to_association(reflection)
+        # this seems counter-intuitive, but the autosave code will assign to attribute bypassing switchman,
+        # after reading the id attribute _without_ bypassing switchman. So we need Shard.current for the
+        # category of the associated record to match Shard.current for the category of self
+        shard.activate(shard_category_for_reflection(reflection)) { super }
+      end
     end
   end
 end

data/lib/switchman/active_record/attribute_methods.rb

@@ -61,6 +61,9 @@ module Switchman
           end
         end
 
+        # see also Base#shard_category_for_reflection
+        # the difference being this will output static strings for the common cases, making them
+        # more performant
         def shard_category_code_for_reflection(reflection)
           if reflection
             if reflection.options[:polymorphic]

data/lib/switchman/active_record/base.rb

@@ -31,20 +31,20 @@ module Switchman
           @integral_id
         end
 
-        def transaction(*args)
+        def transaction(**)
           if self != ::ActiveRecord::Base && current_scope
             current_scope.activate do
               db = Shard.current(shard_category).database_server
-              if ::Shackles.environment != db.shackles_environment
-                db.unshackle { super }
+              if ::GuardRail.environment != db.guard_rail_environment
+                db.unguard { super }
               else
                 super
               end
             end
           else
             db = Shard.current(shard_category).database_server
-            if ::Shackles.environment != db.shackles_environment
-              db.unshackle { super }
+            if ::GuardRail.environment != db.guard_rail_environment
+              db.unguard { super }
             else
               super
             end
@@ -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)
@@ -97,14 +105,14 @@ module Switchman
         end
       end
 
-      def save(*args)
+      def save(*, **)
         @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)
+      def save!(*, **)
         @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
@@ -146,6 +154,36 @@ module Switchman
         # do this the Rails 4.2 way, so that if Shard.current != self.shard, the id gets transposed
         self.class.connection.quote(id)
       end
+
+      def update_columns(*)
+        db = Shard.current(self.class.shard_category).database_server
+        if ::GuardRail.environment != db.guard_rail_environment
+          return db.unguard { super }
+        else
+          super
+        end
+      end
+
+      protected
+
+      # see also AttributeMethods#shard_category_code_for_reflection
+      def shard_category_for_reflection(reflection)
+        if reflection
+          if reflection.options[:polymorphic]
+            begin
+              read_attribute(reflection.foreign_type)&.constantize&.shard_category || :primary
+            rescue NameError
+              # in case someone is abusing foreign_type to not point to an actual class
+              :primary
+            end
+          else
+            # otherwise we can just return a symbol for the statically known type of the association
+            reflection.klass.shard_category
+          end
+        else
+          shard_category
+        end
+      end
     end
   end
 end

data/lib/switchman/active_record/calculations.rb

@@ -121,7 +121,8 @@ module Switchman
           group_fields = group_attrs
         end
 
-        group_aliases = group_fields.map { |field| column_alias_for(field) }
+        # 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.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

@@ -4,8 +4,6 @@ module Switchman
   module ActiveRecord
     module ConnectionHandler
       def self.make_sharing_automagic(config, shard = Shard.current)
-        key = config[:adapter] == 'postgresql' ? :schema_search_path : :database
-
         # only load the shard name from the db if we have to
         if !config[:shard_name]
           # we may not be able to connect to this shard yet, cause it might be an empty database server
@@ -15,15 +13,13 @@ module Switchman
 
           config[:shard_name] ||= shard_name
         end
-
-        if !config[key] || config[key] == shard_name
-          # this may truncate the schema_search_path if it was not specified in database.yml
-          # but that's what our old behavior was anyway, so I guess it's okay
-          config[key] = '%{shard_name}'
-        end
       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,19 +35,27 @@ 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}
+        @shard_connection_pools ||= { [:primary, Shard.default.database_server.shareable? ? ::Rails.env : Shard.default] => pool}
 
         category = pool.spec.name.to_sym
         proxy = ConnectionPoolProxy.new(category,
@@ -60,13 +64,13 @@ module Switchman
         owner_to_pool[pool.spec.name] = proxy
 
         if first_time
-          if Shard.default.database_server.config[:prefer_slave]
-            Shard.default.database_server.shackle!
+          if Shard.default.database_server.config[:prefer_secondary]
+            Shard.default.database_server.guard!
           end
 
-          if Shard.default.is_a?(DefaultShard) && Shard.default.database_server.config[:slave]
-            Shard.default.database_server.shackle!
-            Shard.default(true)
+          if Shard.default.is_a?(DefaultShard) && Shard.default.database_server.config[:secondary]
+            Shard.default.database_server.guard!
+            Shard.default(reload: true)
           end
         end
 
@@ -122,8 +126,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|
@@ -70,18 +84,6 @@ module Switchman
         end
 
         spec.config[:shard_name] = self.shard.name
-        case conn.adapter_name
-          when 'MySQL', 'Mysql2'
-            conn.execute("USE #{spec.config[:database]}")
-          when 'PostgreSQL'
-            if conn.schema_search_path != spec.config[:schema_search_path]
-              conn.schema_search_path = spec.config[:schema_search_path]
-            end
-          when 'SQLite'
-            # This is an artifact of the adapter modifying the path to be an absolute path when it is instantiated; just let it slide
-          else
-            raise("Cannot switch databases on same DatabaseServer with adapter type: #{conn.adapter_name}. Limit one Shard per DatabaseServer.")
-        end
         conn.shard = shard
       end
 

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/migration.rb

@@ -33,5 +33,14 @@ module Switchman
         ::ActiveRecord::Migrator::MIGRATOR_SALT * shard_name_hash
       end
     end
+
+    module MigrationContext
+      def migrations
+        return @migrations if instance_variable_defined?(:@migrations)
+        migrations_cache = Thread.current[:migrations_cache] ||= {}
+        key = Digest::MD5.hexdigest(migration_files.sort.join(','))
+        @migrations = migrations_cache[key] ||= super
+      end
+    end
   end
 end

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

@@ -38,51 +38,26 @@ module Switchman
         select_values("SELECT * FROM unnest(current_schemas(false))")
       end
 
-      def use_qualified_names?
-        @config[:use_qualified_names]
-      end
-
       def tables(name = nil)
-        schema = shard.name if use_qualified_names?
-
         query(<<-SQL, 'SCHEMA').map { |row| row[0] }
           SELECT tablename
           FROM pg_tables
-          WHERE schemaname = #{schema ? "'#{schema}'" : 'ANY (current_schemas(false))'}
+          WHERE schemaname = '#{shard.name}'
         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
+          name.instance_variable_set(:@schema, shard.name)
         end
+        [name.schema, name.identifier]
       end
 
       def view_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?
+        if !name.schema
           name.instance_variable_set(:@schema, shard.name)
         end
 
@@ -92,13 +67,11 @@ module Switchman
           LEFT JOIN pg_namespace n ON n.oid = c.relnamespace
           WHERE c.relkind IN ('v','m') -- (v)iew, (m)aterialized view
           AND c.relname = '#{name.identifier}'
-          AND n.nspname = #{name.schema ? "'#{name.schema}'" : 'ANY (current_schemas(false))'}
+          AND n.nspname = '#{shard.name}'
         SQL
       end
 
       def indexes(table_name)
-        schema = shard.name if use_qualified_names?
-
         result = query(<<-SQL, 'SCHEMA')
            SELECT distinct i.relname, d.indisunique, d.indkey, pg_get_indexdef(d.indexrelid), t.oid
            FROM pg_class t
@@ -107,7 +80,7 @@ module Switchman
            WHERE i.relkind = 'i'
              AND d.indisprimary = 'f'
              AND t.relname = '#{table_name}'
-             AND i.relnamespace IN (SELECT oid FROM pg_namespace WHERE nspname = #{schema ? "'#{schema}'" : 'ANY (current_schemas(false))'} )
+             AND i.relnamespace IN (SELECT oid FROM pg_namespace WHERE nspname = '#{shard.name}' )
           ORDER BY i.relname
         SQL
 
@@ -145,8 +118,6 @@ module Switchman
       end
 
       def index_name_exists?(table_name, index_name, _default = nil)
-        schema = shard.name if use_qualified_names?
-
         exec_query(<<-SQL, 'SCHEMA').rows.first[0].to_i > 0
             SELECT COUNT(*)
             FROM pg_class t
@@ -155,7 +126,7 @@ module Switchman
             WHERE i.relkind = 'i'
               AND i.relname = '#{index_name}'
               AND t.relname = '#{table_name}'
-              AND i.relnamespace IN (SELECT oid FROM pg_namespace WHERE nspname = #{schema ? "'#{schema}'" : 'ANY (current_schemas(false))'} )
+              AND i.relnamespace IN (SELECT oid FROM pg_namespace WHERE nspname = '#{shard.name}' )
         SQL
       end
 
@@ -169,21 +140,21 @@ module Switchman
       def quote_table_name(name)
         return quote_local_table_name(name) if @use_local_table_name
         name = ::ActiveRecord::ConnectionAdapters::PostgreSQL::Utils.extract_schema_qualified_name(name.to_s)
-        if !name.schema && use_qualified_names?
+        if !name.schema
           name.instance_variable_set(:@schema, shard.name)
         end
         name.quoted
       end
 
-      def with_local_table_name
-        @use_local_table_name = true
+      def with_local_table_name(enable = true)
+        old_value = @use_local_table_name
+        @use_local_table_name = enable
         yield
       ensure
-        @use_local_table_name = false
+        @use_local_table_name = old_value
       end
 
       def foreign_keys(table_name)
-        schema = shard.name if use_qualified_names?
 
         # mostly copy-pasted from AR - only change is to the nspname condition for qualified names support
         fk_info = select_all <<-SQL.strip_heredoc
@@ -196,7 +167,7 @@ module Switchman
           JOIN pg_namespace t3 ON c.connamespace = t3.oid
           WHERE c.contype = 'f'
             AND t1.relname = #{quote(table_name)}
-            AND t3.nspname = #{schema ? "'#{schema}'" : 'ANY (current_schemas(false))'}
+            AND t3.nspname = '#{shard.name}'
           ORDER BY c.conname
         SQL
 
@@ -213,7 +184,7 @@ module Switchman
           # strip the schema name from to_table if it matches
           to_table = row['to_table']
           to_table_qualified_name = ::ActiveRecord::ConnectionAdapters::PostgreSQL::Utils.extract_schema_qualified_name(to_table)
-          if use_qualified_names? && to_table_qualified_name.schema == shard.name
+          if to_table_qualified_name.schema == shard.name
             to_table = to_table_qualified_name.identifier
           end
 
@@ -221,7 +192,7 @@ module Switchman
         end
       end
 
-      def add_index_options(_table_name, _column_name, _options = {})
+      def add_index_options(_table_name, _column_name, **)
         index_name, index_type, index_columns, index_options, algorithm, using = super
         algorithm = nil if DatabaseServer.creating_new_shard && algorithm == "CONCURRENTLY"
         [index_name, index_type, index_columns, index_options, algorithm, using]
@@ -248,6 +219,10 @@ module Switchman
 
         execute "ALTER INDEX #{quote_table_name(old_name)} RENAME TO #{quote_local_table_name(new_name)}"
       end
+
+      def columns(*)
+        with_local_table_name(false) { super }
+      end
     end
   end
 end