switchman 2.1.0 → 3.0.6

data/lib/switchman/database_server.rb

@@ -1,6 +1,6 @@
 # frozen_string_literal: true
 
-require "securerandom"
+require 'securerandom'
 
 module Switchman
   class DatabaseServer
@@ -13,16 +13,22 @@ module Switchman
         database_servers.values
       end
 
+      def all_roles
+        @all_roles ||= all.map(&:roles).flatten.uniq
+      end
+
       def find(id_or_all)
-        return self.all if id_or_all == :all
-        return id_or_all.map { |id| self.database_servers[id || ::Rails.env] }.compact.uniq if id_or_all.is_a?(Array)
+        return all if id_or_all == :all
+        return id_or_all.map { |id| database_servers[id || ::Rails.env] }.compact.uniq if id_or_all.is_a?(Array)
+
         database_servers[id_or_all || ::Rails.env]
       end
 
       def create(settings = {})
-        raise "database servers should be set up in database.yml" unless ::Rails.env.test?
+        raise 'database servers should be set up in database.yml' unless ::Rails.env.test?
+
         id = settings[:id]
-        if !id
+        unless id
           @id ||= 0
           @id += 1
           id = @id
@@ -30,25 +36,42 @@ module Switchman
         server = DatabaseServer.new(id.to_s, settings)
         server.instance_variable_set(:@fake, true)
         database_servers[server.id] = server
+        ::ActiveRecord::Base.configurations.configurations <<
+          ::ActiveRecord::DatabaseConfigurations::HashConfig.new(::Rails.env, "#{server.id}/primary", settings)
+        Shard.send(:initialize_sharding)
+        server
       end
 
       def server_for_new_shard
         servers = all.select { |s| s.config[:open] }
         return find(nil) if servers.empty?
+
         servers[rand(servers.length)]
       end
 
       private
+
+      def reference_role(role)
+        return if all_roles.include?(role)
+
+        @all_roles << role
+        Shard.send(:initialize_sharding)
+      end
+
       def database_servers
-        unless @database_servers
+        if !@database_servers || @database_servers.empty?
           @database_servers = {}.with_indifferent_access
-          if ::Rails.version >= '6.0'
-            ::ActiveRecord::Base.configurations.configurations.each do |config|
-              @database_servers[config.env_name] = DatabaseServer.new(config.env_name, config.config)
+          ::ActiveRecord::Base.configurations.configurations.each do |config|
+            if config.name.include?('/')
+              name, role = config.name.split('/')
+            else
+              name, role = config.env_name, config.name
             end
-          else
-            ::ActiveRecord::Base.configurations.each do |(id, config)|
-              @database_servers[id] = DatabaseServer.new(id, config)
+
+            if role == 'primary'
+              @database_servers[name] = DatabaseServer.new(config.env_name, config.configuration_hash)
+            else
+              @database_servers[name].roles << role
             end
           end
         end
@@ -56,15 +79,33 @@ module Switchman
       end
     end
 
+    attr_reader :roles
+
     def initialize(id = nil, config = {})
       @id = id
       @config = config.deep_symbolize_keys
       @configs = {}
+      @roles = [:primary]
+    end
+
+    def connects_to_hash
+      self.class.all_roles.map do |role|
+        config_role = role
+        config_role = :primary unless roles.include?(role)
+        config_name = :"#{id}/#{config_role}"
+        config_name = :primary if id == ::Rails.env && config_role == :primary
+        [role.to_sym, config_name]
+      end.to_h
     end
 
     def destroy
-      raise "database servers should be set up in database.yml" unless ::Rails.env.test?
-      self.class.send(:database_servers).delete(self.id) if self.id
+      self.class.send(:database_servers).delete(id) if id
+      Shard.sharded_models.each do |klass|
+        self.class.all_roles.each do |role|
+          klass.connection_handler.remove_connection_pool(klass.connection_specification_name, role: role,
+                                                                                               shard: id.to_sym)
+        end
+      end
     end
 
     def fake?
@@ -72,20 +113,20 @@ module Switchman
     end
 
     def config(environment = :primary)
-      @configs[environment] ||= begin
-        if @config[environment].is_a?(Array)
+      @configs[environment] ||=
+        case @config[environment]
+        when Array
           @config[environment].map do |config|
             config = @config.merge((config || {}).symbolize_keys)
             # make sure GuardRail doesn't get any brilliant ideas about choosing the first possible server
             config.delete(environment)
             config
           end
-        elsif @config[environment].is_a?(Hash)
+        when Hash
           @config.merge(@config[environment])
         else
           @config
         end
-      end
     end
 
     def guard_rail_environment
@@ -111,172 +152,107 @@ module Switchman
       guard!(old_env)
     end
 
-    def shareable?
-      @shareable_environment_key ||= []
-      environment = guard_rail_environment
-      explicit_user = ::GuardRail.global_config[:username]
-      return @shareable if @shareable_environment_key == [environment, explicit_user]
-      @shareable_environment_key = [environment, explicit_user]
-      if explicit_user
-        username = explicit_user
-      else
-        config = self.config(environment)
-        config = config.first if config.is_a?(Array)
-        username = config[:username]
-      end
-      @shareable = username !~ /%?\{[a-zA-Z0-9_]+\}/
-    end
-
     def shards
-      if self.id == ::Rails.env
-        Shard.where("database_server_id IS NULL OR database_server_id=?", self.id)
+      if id == ::Rails.env
+        Shard.where('database_server_id IS NULL OR database_server_id=?', id)
       else
-        Shard.where(:database_server_id => self.id)
+        Shard.where(database_server_id: id)
       end
     end
 
-    def pool_key
-      self.id == ::Rails.env ? nil : self.id
-    end
-
-    def create_new_shard(options = {})
-      raise NotImplementedError.new("Cannot create new shards when sharding isn't initialized") unless Shard.default.is_a?(Shard)
-
-      name = options[:name]
-      create_schema = options[:schema]
-      # look for another shard associated with this db
-      other_shard = self.shards.where("name<>':memory:' OR name IS NULL").order(:id).first
-
-      case config[:adapter]
-        when 'postgresql'
-          create_statement = lambda { "CREATE SCHEMA #{name}" }
-          password = " PASSWORD #{::ActiveRecord::Base.connection.quote(config[:password])}" if config[:password]
-        else
-          create_statement = lambda { "CREATE DATABASE #{name}" }
+    def create_new_shard(id: nil, name: nil, schema: true)
+      unless Shard.default.is_a?(Shard)
+        raise NotImplementedError,
+              "Cannot create new shards when sharding isn't initialized"
       end
+
+      create_statement = -> { "CREATE SCHEMA #{name}" }
+      password = " PASSWORD #{::ActiveRecord::Base.connection.quote(config[:password])}" if config[:password]
       sharding_config = Switchman.config
       config_create_statement = sharding_config[config[:adapter]]&.[](:create_statement)
       config_create_statement ||= sharding_config[:create_statement]
       if config_create_statement
         create_commands = Array(config_create_statement).dup
         create_statement = lambda {
-          create_commands.map { |statement| statement.gsub('%{name}', name).gsub('%{password}', password || '') }
+          create_commands.map { |statement| format(statement, name: name, password: password) }
         }
       end
 
-      create_shard = lambda do
-        shard_id = options.fetch(:id) do
-          case config[:adapter]
-            when 'postgresql'
-              id_seq = Shard.connection.quote(Shard.connection.quote_table_name('switchman_shards_id_seq'))
-              next_id = Shard.connection.select_value("SELECT nextval(#{id_seq})")
-              next_id.to_i
-            else
-              nil
+      id ||= begin
+        id_seq = Shard.connection.quote(Shard.connection.quote_table_name('switchman_shards_id_seq'))
+        next_id = Shard.connection.select_value("SELECT nextval(#{id_seq})")
+        next_id.to_i
+      end
+
+      name ||= "#{config[:database]}_shard_#{id}"
+
+      schema_already_existed = false
+      shard = nil
+      Shard.connection.transaction do
+        self.class.creating_new_shard = true
+        DatabaseServer.send(:reference_role, :deploy)
+        ::ActiveRecord::Base.connected_to(shard: self.id.to_sym, role: :deploy) do
+          shard = Shard.create!(id: id,
+                                name: name,
+                                database_server_id: self.id)
+          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
           end
-        end
+          if config[:adapter] == 'postgresql'
+            old_proc = ::ActiveRecord::Base.connection.raw_connection.set_notice_processor do
+            end
+          end
+          old_verbose = ::ActiveRecord::Migration.verbose
+          ::ActiveRecord::Migration.verbose = false
 
-        if name.nil?
-          base_name = self.config[:database].to_s % self.config
-          base_name = nil if base_name == ':memory:'
-          base_name << '_' if base_name
-          base_id = shard_id || SecureRandom.uuid
-          name = "#{base_name}shard_#{base_id}"
-        end
+          unless schema == false
+            shard.activate do
+              reset_column_information
 
-        shard = Shard.create!(:id => shard_id,
-                              :name => name,
-                              :database_server => self)
-        schema_already_existed = false
-
-        begin
-          self.class.creating_new_shard = true
-          shard.activate(*Shard.categories) do
-            ::GuardRail.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
-                  # have to disconnect and reconnect to the correct db
-                  if self.shareable? && other_shard
-                    other_shard.activate { ::ActiveRecord::Base.connection }
-                  else
-                    ::ActiveRecord::Base.connection_pool.current_pool.disconnect!
-                  end
-                end
-                old_proc = ::ActiveRecord::Base.connection.raw_connection.set_notice_processor {} if config[:adapter] == 'postgresql'
-                old_verbose = ::ActiveRecord::Migration.verbose
-                ::ActiveRecord::Migration.verbose = false
-
-                unless create_schema == false
-                  reset_column_information
-
-                  migrate = ::Rails.version >= '5.2' ?
-                    -> { ::ActiveRecord::Base.connection.migration_context.migrate } :
-                    -> { ::ActiveRecord::Migrator.migrate(::ActiveRecord::Migrator.migrations_paths) }
-                  if ::ActiveRecord::Base.connection.supports_ddl_transactions?
-                    ::ActiveRecord::Base.connection.transaction(requires_new: true, &migrate)
-                  else
-                    migrate.call
-                  end
-                  reset_column_information
-                  ::ActiveRecord::Base.descendants.reject { |m| m == Shard || !m.table_exists? }.each(&:define_attribute_methods)
-                end
-              ensure
-                ::ActiveRecord::Migration.verbose = old_verbose
-                ::ActiveRecord::Base.connection.raw_connection.set_notice_processor(&old_proc) if old_proc
+              ::ActiveRecord::Base.connection.transaction(requires_new: true) do
+                ::ActiveRecord::Base.connection.migration_context.migrate
               end
+              reset_column_information
+              ::ActiveRecord::Base.descendants.reject do |m|
+                m <= UnshardedRecord || !m.table_exists?
+              end.each(&:define_attribute_methods)
             end
           end
-          shard
-        rescue
-          shard.destroy
-          unless schema_already_existed
-            shard.drop_database rescue nil
-          end
-          reset_column_information unless create_schema == false rescue nil
-          raise
         ensure
-          self.class.creating_new_shard = false
-        end
-      end
-
-      if Shard.connection.supports_ddl_transactions? && self.shareable? && other_shard
-        Shard.transaction do
-          other_shard.activate do
-            ::ActiveRecord::Base.connection.transaction(&create_shard)
-          end
+          ::ActiveRecord::Migration.verbose = old_verbose
+          ::ActiveRecord::Base.connection.raw_connection.set_notice_processor(&old_proc) if old_proc
         end
-      else
-        create_shard.call
+        shard
+      rescue
+        shard&.destroy
+        shard&.drop_database rescue nil unless schema_already_existed
+        reset_column_information unless schema == false rescue nil
+        raise
+      ensure
+        self.class.creating_new_shard = false
       end
     end
 
     def cache_store
-      unless @cache_store
-        @cache_store = Switchman.config[:cache_map][self.id] || Switchman.config[:cache_map][::Rails.env]
-      end
+      @cache_store ||= Switchman.config[:cache_map][id] || Switchman.config[:cache_map][::Rails.env]
       @cache_store
     end
 
     def shard_name(shard)
-      if config[:shard_name]
-        config[:shard_name]
-      elsif config[:adapter] == 'postgresql'
-        if shard == :bootstrap
-          # rescue nil because the database may not exist yet; if it doesn't,
-          # it will shortly, and this will be re-invoked
-          ::ActiveRecord::Base.connection.current_schemas.first rescue nil
-        else
-          shard.activate { ::ActiveRecord::Base.connection_pool.default_schema }
-        end
+      return config[:shard_name] if config[:shard_name]
+
+      if shard == :bootstrap
+        # rescue nil because the database may not exist yet; if it doesn't,
+        # it will shortly, and this will be re-invoked
+        ::ActiveRecord::Base.connection.current_schemas.first rescue nil
       else
-        config[:database]
+        shard.activate { ::ActiveRecord::Base.connection_pool.default_schema }
       end
     end
 
@@ -290,9 +266,9 @@ module Switchman
     end
 
     private
+
     def reset_column_information
-      ::ActiveRecord::Base.descendants.reject { |m| m == Shard }.each(&:reset_column_information)
-      ::ActiveRecord::Base.connection_handler.switchman_connection_pool_proxies.each { |pool| pool.schema_cache.clear! }
+      ::ActiveRecord::Base.descendants.reject { |m| m <= UnshardedRecord }.each(&:reset_column_information)
     end
   end
 end

data/lib/switchman/default_shard.rb

@@ -4,17 +4,44 @@ require 'switchman/database_server'
 
 module Switchman
   class DefaultShard
-    def id; 'default'; end
+    def id
+      'default'
+    end
     alias cache_key id
-    def activate(*categories); yield; end
-    def activate!(*categories); end
-    def default?; true; end
-    def primary?; true; end
-    def relative_id_for(local_id, target = nil); local_id; end
-    def global_id_for(local_id); local_id; end
-    def database_server_id; nil; end
-    def database_server; DatabaseServer.find(nil); end
-    def new_record?; false; end
+    def activate(*_classes)
+      yield
+    end
+
+    def activate!(*classes); end
+
+    def default?
+      true
+    end
+
+    def primary?
+      true
+    end
+
+    def relative_id_for(local_id, _target = nil)
+      local_id
+    end
+
+    def global_id_for(local_id)
+      local_id
+    end
+
+    def database_server_id
+      nil
+    end
+
+    def database_server
+      DatabaseServer.find(nil)
+    end
+
+    def new_record?
+      false
+    end
+
     def name
       unless instance_variable_defined?(:@name)
         @name = nil # prevent taking this branch on recursion
@@ -22,18 +49,27 @@ module Switchman
       end
       @name
     end
-    def description; ::Rails.env; end
+
+    def description
+      ::Rails.env
+    end
+
     # The default's shard is always the default shard
-    def shard; self; end
-    def _dump(depth)
+    def shard
+      self
+    end
+
+    def _dump(_depth)
       ''
     end
-    def self._load(str)
+
+    def self._load(_str)
       Shard.default
     end
 
-    def ==(rhs)
-      return true if rhs.is_a?(DefaultShard) || (rhs.is_a?(Shard) && rhs[:default])
+    def ==(other)
+      return true if other.is_a?(DefaultShard) || (other.is_a?(Shard) && other[:default])
+
       super
     end