switchman 2.0.6 → 3.0.1

data/lib/switchman/active_record/connection_pool.rb

@@ -6,15 +6,18 @@ module Switchman
   module ActiveRecord
     module ConnectionPool
       def shard
-        Thread.current[tls_key] || Shard.default
+        shard_stack.last || Shard.default
       end
 
-      def shard=(value)
-        Thread.current[tls_key] = value
+      def shard_stack
+        unless (shard_stack = Thread.current.thread_variable_get(tls_key))
+          shard_stack = Concurrent::Array.new
+          Thread.current.thread_variable_set(tls_key, shard_stack)
+        end
+        shard_stack
       end
 
       def default_schema
-        raise "Not postgres!" unless self.spec.config[:adapter] == 'postgresql'
         connection unless @schemas
         # default shard will not switch databases immediately, so it won't be set yet
         @schemas ||= connection.current_schemas
@@ -22,39 +25,29 @@ module Switchman
       end
 
       def checkout_new_connection
-        conn = synchronize do
-          # ideally I would just keep a thread-local spec that I could modify
-          # without locking anything, but if spec returns not-the-object passed
-          # to initialize this pool, things break
-          spec.config[:shard_name] = self.shard.name
-
-          super
-        end
-        conn.shard = self.shard
+        conn = super
+        conn.shard = shard
         conn
       end
 
       def connection(switch_shard: true)
         conn = super()
         raise NonExistentShardError if shard.new_record?
-        switch_database(conn) if conn.shard != self.shard && switch_shard
+
+        switch_database(conn) if conn.shard != shard && switch_shard
         conn
       end
 
       def release_connection(with_id = Thread.current)
         super(with_id)
 
-        if spec.config[:idle_timeout]
-          clear_idle_connections!(Time.now - spec.config[:idle_timeout].to_i)
-        end
+        flush
       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
+          self.shard = Shard.default if self.shard == shard
           # 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
@@ -63,29 +56,9 @@ module Switchman
         end
       end
 
-      def clear_idle_connections!(since_when)
-        synchronize do
-          @connections.reject! do |conn|
-            if conn.last_query_at < since_when && !conn.in_use?
-              conn.disconnect!
-              true
-            else
-              false
-            end
-          end
-          @available.clear
-          @connections.each do |conn|
-            @available.add conn
-          end
-        end
-      end
-
       def switch_database(conn)
-        if !@schemas && conn.adapter_name == 'PostgreSQL' && !self.shard.database_server.config[:shard_name]
-          @schemas = conn.current_schemas
-        end
+        @schemas = conn.current_schemas if !@schemas && conn.adapter_name == 'PostgreSQL' && !shard.database_server.config[:shard_name]
 
-        spec.config[:shard_name] = self.shard.name
         conn.shard = shard
       end
 

data/lib/switchman/active_record/database_configurations.rb

@@ -0,0 +1,34 @@
+# frozen_string_literal: true
+
+module Switchman
+  module ActiveRecord
+    module DatabaseConfigurations
+      private
+
+      # key difference: assumes a hybrid two-tier structure; each third tier
+      # is implicitly named, and their config is constructing by merging into
+      # its parent
+      def build_configs(configs)
+        return configs.configurations if configs.is_a?(DatabaseConfigurations)
+        return configs if configs.is_a?(Array)
+
+        db_configs = configs.flat_map do |env_name, config|
+          roles = config.keys.select { |k| config[k].is_a?(Hash) }
+          base_config = config.except(*roles)
+
+          name = "#{env_name}/primary"
+          name = 'primary' if env_name == default_env
+          base_db = build_db_config_from_raw_config(env_name, name, base_config)
+          [base_db] + roles.map do |role|
+            build_db_config_from_raw_config(env_name, "#{env_name}/#{role}",
+                                            base_config.merge(config[role]).merge(replica: true))
+          end
+        end
+
+        db_configs << environment_url_config(default_env, 'primary', {}) unless db_configs.find(&:for_current_env?)
+
+        merge_db_environment_variables(default_env, db_configs.compact)
+      end
+    end
+  end
+end

data/lib/switchman/active_record/database_configurations/database_config.rb

@@ -0,0 +1,13 @@
+# frozen_string_literal: true
+
+module Switchman
+  module ActiveRecord
+    module DatabaseConfigurations
+      module DatabaseConfig
+        def for_current_env?
+          true
+        end
+      end
+    end
+  end
+end

data/lib/switchman/active_record/finder_methods.rb

@@ -7,19 +7,18 @@ module Switchman
         return super(id) unless klass.integral_id?
 
         if shard_source_value != :implicit
-          current_shard = Shard.current(klass.shard_category)
-          result = self.activate do |relation, shard|
+          current_shard = Shard.current(klass.connection_classes)
+          result = activate do |relation, shard|
             current_id = Shard.relative_id_for(id, current_shard, shard)
             # current_id will be nil for non-integral id
             next unless current_id
             # skip the shard if the object can't be on it. unless we're only looking at one shard;
             # we might be expecting a shadow object
-            next if current_id > Shard::IDS_PER_SHARD && self.all_shards.length > 1
+            next if current_id > Shard::IDS_PER_SHARD && all_shards.length > 1
+
             relation.call_super(:find_one, FinderMethods, current_id)
           end
-          if result.is_a?(Array)
-            result = result.first
-          end
+          result = result.first if result.is_a?(Array)
           # we may have skipped all shards
           raise_record_not_found_exception!(id, 0, 1) unless result
           return result
@@ -34,8 +33,8 @@ module Switchman
       end
 
       def find_some_ordered(ids)
-        current_shard = Shard.current(klass.shard_category)
-        ids = ids.map{|id| Shard.relative_id_for(id, current_shard, current_shard)}
+        current_shard = Shard.current(klass.connection_classes)
+        ids = ids.map { |id| Shard.relative_id_for(id, current_shard, current_shard) }
         super(ids)
       end
 
@@ -45,14 +44,12 @@ module Switchman
 
       def exists?(conditions = :none)
         conditions = conditions.id if ::ActiveRecord::Base === conditions
-        return false if !conditions
+        return false unless conditions
 
-        relation = ::Rails.version >= "5.2" ?
-          apply_join_dependency(eager_loading: false) :
-          apply_join_dependency(self, construct_join_dependency)
+        relation = apply_join_dependency(eager_loading: false)
         return false if ::ActiveRecord::NullRelation === relation
 
-        relation = relation.except(:select, :order).select("1 AS one").limit(1)
+        relation = relation.except(:select, :order).select('1 AS one').limit(1)
 
         case conditions
         when Array, Hash
@@ -62,9 +59,7 @@ module Switchman
         end
 
         relation.activate do |shard_rel|
-          return true if ::Rails.version >= "5.2" ?
-            connection.select_value(shard_rel.arel, "#{name} Exists") :
-            connection.select_value(shard_rel, "#{name} Exists", shard_rel.bound_attributes)
+          return true if connection.select_value(shard_rel.arel, "#{name} Exists")
         end
         false
       end

data/lib/switchman/active_record/log_subscriber.rb

@@ -14,20 +14,16 @@ module Switchman
 
         name  = "#{payload[:name]} (#{event.duration.round(1)}ms)"
         name  = "CACHE #{name}" if payload[:cached]
-        sql   = payload[:sql].squeeze(' '.freeze)
+        sql   = payload[:sql].squeeze(' ')
         binds = nil
         shard = payload[:shard]
         shard = "  [#{shard[:database_server_id]}:#{shard[:id]} #{shard[:env]}]" if shard
 
         unless (payload[:binds] || []).empty?
-          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|
+          casted_params = type_casted_binds(payload[:type_casted_binds])
+          binds = '  ' + payload[:binds].zip(casted_params).map do |attr, value|
             render_bind(attr, value)
-          }.inspect
+          end.inspect
         end
 
         name = colorize_payload_name(name, payload[:name])

data/lib/switchman/active_record/migration.rb

@@ -4,41 +4,44 @@ module Switchman
   module ActiveRecord
     module Migration
       module Compatibility
-        module V5_0
+        module V5_0 # rubocop:disable Naming/ClassAndModuleCamelCase
           def create_table(*args, **options)
-            unless options.key?(:id)
-              options[:id] = :bigserial
-            end
-            if block_given?
-              super do |td|
-                yield td
-              end
-            else
-              super
-            end
+            options[:id] = :bigserial unless options.key?(:id)
+            super
           end
         end
       end
 
       def connection
         conn = super
-        if conn.shard != ::ActiveRecord::Base.connection_pool.current_pool.shard
-          ::ActiveRecord::Base.connection_pool.current_pool.switch_database(conn)
-        end
+        ::ActiveRecord::Base.connection_pool.switch_database(conn) if conn.shard != ::ActiveRecord::Base.connection_pool.shard
         conn
       end
     end
 
     module Migrator
+      # significant change: hash shard name, not database name
       def generate_migrator_advisory_lock_id
-        shard_name_hash = Zlib.crc32("#{Shard.current.id}:#{Shard.current.name}")
+        shard_name_hash = Zlib.crc32(Shard.current.name)
         ::ActiveRecord::Migrator::MIGRATOR_SALT * shard_name_hash
       end
+
+      # significant change: strip out prefer_secondary from config
+      def with_advisory_lock_connection
+        pool = ::ActiveRecord::ConnectionAdapters::ConnectionHandler.new.establish_connection(
+          ::ActiveRecord::Base.connection_db_config.configuration_hash.except(:prefer_secondary)
+        )
+
+        pool.with_connection { |connection| yield(connection) } # rubocop:disable Style/ExplicitBlockArgument
+      ensure
+        pool&.disconnect!
+      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

data/lib/switchman/active_record/model_schema.rb

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

data/lib/switchman/active_record/persistence.rb

@@ -5,14 +5,12 @@ module Switchman
     module Persistence
       # touch reads the id attribute directly, so it's not relative to the current shard
       def touch(*, **)
-        shard.activate(self.class.shard_category) { super }
+        shard.activate(self.class.connection_classes) { super }
       end
 
-      if ::Rails.version >= '5.2'
-        def update_columns(*)
-          shard.activate(self.class.shard_category) { super }
-        end
+      def update_columns(*)
+        shard.activate(self.class.connection_classes) { super }
       end
     end
   end
-end
+end

data/lib/switchman/active_record/postgresql_adapter.rb

@@ -9,22 +9,22 @@ module Switchman
 
         option_string = options.sum do |key, value|
           case key
-            when :owner
-              " OWNER = \"#{value}\""
-            when :template
-              " TEMPLATE = \"#{value}\""
-            when :encoding
-              " ENCODING = '#{value}'"
-            when :collation
-              " LC_COLLATE = '#{value}'"
-            when :ctype
-              " LC_CTYPE = '#{value}'"
-            when :tablespace
-              " TABLESPACE = \"#{value}\""
-            when :connection_limit
-              " CONNECTION LIMIT = #{value}"
-            else
-              ""
+          when :owner
+            " OWNER = \"#{value}\""
+          when :template
+            " TEMPLATE = \"#{value}\""
+          when :encoding
+            " ENCODING = '#{value}'"
+          when :collation
+            " LC_COLLATE = '#{value}'"
+          when :ctype
+            " LC_CTYPE = '#{value}'"
+          when :tablespace
+            " TABLESPACE = \"#{value}\""
+          when :connection_limit
+            " CONNECTION LIMIT = #{value}"
+          else
+            ''
           end
         end
 
@@ -37,99 +37,39 @@ module Switchman
       end
 
       def current_schemas
-        select_values("SELECT * FROM unnest(current_schemas(false))")
-      end
-
-      def tables(name = nil)
-        query(<<-SQL, 'SCHEMA').map { |row| row[0] }
-          SELECT tablename
-          FROM pg_tables
-          WHERE schemaname = '#{shard.name}'
-        SQL
+        select_values('SELECT * FROM unnest(current_schemas(false))')
       end
 
       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.instance_variable_set(:@schema, shard.name) if string && !name.schema
         [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
-          name.instance_variable_set(:@schema, shard.name)
-        end
-
-        select_values(<<-SQL, 'SCHEMA').any?
-          SELECT c.relname
-          FROM pg_class c
-          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 = '#{shard.name}'
-        SQL
-      end
-
-      def indexes(table_name)
-        result = query(<<-SQL, 'SCHEMA')
-           SELECT distinct i.relname, d.indisunique, d.indkey, pg_get_indexdef(d.indexrelid), t.oid
-           FROM pg_class t
-           INNER JOIN pg_index d ON t.oid = d.indrelid
-           INNER JOIN pg_class i ON d.indexrelid = i.oid
-           WHERE i.relkind = 'i'
-             AND d.indisprimary = 'f'
-             AND t.relname = '#{table_name}'
-             AND i.relnamespace IN (SELECT oid FROM pg_namespace WHERE nspname = '#{shard.name}' )
-          ORDER BY i.relname
-        SQL
-
-
-        result.map do |row|
-          index_name = row[0]
-          unique = row[1] == true || row[1] == 't'
-          indkey = row[2].split(" ")
-          inddef = row[3]
-          oid = row[4]
-
-          columns = Hash[query(<<-SQL, "SCHEMA")]
-          SELECT a.attnum, a.attname
-          FROM pg_attribute a
-          WHERE a.attrelid = #{oid}
-          AND a.attnum IN (#{indkey.join(",")})
-          SQL
-
-          column_names = columns.stringify_keys.values_at(*indkey).compact
-
-          unless column_names.empty?
-            # add info on sort order for columns (only desc order is explicitly specified, asc is the default)
-            desc_order_columns = inddef.scan(/(\w+) DESC/).flatten
-            orders = desc_order_columns.any? ? Hash[desc_order_columns.map {|order_column| [order_column, :desc]}] : {}
-            where = inddef.scan(/WHERE (.+)$/).flatten[0]
-            using = inddef.scan(/USING (.+?) /).flatten[0].to_sym
-
-            if ::Rails.version >= "5.2"
-              ::ActiveRecord::ConnectionAdapters::IndexDefinition.new(table_name, index_name, unique, column_names, orders: orders, where: where, using: using)
-            else
-              ::ActiveRecord::ConnectionAdapters::IndexDefinition.new(table_name, index_name, unique, column_names, [], orders, where, nil, using)
-            end
+      # significant change: use the shard name if no explicit schema
+      def quoted_scope(name = nil, type: nil)
+        schema, name = extract_schema_qualified_name(name)
+        type = \
+          case type # rubocop:disable Style/HashLikeCase
+          when 'BASE TABLE'
+            "'r','p'"
+          when 'VIEW'
+            "'v','m'"
+          when 'FOREIGN TABLE'
+            "'f'"
           end
-        end.compact
+        scope = {}
+        scope[:schema] = quote(schema || shard.name)
+        scope[:name] = quote(name) if name
+        scope[:type] = type if type
+        scope
       end
 
-      def index_name_exists?(table_name, index_name, _default = nil)
-        exec_query(<<-SQL, 'SCHEMA').rows.first[0].to_i > 0
-            SELECT COUNT(*)
-            FROM pg_class t
-            INNER JOIN pg_index d ON t.oid = d.indrelid
-            INNER JOIN pg_class i ON d.indexrelid = i.oid
-            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 = '#{shard.name}' )
-        SQL
+      def foreign_keys(table_name)
+        super.each do |fk|
+          to_table_qualified_name = ::ActiveRecord::ConnectionAdapters::PostgreSQL::Utils.extract_schema_qualified_name(fk.to_table)
+          fk.to_table = to_table_qualified_name.identifier if to_table_qualified_name.schema == shard.name
+        end
       end
 
       def quote_local_table_name(name)
@@ -141,14 +81,13 @@ 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
-          name.instance_variable_set(:@schema, shard.name)
-        end
+        name.instance_variable_set(:@schema, shard.name) unless name.schema
         name.quoted
       end
 
-      def with_local_table_name(enable = true)
+      def with_local_table_name(enable = true) # rubocop:disable Style/OptionalBooleanParameter
         old_value = @use_local_table_name
         @use_local_table_name = enable
         yield
@@ -156,48 +95,10 @@ module Switchman
         @use_local_table_name = old_value
       end
 
-      def foreign_keys(table_name)
-
-        # mostly copy-pasted from AR - only change is to the nspname condition for qualified names support
-        fk_info = select_all <<-SQL.strip_heredoc
-          SELECT t2.oid::regclass::text AS to_table, a1.attname AS column, a2.attname AS primary_key, c.conname AS name, c.confupdtype AS on_update, c.confdeltype AS on_delete
-          FROM pg_constraint c
-          JOIN pg_class t1 ON c.conrelid = t1.oid
-          JOIN pg_class t2 ON c.confrelid = t2.oid
-          JOIN pg_attribute a1 ON a1.attnum = c.conkey[1] AND a1.attrelid = t1.oid
-          JOIN pg_attribute a2 ON a2.attnum = c.confkey[1] AND a2.attrelid = t2.oid
-          JOIN pg_namespace t3 ON c.connamespace = t3.oid
-          WHERE c.contype = 'f'
-            AND t1.relname = #{quote(table_name)}
-            AND t3.nspname = '#{shard.name}'
-          ORDER BY c.conname
-        SQL
-
-        fk_info.map do |row|
-          options = {
-            column: row['column'],
-            name: row['name'],
-            primary_key: row['primary_key']
-          }
-
-          options[:on_delete] = extract_foreign_key_action(row['on_delete'])
-          options[:on_update] = extract_foreign_key_action(row['on_update'])
-
-          # 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 to_table_qualified_name.schema == shard.name
-            to_table = to_table_qualified_name.identifier
-          end
-
-          ::ActiveRecord::ConnectionAdapters::ForeignKeyDefinition.new(table_name, to_table, options)
-        end
-      end
-
       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]
+        index, algorithm, if_not_exists = super
+        algorithm = nil if DatabaseServer.creating_new_shard && algorithm == 'CONCURRENTLY'
+        [index, algorithm, if_not_exists]
       end
 
       def rename_table(table_name, new_name)