switchman 2.0.2 → 3.0.0

data/lib/switchman/rails.rb

@@ -1,18 +1,20 @@
 # frozen_string_literal: true
 
-module Switchman::Rails
-  module ClassMethods
-    def self.prepended(klass)
-      # in Rails 4+, the Rails.cache= method was used during bootstrap to set
-      # Rails.cache(_without_sharding) to the value from the config file. but now
-      # that that's done (the bootstrap happened before this module is included
-      # into Rails), we want to make sure no one tries to assign to Rails.cache,
-      # because it would be wrong w.r.t. sharding.
-      klass.send(:remove_method, :cache=)
-    end
+module Switchman
+  module Rails
+    module ClassMethods
+      def self.prepended(klass)
+        # in Rails 4+, the Rails.cache= method was used during bootstrap to set
+        # Rails.cache(_without_sharding) to the value from the config file. but now
+        # that that's done (the bootstrap happened before this module is included
+        # into Rails), we want to make sure no one tries to assign to Rails.cache,
+        # because it would be wrong w.r.t. sharding.
+        klass.send(:remove_method, :cache=)
+      end
 
-    def cache
-      Switchman::Shard.current.database_server.cache_store
+      def cache
+        Switchman::Shard.current.database_server.cache_store
+      end
     end
   end
 end

data/lib/switchman/schema_cache.rb

@@ -12,7 +12,7 @@ module Switchman
 
     def copy_values(other_cache)
       # use the same cached values but still fall back to the correct pool
-      [:@columns, :@columns_hash, :@primary_keys, :@data_sources].each do |iv|
+      %i[@columns @columns_hash @primary_keys @data_sources].each do |iv|
         instance_variable_set(iv, other_cache.instance_variable_get(iv))
       end
     end

data/lib/switchman/sharded_instrumenter.rb

@@ -7,7 +7,7 @@ module Switchman
       @shard_host = shard_host
     end
 
-    def instrument(name, payload={})
+    def instrument(name, payload = {})
       shard = @shard_host&.shard
       # attribute_methods_generated? will be false during a reload -
       # when we might be doing a query while defining attribute methods,

data/lib/switchman/standard_error.rb

@@ -3,12 +3,24 @@
 module Switchman
   module StandardError
     def initialize(*args)
-      @active_shards = Shard.send(:active_shards).dup
+      # Shard.current can throw this when switchman isn't working right; if we try to
+      # do our stuff here, it'll cause a SystemStackError, which is a pain to deal with
+      if is_a?(::ActiveRecord::ConnectionNotEstablished)
+        super
+        return
+      end
+
+      if defined?(Shard)
+        @active_shards = Shard.sharded_models.map do |klass|
+          [klass, Shard.current(klass)]
+        end.compact.to_h
+      end
+
       super
     end
 
-    def current_shard(category = :primary)
-      @active_shards[category] || Shard.default
+    def current_shard(klass = ::ActiveRecord::Base)
+      @active_shards&.[](klass) || Shard.default
     end
   end
 end

data/lib/switchman/test_helper.rb

@@ -11,19 +11,14 @@ module Switchman
             Shard.create!(default: true)
           rescue
             raise unless dont_create
+
             # database doesn't exist yet, presumably cause we're creating it right now
             return [nil, nil]
           end
           Shard.default(reload: true)
         end
 
-        # can't auto-create a new shard on the default shard's db server if the
-        # default shard is split across multiple db servers
-        if ::ActiveRecord::Base.connection_handler.connection_pool_list.length > 1
-          server1 = DatabaseServer.create(Shard.default.database_server.config)
-        else
-          server1 = Shard.default.database_server
-        end
+        server1 = Shard.default.database_server
         server2 = DatabaseServer.create(Shard.default.database_server.config)
 
         if server1 == Shard.default.database_server && server1.config[:shard1] && server1.config[:shard2]
@@ -55,8 +50,8 @@ module Switchman
             shard1.destroy
             shard2.drop_database rescue nil
             shard2.destroy
-            shard1 = server1.create_new_shard(:name => server1.config[:shard1])
-            shard2 = server2.create_new_shard(:name => server1.config[:shard2])
+            shard1 = server1.create_new_shard(name: server1.config[:shard1])
+            shard2 = server2.create_new_shard(name: server1.config[:shard2])
           end
           [shard1, shard2]
         else
@@ -65,11 +60,12 @@ module Switchman
       end
 
       private
+
       def find_existing_test_shard(server, name)
         if server == Shard.default.database_server
           server.shards.where(name: name).first
         else
-          shard = Shard.where("database_server_id IS NOT NULL AND name=?", name).first
+          shard = Shard.where('database_server_id IS NOT NULL AND name=?', name).first
           # if somehow databases got created in a different order, change the shard to match
           shard.database_server = server if shard
           shard

data/lib/switchman/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Switchman
-  VERSION = "2.0.2"
+  VERSION = '3.0.0'
 end

data/lib/tasks/switchman.rake

@@ -1,20 +1,22 @@
+# frozen_string_literal: true
+
 module Switchman
   module Rake
     def self.filter_database_servers(&block)
       chain = filter_database_servers_chain # use a local variable so that the current chain is closed over in the following lambda
-      @filter_database_servers_chain = lambda { |servers| block.call(servers, chain) }
+      @filter_database_servers_chain = ->(servers) { block.call(servers, chain) }
     end
 
     def self.scope(base_scope = Shard,
-      database_server: ENV['DATABASE_SERVER'],
-      shard: ENV['SHARD'])
+                   database_server: ENV['DATABASE_SERVER'],
+                   shard: ENV['SHARD'])
       servers = DatabaseServer.all
 
       if database_server
         servers = database_server
         if servers.first == '-'
           negative = true
-          servers = servers[1..-1]
+          servers = servers[1..]
         end
         servers = servers.split(',')
         open = servers.delete('open')
@@ -31,16 +33,14 @@ module Switchman
 
       servers = filter_database_servers_chain.call(servers)
 
-      scope = base_scope.order(::Arel.sql("database_server_id IS NOT NULL, database_server_id, id"))
+      scope = base_scope.order(::Arel.sql('database_server_id IS NOT NULL, database_server_id, id'))
       if servers != DatabaseServer.all
-        conditions = ["database_server_id IN (?)", servers.map(&:id)]
-        conditions.first << " OR database_server_id IS NULL" if servers.include?(Shard.default.database_server)
+        conditions = ['database_server_id IN (?)', servers.map(&:id)]
+        conditions.first << ' OR database_server_id IS NULL' if servers.include?(Shard.default.database_server)
         scope = scope.where(conditions)
       end
 
-      if shard
-        scope = shard_scope(scope, shard)
-      end
+      scope = shard_scope(scope, shard) if shard
 
       scope
     end
@@ -49,11 +49,11 @@ module Switchman
       { parallel: ENV['PARALLEL'].to_i, max_procs: ENV['MAX_PARALLEL_PROCS'] }
     end
 
-    # categories - an array or proc, to activate as the current shard during the
+    # classes - an array or proc, to activate as the current shard during the
     # task. tasks which modify the schema may want to pass all categories in
     # so that schema updates for non-default tables happen against all shards.
     # this is handled automatically for the default migration tasks, below.
-    def self.shardify_task(task_name, categories: [:primary])
+    def self.shardify_task(task_name, classes: [::ActiveRecord::Base])
       old_task = ::Rake::Task[task_name]
       old_actions = old_task.actions.dup
       old_task.actions.clear
@@ -66,43 +66,30 @@ module Switchman
 
         ::GuardRail.activate(:deploy) do
           Shard.default.database_server.unguard do
-            begin
-              categories = categories.call if categories.respond_to?(:call)
-              Shard.with_each_shard(scope, categories, options) do
-                shard = Shard.current
-                puts "#{shard.id}: #{shard.description}"
-                ::ActiveRecord::Base.connection_pool.spec.config[:shard_name] = Shard.current.name
-                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.unguard do
-                  old_actions.each { |action| action.call(*task_args) }
-                end
-                nil
+            classes = classes.call if classes.respond_to?(:call)
+            Shard.with_each_shard(scope, classes, **options) do
+              shard = Shard.current
+              puts "#{shard.id}: #{shard.description}"
+
+              shard.database_server.unguard do
+                old_actions.each { |action| action.call(*task_args) }
               end
-            rescue => e
-              puts "Exception from #{e.current_shard.id}: #{e.current_shard.description}" if options[:parallel] != 0
-              raise
+              nil
             end
+          rescue => e
+            puts "Exception from #{e.current_shard.id}: #{e.current_shard.description}" if options[:parallel] != 0
+            raise
+
+            #::ActiveRecord::Base.configurations = old_configurations
           end
         end
       end
     end
 
-    %w{db:migrate db:migrate:up db:migrate:down db:rollback}.each do |task_name|
-      shardify_task(task_name, categories: ->{ Shard.categories })
+    %w[db:migrate db:migrate:up db:migrate:down db:rollback].each do |task_name|
+      shardify_task(task_name, classes: -> { Shard.sharded_models })
     end
 
-    private
-
     def self.shard_scope(scope, raw_shard_ids)
       raw_shard_ids = raw_shard_ids.split(',')
 
@@ -125,6 +112,7 @@ module Switchman
         when /^(-?)(\d+)?\.\.(\.)?(\d+)?$/
           negative, start, open, finish = $1.present?, $2, $3.present?, $4
           raise "Invalid shard id or range: #{id}" unless start || finish
+
           range = []
           range << "id>=#{start}" if start
           range << "id<#{'=' unless open}#{finish}" if finish
@@ -136,12 +124,12 @@ module Switchman
         when %r{^(-?\d+)/(\d+)$}
           numerator = $1.to_i
           denominator = $2.to_i
-          if numerator == 0 || numerator.abs > denominator
-            raise "Invalid fractional chunk: #{id}"
-          end
+          raise "Invalid fractional chunk: #{id}" if numerator.zero? || numerator.abs > denominator
+
           # one chunk means everything
           if denominator == 1
             next if numerator == 1
+
             return scope.none
           end
 
@@ -156,24 +144,25 @@ module Switchman
           select = []
           if index != 1
             subscope = subscope.offset(per_chunk * (index - 1))
-            select << "MIN(id) AS min_id"
+            select << 'MIN(id) AS min_id'
           end
           if index != denominator
             subscope = subscope.limit(per_chunk)
-            select << "MAX(id) AS max_id"
-          end
-
-          result = Shard.from(subscope).select(select.join(", ")).to_a.first
-          if index == 1
-            range = "id<=#{result['max_id']}"
-          elsif index == denominator
-            range = "id>=#{result['min_id']}"
-          else
-            range = "(id>=#{result['min_id']} AND id<=#{result['max_id']})"
+            select << 'MAX(id) AS max_id'
           end
 
-          (numerator < 0 ? negative_ranges : ranges) <<  range
-          else
+          result = Shard.from(subscope).select(select.join(', ')).to_a.first
+          range = case index
+                  when 1
+                    "id<=#{result['max_id']}"
+                  when denominator
+                    "id>=#{result['min_id']}"
+                  else
+                    "(id>=#{result['min_id']} AND id<=#{result['max_id']})"
+                  end
+
+          (numerator.negative? ? negative_ranges : ranges) << range
+        else
           raise "Invalid shard id or range: #{id}"
         end
       end
@@ -182,27 +171,24 @@ module Switchman
       negative_shard_ids.uniq!
       unless shard_ids.empty?
         shard_ids -= negative_shard_ids
-        if shard_ids.empty? && ranges.empty?
-          return scope.none
-        end
+        return scope.none if shard_ids.empty? && ranges.empty?
+
         # we already trimmed them all out; no need to make the server do it as well
         negative_shard_ids = [] if ranges.empty?
       end
 
       conditions = []
       positive_queries = []
-      unless ranges.empty?
-        positive_queries << ranges.join(" OR ")
-      end
+      positive_queries << ranges.join(' OR ') unless ranges.empty?
       unless shard_ids.empty?
-        positive_queries << "id IN (?)"
+        positive_queries << 'id IN (?)'
         conditions << shard_ids
       end
-      positive_query = positive_queries.join(" OR ")
+      positive_query = positive_queries.join(' OR ')
       scope = scope.where(positive_query, *conditions) unless positive_queries.empty?
 
-      scope = scope.where("NOT (#{negative_ranges.join(" OR")})") unless negative_ranges.empty?
-      scope = scope.where("id NOT IN (?)", negative_shard_ids) unless negative_shard_ids.empty?
+      scope = scope.where("NOT (#{negative_ranges.join(' OR')})") unless negative_ranges.empty?
+      scope = scope.where('id NOT IN (?)', negative_shard_ids) unless negative_shard_ids.empty?
       scope
     end
 
@@ -213,18 +199,17 @@ module Switchman
 
   module ActiveRecord
     module PostgreSQLDatabaseTasks
-      def structure_dump(filename, extra_flags=nil)
+      def structure_dump(filename, extra_flags = nil)
         set_psql_env
         args = ['-s', '-x', '-O', '-f', filename]
         args.concat(Array(extra_flags)) if extra_flags
-        search_path = configuration['schema_search_path']
         shard = Shard.current.name
         serialized_search_path = shard
         args << "--schema=#{Shellwords.escape(shard)}"
 
         args << configuration['database']
         run_cmd('pg_dump', args, 'dumping')
-        File.open(filename, "a") { |f| f << "SET search_path TO #{serialized_search_path};\n\n" }
+        File.open(filename, 'a') { |f| f << "SET search_path TO #{serialized_search_path};\n\n" }
       end
     end
   end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: switchman
 version: !ruby/object:Gem::Version
-  version: 2.0.2
+  version: 3.0.0
 platform: ruby
 authors:
 - Cody Cutrer
@@ -10,76 +10,76 @@ authors:
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2020-11-06 00:00:00.000000000 Z
+date: 2021-03-17 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
-  name: railties
+  name: activerecord
   requirement: !ruby/object:Gem::Requirement
     requirements:
     - - ">="
       - !ruby/object:Gem::Version
-        version: '5.1'
+        version: '6.1'
     - - "<"
       - !ruby/object:Gem::Version
-        version: '6.1'
+        version: '6.2'
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - ">="
       - !ruby/object:Gem::Version
-        version: '5.1'
+        version: '6.1'
     - - "<"
       - !ruby/object:Gem::Version
-        version: '6.1'
+        version: '6.2'
 - !ruby/object:Gem::Dependency
-  name: activerecord
+  name: guardrail
   requirement: !ruby/object:Gem::Requirement
     requirements:
-    - - ">="
-      - !ruby/object:Gem::Version
-        version: '5.1'
-    - - "<"
+    - - "~>"
       - !ruby/object:Gem::Version
-        version: '6.1'
+        version: 3.0.0
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
-    - - ">="
-      - !ruby/object:Gem::Version
-        version: '5.1'
-    - - "<"
+    - - "~>"
       - !ruby/object:Gem::Version
-        version: '6.1'
+        version: 3.0.0
 - !ruby/object:Gem::Dependency
-  name: guardrail
+  name: open4
   requirement: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: 2.0.0
+        version: 1.3.0
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: 2.0.0
+        version: 1.3.0
 - !ruby/object:Gem::Dependency
-  name: open4
+  name: railties
   requirement: !ruby/object:Gem::Requirement
     requirements:
-    - - "~>"
+    - - ">="
       - !ruby/object:Gem::Version
-        version: 1.3.0
+        version: '6.1'
+    - - "<"
+      - !ruby/object:Gem::Version
+        version: '6.2'
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
-    - - "~>"
+    - - ">="
       - !ruby/object:Gem::Version
-        version: 1.3.0
+        version: '6.1'
+    - - "<"
+      - !ruby/object:Gem::Version
+        version: '6.2'
 - !ruby/object:Gem::Dependency
   name: appraisal
   requirement: !ruby/object:Gem::Requirement
@@ -108,6 +108,20 @@ dependencies:
     - - ">="
       - !ruby/object:Gem::Version
         version: '0'
+- !ruby/object:Gem::Dependency
+  name: pg
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '1.2'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '1.2'
 - !ruby/object:Gem::Dependency
   name: pry
   requirement: !ruby/object:Gem::Requirement
@@ -123,61 +137,103 @@ dependencies:
       - !ruby/object:Gem::Version
         version: '0'
 - !ruby/object:Gem::Dependency
-  name: pg
+  name: rake
   requirement: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '0'
+        version: '13.0'
   type: :development
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '0'
+        version: '13.0'
+- !ruby/object:Gem::Dependency
+  name: rspec-mocks
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '3.5'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '3.5'
 - !ruby/object:Gem::Dependency
   name: rspec-rails
   requirement: !ruby/object:Gem::Requirement
     requirements:
-    - - '='
+    - - "~>"
       - !ruby/object:Gem::Version
-        version: 3.5.2
+        version: '4.0'
   type: :development
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
-    - - '='
+    - - "~>"
       - !ruby/object:Gem::Version
-        version: 3.5.2
+        version: '4.0'
 - !ruby/object:Gem::Dependency
-  name: simplecov
+  name: rubocop
   requirement: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '0.15'
+        version: '1.10'
   type: :development
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '0.15'
+        version: '1.10'
 - !ruby/object:Gem::Dependency
-  name: rake
+  name: rubocop-rake
   requirement: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '12.0'
+        version: '0.5'
   type: :development
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '12.0'
+        version: '0.5'
+- !ruby/object:Gem::Dependency
+  name: rubocop-rspec
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '2.2'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '2.2'
+- !ruby/object:Gem::Dependency
+  name: simplecov
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '0.15'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '0.15'
 description: Sharding
 email:
 - cody@instructure.com
@@ -187,6 +243,7 @@ extra_rdoc_files: []
 files:
 - Rakefile
 - app/models/switchman/shard.rb
+- app/models/switchman/unsharded_record.rb
 - db/migrate/20130328212039_create_switchman_shards.rb
 - db/migrate/20130328224244_create_default_shard.rb
 - db/migrate/20161206323434_add_back_default_string_limits_switchman.rb
@@ -199,10 +256,10 @@ files:
 - lib/switchman/active_record/association.rb
 - lib/switchman/active_record/attribute_methods.rb
 - lib/switchman/active_record/base.rb
-- lib/switchman/active_record/batches.rb
 - lib/switchman/active_record/calculations.rb
-- lib/switchman/active_record/connection_handler.rb
 - lib/switchman/active_record/connection_pool.rb
+- lib/switchman/active_record/database_configurations.rb
+- lib/switchman/active_record/database_configurations/database_config.rb
 - lib/switchman/active_record/finder_methods.rb
 - lib/switchman/active_record/log_subscriber.rb
 - lib/switchman/active_record/migration.rb
@@ -217,12 +274,11 @@ files:
 - lib/switchman/active_record/spawn_methods.rb
 - lib/switchman/active_record/statement_cache.rb
 - lib/switchman/active_record/table_definition.rb
+- lib/switchman/active_record/tasks/database_tasks.rb
 - lib/switchman/active_record/type_caster.rb
-- lib/switchman/active_record/where_clause_factory.rb
 - lib/switchman/active_support/cache.rb
 - lib/switchman/arel.rb
 - lib/switchman/call_super.rb
-- lib/switchman/connection_pool_proxy.rb
 - lib/switchman/database_server.rb
 - lib/switchman/default_shard.rb
 - lib/switchman/engine.rb
@@ -251,7 +307,7 @@ required_ruby_version: !ruby/object:Gem::Requirement
   requirements:
   - - ">="
     - !ruby/object:Gem::Version
-      version: '2.5'
+      version: '2.6'
 required_rubygems_version: !ruby/object:Gem::Requirement
   requirements:
   - - ">="