switchman 2.1.0 → 3.0.6

data/lib/switchman/active_record/base.rb

@@ -6,49 +6,40 @@ module Switchman
       module ClassMethods
         delegate :shard, to: :all
 
-        def find_ids_in_ranges(opts={}, &block)
-          opts.reverse_merge!(:loose => true)
+        def find_ids_in_ranges(opts = {}, &block)
+          opts.reverse_merge!(loose: true)
           all.find_ids_in_ranges(opts, &block)
         end
 
-        def shard_category
-          connection_specification_name.to_sym
-        end
+        def sharded_model
+          self.abstract_class = true
 
-        def shard_category=(category)
-          categories = Shard.const_get(:CATEGORIES)
-          if categories[shard_category]
-            categories[shard_category].delete(self)
-            categories.delete(shard_category) if categories[shard_category].empty?
-          end
-          categories[category] ||= []
-          categories[category] << self
-          self.connection_specification_name = category.to_s
+          return if self == UnshardedRecord
+
+          Shard.send(:add_sharded_model, self)
         end
 
         def integral_id?
-          if @integral_id == nil
-            @integral_id = columns_hash[primary_key]&.type == :integer
-          end
+          @integral_id = columns_hash[primary_key]&.type == :integer if @integral_id.nil?
           @integral_id
         end
 
         def transaction(**)
           if self != ::ActiveRecord::Base && current_scope
             current_scope.activate do
-              db = Shard.current(shard_category).database_server
-              if ::GuardRail.environment != db.guard_rail_environment
-                db.unguard { super }
-              else
+              db = Shard.current(connection_classes).database_server
+              if ::GuardRail.environment == db.guard_rail_environment
                 super
+              else
+                db.unguard { super }
               end
             end
           else
-            db = Shard.current(shard_category).database_server
-            if ::GuardRail.environment != db.guard_rail_environment
-              db.unguard { super }
-            else
+            db = Shard.current(connection_classes).database_server
+            if ::GuardRail.environment == db.guard_rail_environment
               super
+            else
+              db.unguard { super }
             end
           end
         end
@@ -72,53 +63,70 @@ module Switchman
         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
+          ::ActiveRecord::Base.connection_handler.connection_pool_list.each do |pool|
+            pool.connection(switch_shard: false).clear_query_cache if pool.active_connection?
           end
         end
-      end
 
-      def self.included(klass)
-        klass.extend(ClassMethods)
-        klass.set_callback(:initialize, :before) do
-          unless @shard
-            if self.class.sharded_primary_key?
-              @shard = Shard.shard_for(self[self.class.primary_key], Shard.current(self.class.shard_category))
-            else
-              @shard = Shard.current(self.class.shard_category)
-            end
+        # significant change: _don't_ check if klasses.include?(Base)
+        # i.e. other sharded models don't inherit the current shard of Base
+        def current_shard
+          connected_to_stack.reverse_each do |hash|
+            return hash[:shard] if hash[:shard] && hash[:klasses].include?(connection_classes)
           end
+
+          default_shard
         end
+
+        def current_switchman_shard
+          connected_to_stack.reverse_each do |hash|
+            return hash[:switchman_shard] if hash[:switchman_shard] && hash[:klasses].include?(connection_classes)
+          end
+
+          Shard.default
+        end
+      end
+
+      def self.prepended(klass)
+        klass.singleton_class.prepend(ClassMethods)
+      end
+
+      def _run_initialize_callbacks
+        @shard ||= if self.class.sharded_primary_key?
+                     Shard.shard_for(self[self.class.primary_key], Shard.current(self.class.connection_classes))
+                   else
+                     Shard.current(self.class.connection_classes)
+                   end
+        super
       end
 
       def shard
-        @shard || Shard.current(self.class.shard_category) || Shard.default
+        @shard || Shard.current(self.class.connection_classes) || Shard.default
       end
 
       def shard=(new_shard)
-        raise ::ActiveRecord::ReadOnlyRecord if !self.new_record? || @shard_set_in_stone
-        if shard != new_shard
-          attributes.each do |attr, value|
-            self[attr] = Shard.relative_id_for(value, shard, new_shard) if self.class.sharded_column?(attr)
-          end
-          @shard = new_shard
+        raise ::ActiveRecord::ReadOnlyRecord if !new_record? || @shard_set_in_stone
+
+        return if shard == new_shard
+
+        attributes.each do |attr, value|
+          self[attr] = Shard.relative_id_for(value, shard, new_shard) if self.class.sharded_column?(attr)
         end
+        @shard = new_shard
       end
 
       def save(*, **)
         @shard_set_in_stone = true
-        (self.class.current_scope || self.class.default_scoped).shard(shard, :implicit).scoping { super }
+        super
       end
 
       def save!(*, **)
         @shard_set_in_stone = true
-        (self.class.current_scope || self.class.default_scoped).shard(shard, :implicit).scoping { super }
+        super
       end
 
       def destroy
-        self.class.shard(shard, :implicit).scoping { super }
+        shard.activate(self.class.connection_classes) { super }
       end
 
       def clone
@@ -126,23 +134,29 @@ module Switchman
         # TODO: adjust foreign keys
         # don't use the setter, cause the foreign keys are already
         # relative to this shard
-        result.instance_variable_set(:@shard, self.shard)
+        result.instance_variable_set(:@shard, shard)
         result
       end
 
       def transaction(**kwargs, &block)
-        shard.activate(self.class.shard_category) do
+        shard.activate(self.class.connection_classes) do
           self.class.transaction(**kwargs, &block)
         end
       end
 
+      def with_transaction_returning_status
+        shard.activate(self.class.connection_classes) do
+          super
+        end
+      end
+
       def hash
-        self.class.sharded_primary_key? ? self.class.hash ^ Shard.global_id_for(id).hash : super
+        self.class.sharded_primary_key? ? self.class.hash ^ global_id.hash : super
       end
 
       def to_param
         short_id = Shard.short_id_for(id)
-        short_id && short_id.to_s
+        short_id&.to_s
       end
 
       def initialize_dup(*args)
@@ -151,16 +165,18 @@ module Switchman
         copy
       end
 
-      def quoted_id
-        return super unless self.class.sharded_primary_key?
-        # do this the Rails 4.2 way, so that if Shard.current != self.shard, the id gets transposed
-        self.class.connection.quote(id)
+      def update_columns(*)
+        db = shard.database_server
+        return db.unguard { super } if ::GuardRail.environment != db.guard_rail_environment
+
+        super
       end
 
-      def update_columns(*)
-        db = Shard.current(self.class.shard_category).database_server
-        if ::GuardRail.environment != db.guard_rail_environment
-          return db.unguard { super }
+      def id_for_database
+        if self.class.sharded_primary_key?
+          # It's an int, so so it's safe to just return it without passing it through anything else
+          # In theory we should do `@attributes[@primary_key].type.serialize(id)`, but that seems to have surprising side-effects
+          id
         else
           super
         end
@@ -168,22 +184,22 @@ module Switchman
 
       protected
 
-      # see also AttributeMethods#shard_category_code_for_reflection
-      def shard_category_for_reflection(reflection)
+      # see also AttributeMethods#connection_classes_code_for_reflection
+      def connection_classes_for_reflection(reflection)
         if reflection
           if reflection.options[:polymorphic]
             begin
-              read_attribute(reflection.foreign_type)&.constantize&.shard_category || :primary
+              read_attribute(reflection.foreign_type)&.constantize&.connection_classes
             rescue NameError
               # in case someone is abusing foreign_type to not point to an actual class
-              :primary
+              ::ActiveRecord::Base
             end
           else
             # otherwise we can just return a symbol for the statically known type of the association
-            reflection.klass.shard_category
+            reflection.klass.connection_classes
           end
         else
-          shard_category
+          self.class.connection_classes
         end
       end
     end

data/lib/switchman/active_record/calculations.rb

@@ -3,45 +3,44 @@
 module Switchman
   module ActiveRecord
     module Calculations
-
       def pluck(*column_names)
-        target_shard = Shard.current(klass.shard_category)
+        target_shard = Shard.current(klass.connection_classes)
         shard_count = 0
-        result = self.activate do |relation, shard|
+        result = activate do |relation, shard|
           shard_count += 1
           results = relation.call_super(:pluck, Calculations, *column_names)
           if column_names.length > 1
             column_names.each_with_index do |column_name, idx|
-              if klass.sharded_column?(column_name)
-                results.each{|result| result[idx] = Shard.relative_id_for(result[idx], shard, target_shard)}
+              next unless klass.sharded_column?(column_name)
+
+              results.each do |r|
+                r[idx] = Shard.relative_id_for(r[idx], shard, target_shard)
               end
             end
-          else
-            if klass.sharded_column?(column_names.first.to_s)
-              results = results.map{|result| Shard.relative_id_for(result, shard, target_shard)}
-            end
+          elsif klass.sharded_column?(column_names.first.to_s)
+            results = results.map { |r| Shard.relative_id_for(r, shard, target_shard) }
           end
           results
         end
-        if distinct_value && shard_count > 1
-          result.uniq!
-        end
+        result.uniq! if distinct_value && shard_count > 1
         result
       end
 
       def execute_simple_calculation(operation, column_name, distinct)
         operation = operation.to_s.downcase
-        if operation == "average"
+        if operation == 'average'
           result = calculate_simple_average(column_name, distinct)
         else
-          result = self.activate{ |relation| relation.call_super(:execute_simple_calculation, Calculations, operation, column_name, distinct) }
+          result = activate do |relation|
+            relation.call_super(:execute_simple_calculation, Calculations, operation, column_name, distinct)
+          end
           if result.is_a?(Array)
             case operation
-            when "count", "sum"
+            when 'count', 'sum'
               result = result.sum
-            when "minimum"
+            when 'minimum'
               result = result.min
-            when "maximum"
+            when 'maximum'
               result = result.max
             end
           end
@@ -53,18 +52,20 @@ module Switchman
         # See activerecord#execute_simple_calculation
         relation = except(:order)
         column = aggregate_column(column_name)
-        relation.select_values = [operation_over_aggregate_column(column, "average", distinct).as("average"),
-                                  operation_over_aggregate_column(column, "count", distinct).as("count")]
+        relation.select_values = [operation_over_aggregate_column(column, 'average', distinct).as('average'),
+                                  operation_over_aggregate_column(column, 'count', distinct).as('count')]
 
-        initial_results = relation.activate{ |rel| klass.connection.select_all(rel) }
+        initial_results = relation.activate { |rel| klass.connection.select_all(rel) }
         if initial_results.is_a?(Array)
           initial_results.each do |r|
-            r["average"] = type_cast_calculated_value(r["average"], type_for(column_name), "average")
-            r["count"] = type_cast_calculated_value(r["count"], type_for(column_name), "count")
+            r['average'] = type_cast_calculated_value_switchman(r['average'], column_name, 'average')
+            r['count'] = type_cast_calculated_value_switchman(r['count'], column_name, 'count')
           end
-          result = initial_results.map{|r| r["average"] * r["count"]}.sum / initial_results.map{|r| r["count"]}.sum
+          result = initial_results.map { |r| r['average'] * r['count'] }.sum / initial_results.map do |r|
+                                                                                 r['count']
+                                                                               end.sum
         else
-          result = type_cast_calculated_value(initial_results.first["average"], type_for(column_name), "average")
+          result = type_cast_calculated_value_switchman(initial_results.first['average'], column_name, 'average')
         end
         result
       end
@@ -74,27 +75,28 @@ module Switchman
         opts = grouped_calculation_options(operation.to_s.downcase, column_name, distinct)
 
         relation = build_grouped_calculation_relation(opts)
-        target_shard = Shard.current(:primary)
+        target_shard = Shard.current
 
         rows = relation.activate do |rel, shard|
           calculated_data = klass.connection.select_all(rel)
 
           if opts[:association]
             key_ids     = calculated_data.collect { |row| row[opts[:group_aliases].first] }
-            key_records = opts[:association].klass.base_class.where(:id => key_ids)
-            key_records = Hash[key_records.map { |r| [Shard.relative_id_for(r, shard, target_shard), r] }]
+            key_records = opts[:association].klass.base_class.where(id: key_ids)
+            key_records = key_records.map { |r| [Shard.relative_id_for(r, shard, target_shard), r] }.to_h
           end
 
           calculated_data.map do |row|
-            row[opts[:aggregate_alias]] = type_cast_calculated_value(
-                row[opts[:aggregate_alias]], type_for(opts[:column_name]), opts[:operation])
+            row[opts[:aggregate_alias]] = type_cast_calculated_value_switchman(
+              row[opts[:aggregate_alias]], column_name, opts[:operation]
+            )
             row['count'] = row['count'].to_i if opts[:operation] == 'average'
 
-            opts[:group_columns].each do |aliaz, type, column_name|
+            opts[:group_columns].each do |aliaz, _type, group_column_name|
               if opts[:associated] && (aliaz == opts[:group_aliases].first)
                 row[aliaz] = key_records[Shard.relative_id_for(row[aliaz], shard, target_shard)]
-              elsif column_name && @klass.sharded_column?(column_name)
-                row[aliaz] = Shard.relative_id_for(type_cast_calculated_value(row[aliaz], type), shard, target_shard)
+              elsif group_column_name && @klass.sharded_column?(group_column_name)
+                row[aliaz] = Shard.relative_id_for(row[aliaz], shard, target_shard)
               end
             end
             row
@@ -106,12 +108,21 @@ module Switchman
 
       private
 
+      def type_cast_calculated_value_switchman(value, column_name, operation)
+        type_cast_calculated_value(value, operation) do |val|
+          column = aggregate_column(column_name)
+          type ||= column.try(:type_caster) ||
+                   lookup_cast_type_from_join_dependencies(column_name.to_s) || Type.default_value
+          type.deserialize(val)
+        end
+      end
+
       def column_name_for(field)
         field.respond_to?(:name) ? field.name.to_s : field.to_s.split('.').last
       end
 
       def grouped_calculation_options(operation, column_name, distinct)
-        opts = {:operation => operation, :column_name => column_name, :distinct => distinct}
+        opts = { operation: operation, column_name: column_name, distinct: distinct }
 
         opts[:aggregate_alias] = aggregate_alias_for(operation, column_name)
         group_attrs = group_values
@@ -125,12 +136,12 @@ module Switchman
 
         # 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|
+        group_columns = group_aliases.zip(group_fields).map do |aliaz, field|
           [aliaz, type_for(field), column_name_for(field)]
-        }
-        opts.merge!(:association => association, :associated => associated,
-         :group_aliases => group_aliases, :group_columns => group_columns,
-         :group_fields => group_fields)
+        end
+        opts.merge!(association: association, associated: associated,
+                    group_aliases: group_aliases, group_columns: group_columns,
+                    group_fields: group_fields)
 
         opts
       end
@@ -149,28 +160,30 @@ module Switchman
         group = opts[:group_fields]
 
         select_values = [
-            operation_over_aggregate_column(
-                aggregate_column(opts[:column_name]),
-                opts[:operation],
-                opts[:distinct]).as(opts[:aggregate_alias])
+          operation_over_aggregate_column(
+            aggregate_column(opts[:column_name]),
+            opts[:operation],
+            opts[:distinct]
+          ).as(opts[:aggregate_alias])
         ]
-        if opts[:operation ]== 'average'
+        if opts[:operation] == 'average'
           # include count in average so we can recalculate the average
           # across all shards if needed
           select_values << operation_over_aggregate_column(
-              aggregate_column(opts[:column_name]),
-              'count', opts[:distinct]).as('count')
+            aggregate_column(opts[:column_name]),
+            'count', opts[:distinct]
+          ).as('count')
         end
 
         haves = having_clause.send(:predicates)
         select_values += select_values unless haves.empty?
-        select_values.concat opts[:group_fields].zip(opts[:group_aliases]).map { |field,aliaz|
+        select_values.concat(opts[:group_fields].zip(opts[:group_aliases]).map do |field, aliaz|
           if field.respond_to?(:as)
             field.as(aliaz)
           else
             "#{field} AS #{aliaz}"
           end
-        }
+        end)
 
         relation = except(:group)
         relation.group_values = group
@@ -181,12 +194,12 @@ module Switchman
       def compact_grouped_calculation_rows(rows, opts)
         result = ::ActiveSupport::OrderedHash.new
         rows.each do |row|
-          key = opts[:group_columns].map { |aliaz, column| row[aliaz] }
+          key = opts[:group_columns].map { |aliaz, _column| row[aliaz] }
           key = key.first if key.size == 1
           value = row[opts[:aggregate_alias]]
 
           if opts[:operation] == 'average'
-            if result.has_key?(key)
+            if result.key?(key)
               old_value, old_count = result[key]
               new_count = old_count + row['count']
               new_value = ((old_value * old_count) + (value * row['count'])) / new_count
@@ -194,30 +207,24 @@ module Switchman
             else
               result[key] = [value, row['count']]
             end
-          else
-            if result.has_key?(key)
-              case opts[:operation]
-              when "count", "sum"
-                result[key] += value
-              when "minimum"
-                result[key] = value if value < result[key]
-              when "maximum"
-                result[key] = value if value > result[key]
-              end
-            else
-              result[key] = value
+          elsif result.key?(key)
+            case opts[:operation]
+            when 'count', 'sum'
+              result[key] += value
+            when 'minimum'
+              result[key] = value if value < result[key]
+            when 'maximum'
+              result[key] = value if value > result[key]
             end
+          else
+            result[key] = value
           end
         end
 
-        if opts[:operation] == 'average'
-          result = Hash[result.map{|k, v| [k, v.first]}]
-        end
+        result.transform_values!(&:first) if opts[:operation] == 'average'
 
         result
       end
-
-
     end
   end
 end

data/lib/switchman/active_record/connection_pool.rb

@@ -5,16 +5,7 @@ require 'switchman/errors'
 module Switchman
   module ActiveRecord
     module ConnectionPool
-      def shard
-        Thread.current[tls_key] || Shard.default
-      end
-
-      def shard=(value)
-        Thread.current[tls_key] = value
-      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,75 +13,37 @@ 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 = current_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
+        raise NonExistentShardError if current_shard.new_record?
+
+        switch_database(conn) if conn.shard != current_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
-      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|
-            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
+        flush
       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' && !current_shard.database_server.config[:shard_name]
 
-        spec.config[:shard_name] = self.shard.name
-        conn.shard = shard
+        conn.shard = current_shard
       end
 
       private
 
+      def current_shard
+        connection_klass.current_switchman_shard
+      end
+
       def tls_key
         "#{object_id}_shard".to_sym
       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/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