switchman 2.0.5 → 2.0.10

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: c544d2ed8498a6aecf33823f39b304393a65f213f6a68bf04f4d36610ebe3ac2
-  data.tar.gz: be685b315997b103dd13376f5223196f1b66a5be490502aac6bc401408899abc
+  metadata.gz: bce5f0819ee2570b0bf51ba3872af1f34c23273712cc33c35559225723fb13a0
+  data.tar.gz: 11e7c3d612263f94ea17bc5213580ae42e2c12743a13ff36e9f988feddda0a7e
 SHA512:
-  metadata.gz: fbdd769d90c831e1bbe6dd39def12775537989346daae7dbbab3103a330b89151e9d1c3b9fe216ab9ade75e3c01ee1b1c91062cc622707c3c2697f7dc3239ea8
-  data.tar.gz: 27d32eefa842a76debb1a1fc7875407b4df1398a45a71419c7cc13045fd8705ff5e4af724478a9cb854d517433ae8d16420baa13fc8e7459ab1953e626c0c80a
+  metadata.gz: 0fad57a01e86f34533c5cf539cf92360605b129bdfdf77878fec72663c3783bad56e7a7a324f3b388e37a89296d05f309dc2c84084e30c1c44d20917e15d35e3
+  data.tar.gz: 002bcec40325a407c8be94432d2396c8ae812ad2828695266e591956d1e1e8195628906fce710c4681a3c72cc5f6ea8f1859a99aac22cbd3462be5dd6118bbdf

data/app/models/switchman/shard.rb

@@ -718,6 +718,7 @@ module Switchman
         Switchman.cache.delete(['shard', id].join('/'))
         Switchman.cache.delete("default_shard") if default?
       end
+      self.class.clear_cache
     end
 
     def default_name

data/db/migrate/20180828183945_add_default_shard_index.rb

@@ -10,6 +10,6 @@ class AddDefaultShardIndex < ActiveRecord::Migration[4.2]
               else
                 {}
               end
-    add_index :switchman_shards, :default, options
+    add_index :switchman_shards, :default, **options
   end
 end

data/lib/switchman.rb

@@ -17,4 +17,6 @@ module Switchman
   def self.cache=(cache)
     @cache = cache
   end
+
+  class OrderOnMultiShardQuery < RuntimeError; end
 end

data/lib/switchman/active_record/association.rb

@@ -87,7 +87,7 @@ module Switchman
           # 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
+              global_key = if model.shard_category == :unsharded
                              convert_key(record[association_key_name])
                            else
                              Shard.global_id_for(record[association_key_name], record.shard)

data/lib/switchman/active_record/calculations.rb

@@ -51,7 +51,7 @@ module Switchman
 
       def calculate_simple_average(column_name, distinct)
         # See activerecord#execute_simple_calculation
-        relation = reorder(nil)
+        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")]

data/lib/switchman/active_record/connection_handler.rb

@@ -135,7 +135,7 @@ module Switchman
             primary_pool = retrieve_connection_pool("primary")
             if primary_pool.is_a?(ConnectionPoolProxy)
               pool = ConnectionPoolProxy.new(spec_name.to_sym, primary_pool.default_pool, @shard_connection_pools)
-              pool.schema_cache.copy_values(primary_pool.schema_cache)
+              pool.schema_cache.copy_references(primary_pool.schema_cache)
               pool
             else
               primary_pool

data/lib/switchman/active_record/migration.rb

@@ -30,10 +30,42 @@ module Switchman
     end
 
     module Migrator
+      # significant change: hash shard id, 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
+
+      if ::Rails.version >= '6.0'
+        # copy/paste from Rails 6.1
+        def with_advisory_lock
+          lock_id = generate_migrator_advisory_lock_id
+
+          with_advisory_lock_connection do |connection|
+            got_lock = connection.get_advisory_lock(lock_id)
+            raise ::ActiveRecord::ConcurrentMigrationError unless got_lock
+            load_migrated # reload schema_migrations to be sure it wasn't changed by another process before we got the lock
+            yield
+          ensure
+            if got_lock && !connection.release_advisory_lock(lock_id)
+              raise ::ActiveRecord::ConcurrentMigrationError.new(
+                ::ActiveRecord::ConcurrentMigrationError::RELEASE_LOCK_FAILED_MESSAGE
+              )
+            end
+          end
+        end
+
+        # significant change: strip out prefer_secondary from config
+        def with_advisory_lock_connection
+          pool = ::ActiveRecord::ConnectionAdapters::ConnectionHandler.new.establish_connection(
+            ::ActiveRecord::Base.connection_config.except(:prefer_secondary)
+          )
+
+          pool.with_connection { |connection| yield(connection) }
+        ensure
+          pool&.disconnect!
+        end
+      end
     end
 
     module MigrationContext

data/lib/switchman/active_record/postgresql_adapter.rb

@@ -40,14 +40,6 @@ module Switchman
         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
-      end
-
       def extract_schema_qualified_name(string)
         name = ::ActiveRecord::ConnectionAdapters::PostgreSQL::Utils.extract_schema_qualified_name(string.to_s)
         if string && !name.schema
@@ -56,80 +48,155 @@ module Switchman
         [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
+      # 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
+          when "BASE TABLE"
+            "'r','p'"
+          when "VIEW"
+            "'v','m'"
+          when "FOREIGN TABLE"
+            "'f'"
+          end
+        scope = {}
+        scope[:schema] = quote(schema || shard.name)
+        scope[:name] = quote(name) if name
+        scope[:type] = type if type
+        scope
       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(",")})
+      if ::Rails.version < '6.0'
+        def tables(name = nil)
+          query(<<-SQL, 'SCHEMA').map { |row| row[0] }
+            SELECT tablename
+            FROM pg_tables
+            WHERE schemaname = '#{shard.name}'
           SQL
+        end
 
-          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
+        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
-        end.compact
-      end
 
-      def index_name_exists?(table_name, index_name, _default = nil)
-        exec_query(<<-SQL, 'SCHEMA').rows.first[0].to_i > 0
-            SELECT COUNT(*)
+          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 i.relname = '#{index_name}'
+              AND d.indisprimary = 'f'
               AND t.relname = '#{table_name}'
               AND i.relnamespace IN (SELECT oid FROM pg_namespace WHERE nspname = '#{shard.name}' )
-        SQL
+            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
+            end
+          end.compact
+        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
+        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
+      else
+        def foreign_keys(table_name)
+          super.each do |fk|
+            to_table_qualified_name = ::ActiveRecord::ConnectionAdapters::PostgreSQL::Utils.extract_schema_qualified_name(fk.to_table)
+            if to_table_qualified_name.schema == shard.name
+              fk.to_table = to_table_qualified_name.identifier
+            end
+          end
+        end
       end
 
       def quote_local_table_name(name)
@@ -156,44 +223,6 @@ 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"

data/lib/switchman/active_record/query_methods.rb

@@ -78,10 +78,6 @@ module Switchman
         end
       end
 
-      def or(other)
-        super(other.shard(self.primary_shard))
-      end
-
       private
 
       if ::Rails.version >= '5.2'
@@ -246,135 +242,92 @@ module Switchman
                                binds: nil,
                                dup_binds_on_mutation: false)
         result = predicates.map do |predicate|
-          transposed, binds = transpose_single_predicate(predicate, source_shard, target_shard, remove_nonlocal_primary_keys,
-                                     binds: binds, dup_binds_on_mutation: dup_binds_on_mutation)
-          transposed
-        end
-        result = [result, binds]
-        result
-      end
+          next predicate unless predicate.is_a?(::Arel::Nodes::Binary)
+          next predicate unless predicate.left.is_a?(::Arel::Attributes::Attribute)
+          relation, column = relation_and_column(predicate.left)
+          next predicate unless (type = transposable_attribute_type(relation, column))
 
-      def transpose_single_predicate(predicate,
-                                     source_shard,
-                                     target_shard,
-                                     remove_nonlocal_primary_keys = false,
-                                     binds: nil,
-                                     dup_binds_on_mutation: false)
-        if predicate.is_a?(::Arel::Nodes::Grouping)
-          return predicate, binds unless predicate.expr.is_a?(::Arel::Nodes::Or)
-          # Dang, we have an OR.  OK, that means we have other epxressions below this
-          # level, perhaps many, that may need transposition.
-          # the left side and right side must each be treated as predicate lists and
-          # transformed in kind, if neither of them changes we can just return the grouping as is.
-          # hold on, it's about to get recursive...
-          #
-          # TODO: "binds" is getting passed up and down
-          # this stack purely because of the necessary handling for rails <5.2
-          #  Dropping support for 5.2 means we can remove the "binds" argument from
-          # all of this and yank the conditional below where we monkey with their instance state.
-          or_expr = predicate.expr
-          left_node = or_expr.left
-          right_node = or_expr.right
-          left_predicates = left_node.children
-          right_predicates = right_node.children
-          new_left_predicates, binds = transpose_predicates(left_predicates, source_shard,
-                                                               target_shard, remove_nonlocal_primary_keys,
-                                                               binds: binds, dup_binds_on_mutation: dup_binds_on_mutation)
-          new_right_predicates, binds = transpose_predicates(right_predicates, source_shard,
-                                                               target_shard, remove_nonlocal_primary_keys,
-                                                               binds: binds, dup_binds_on_mutation: dup_binds_on_mutation)
-          if new_left_predicates != left_predicates
-            left_node.instance_variable_set(:@children, new_left_predicates)
-          end
-          if new_right_predicates != right_predicates
-            right_node.instance_variable_set(:@children, new_right_predicates)
-          end
-          return predicate, binds
-        end
-        return predicate, binds unless predicate.is_a?(::Arel::Nodes::Binary)
-        return predicate, binds unless predicate.left.is_a?(::Arel::Attributes::Attribute)
-        relation, column = relation_and_column(predicate.left)
-        return predicate, binds unless (type = transposable_attribute_type(relation, column))
+          remove = true if type == :primary &&
+              remove_nonlocal_primary_keys &&
+              predicate.left.relation.model == klass &&
+              (predicate.is_a?(::Arel::Nodes::Equality) || predicate.is_a?(::Arel::Nodes::In))
 
-        remove = true if type == :primary &&
-            remove_nonlocal_primary_keys &&
-            predicate.left.relation.model == klass &&
-            predicate.is_a?(::Arel::Nodes::Equality)
-
-        current_source_shard =
-            if source_shard
-              source_shard
-            elsif type == :primary
-              Shard.current(klass.shard_category)
-            elsif type == :foreign
-              source_shard_for_foreign_key(relation, column)
-            end
+          current_source_shard =
+              if source_shard
+                source_shard
+              elsif type == :primary
+                Shard.current(klass.shard_category)
+              elsif type == :foreign
+                source_shard_for_foreign_key(relation, column)
+              end
 
-        if ::Rails.version >= "5.2"
-          new_right_value =
-            case predicate.right
+          if ::Rails.version >= "5.2"
+            new_right_value =
+              case predicate.right
+              when Array
+                predicate.right.map {|val| transpose_predicate_value(val, current_source_shard, target_shard, type, remove).presence }.compact
+              else
+                transpose_predicate_value(predicate.right, current_source_shard, target_shard, type, remove)
+              end
+          else
+            new_right_value = case predicate.right
             when Array
-              predicate.right.map {|val| transpose_predicate_value(val, current_source_shard, target_shard, type, remove) }
-            else
-              transpose_predicate_value(predicate.right, current_source_shard, target_shard, type, remove)
-            end
-        else
-          new_right_value = case predicate.right
-          when Array
-            local_ids = []
-            predicate.right.each do |value|
-              local_id = Shard.relative_id_for(value, current_source_shard, target_shard)
-              next unless local_id
-              unless remove && local_id > Shard::IDS_PER_SHARD
-                if value.is_a?(::Arel::Nodes::Casted)
-                  if local_id == value.val
-                    local_id = value
-                  elsif local_id != value
-                    local_id = value.class.new(local_id, value.attribute)
+              local_ids = []
+              predicate.right.each do |value|
+                local_id = Shard.relative_id_for(value, current_source_shard, target_shard)
+                next unless local_id
+                unless remove && local_id > Shard::IDS_PER_SHARD
+                  if value.is_a?(::Arel::Nodes::Casted)
+                    if local_id == value.val
+                      local_id = value
+                    elsif local_id != value
+                      local_id = value.class.new(local_id, value.attribute)
+                    end
                   end
+                  local_ids << local_id
                 end
-                local_ids << local_id
-              end
-            end
-            local_ids
-          when ::Arel::Nodes::BindParam
-            # look for a bind param with a matching column name
-            if binds && bind = binds.detect{|b| b&.name.to_s == predicate.left.name.to_s}
-              # before we mutate, dup
-              if dup_binds_on_mutation
-                binds = binds.map(&:dup)
-                dup_binds_on_mutation = false
-                bind = binds.find { |b| b&.name.to_s == predicate.left.name.to_s }
               end
-              if bind.value.is_a?(::ActiveRecord::StatementCache::Substitute)
-                bind.value.sharded = true # mark for transposition later
-                bind.value.primary = true if type == :primary
-              else
-                local_id = Shard.relative_id_for(bind.value, current_source_shard, target_shard)
-                local_id = [] if remove && local_id > Shard::IDS_PER_SHARD
-                bind.instance_variable_set(:@value, local_id)
-                bind.instance_variable_set(:@value_for_database, nil)
+              local_ids
+            when ::Arel::Nodes::BindParam
+              # look for a bind param with a matching column name
+              if binds && bind = binds.detect{|b| b&.name.to_s == predicate.left.name.to_s}
+                # before we mutate, dup
+                if dup_binds_on_mutation
+                  binds = binds.map(&:dup)
+                  dup_binds_on_mutation = false
+                  bind = binds.find { |b| b&.name.to_s == predicate.left.name.to_s }
+                end
+                if bind.value.is_a?(::ActiveRecord::StatementCache::Substitute)
+                  bind.value.sharded = true # mark for transposition later
+                  bind.value.primary = true if type == :primary
+                else
+                  local_id = Shard.relative_id_for(bind.value, current_source_shard, target_shard)
+                  local_id = [] if remove && local_id > Shard::IDS_PER_SHARD
+                  bind.instance_variable_set(:@value, local_id)
+                  bind.instance_variable_set(:@value_for_database, nil)
+                end
               end
+              predicate.right
+            else
+              local_id = Shard.relative_id_for(predicate.right, current_source_shard, target_shard) || predicate.right
+              local_id = [] if remove && local_id.is_a?(Integer) && local_id > Shard::IDS_PER_SHARD
+              local_id
             end
-            predicate.right
-          else
-            local_id = Shard.relative_id_for(predicate.right, current_source_shard, target_shard) || predicate.right
-            local_id = [] if remove && local_id.is_a?(Integer) && local_id > Shard::IDS_PER_SHARD
-            local_id
           end
-        end
-        out_predicate = if new_right_value == predicate.right
-          predicate
-        elsif predicate.right.is_a?(::Arel::Nodes::Casted)
-          if new_right_value == predicate.right.val
+          if new_right_value == predicate.right
             predicate
+          elsif predicate.right.is_a?(::Arel::Nodes::Casted)
+            if new_right_value == predicate.right.val
+              predicate
+            else
+              predicate.class.new(predicate.left, predicate.right.class.new(new_right_value, predicate.right.attribute))
+            end
           else
-            predicate.class.new(predicate.left, predicate.right.class.new(new_right_value, predicate.right.attribute))
+            predicate.class.new(predicate.left, new_right_value)
           end
-        else
-          predicate.class.new(predicate.left, new_right_value)
         end
-        return out_predicate, binds
+        result = [result, binds]
+        result
       end
 
       def transpose_predicate_value(value, current_shard, target_shard, attribute_type, remove_non_local_ids)
@@ -387,7 +340,7 @@ module Switchman
             value
           else
             local_id = Shard.relative_id_for(current_id, current_shard, target_shard) || current_id
-            local_id = [] if remove_non_local_ids && local_id.is_a?(Integer) && local_id > Shard::IDS_PER_SHARD
+            return nil if remove_non_local_ids && local_id.is_a?(Integer) && local_id > Shard::IDS_PER_SHARD
             if current_id != local_id
               # make a new bind param
               ::Arel::Nodes::BindParam.new(query_att.class.new(query_att.name, local_id, query_att.type))

data/lib/switchman/active_record/relation.rb

@@ -106,10 +106,62 @@ module Switchman
             shards.first.activate(klass.shard_category) { yield(self, shards.first) }
           end
         else
-          # TODO: implement local limit to avoid querying extra shards
-          Shard.with_each_shard(shards, [klass.shard_category]) do
-            shard(Shard.current(klass.shard_category), :to_a).activate(&block)
+          result_count = 0
+          can_order = false
+          result = Shard.with_each_shard(shards, [klass.shard_category]) do
+            # don't even query other shards if we're already past the limit
+            next if limit_value && result_count >= limit_value && order_values.empty?
+
+            relation = shard(Shard.current(klass.shard_category), :to_a)
+            # do a minimal query if possible
+            relation = relation.limit(limit_value - result_count) if limit_value && !result_count.zero? && order_values.empty?
+
+            shard_results = relation.activate(&block)
+
+            if shard_results.present?
+              can_order ||= can_order_cross_shard_results? unless order_values.empty?
+              raise OrderOnMultiShardQuery if !can_order && !order_values.empty? && result_count.positive?
+
+              result_count += shard_results.is_a?(Array) ? shard_results.length : 1
+            end
+            shard_results
+          end
+
+          result = reorder_cross_shard_results(result) if can_order
+          result.slice!(limit_value..-1) if limit_value
+          result
+        end
+      end
+
+      def can_order_cross_shard_results?
+        # we only presume to be able to post-sort the most basic of orderings
+        order_values.all? { |ov| ov.is_a?(::Arel::Nodes::Ordering) && ov.expr.is_a?(::Arel::Attributes::Attribute) }
+      end
+
+      def reorder_cross_shard_results(results)
+        results.sort! do |l, r|
+          result = 0
+          order_values.each do |ov|
+            if l.respond_to?(ov.expr.name)
+              a = l.send(ov.expr.name)
+              b = r.send(ov.expr.name)
+            else
+              a = l.attributes[ov.expr.name]
+              b = r.attributes[ov.expr.name]
+            end
+            next if a == b
+
+            if a.nil? || b.nil?
+              result = 1 if a.nil?
+              result *= -1 if ov.is_a?(::Arel::Nodes::Descending)
+            else
+              result = a <=> b
+            end
+
+            result *= -1 if ov.is_a?(::Arel::Nodes::Descending)
+            break unless result.zero?
           end
+          result
         end
       end
     end

data/lib/switchman/connection_pool_proxy.rb

@@ -84,6 +84,10 @@ module Switchman
       @schema_cache
     end
 
+    def set_schema_cache(cache)
+      @schema_cache.copy_values(cache)
+    end
+
     %w{release_connection
        disconnect!
        flush!

data/lib/switchman/schema_cache.rb

@@ -5,14 +5,22 @@ module Switchman
     delegate :connection, to: :pool
     attr_reader :pool
 
+    SHARED_IVS = %i{@columns @columns_hash @primary_keys @data_sources @indexes}.freeze
+
     def initialize(pool)
       @pool = pool
       super(nil)
     end
 
     def copy_values(other_cache)
+      SHARED_IVS.each do |iv|
+        instance_variable_get(iv).replace(other_cache.instance_variable_get(iv))
+      end
+    end
+
+    def copy_references(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|
+      SHARED_IVS.each do |iv|
         instance_variable_set(iv, other_cache.instance_variable_get(iv))
       end
     end

data/lib/switchman/standard_error.rb

@@ -8,7 +8,7 @@ module Switchman
     end
 
     def current_shard(category = :primary)
-      @active_shards[category] || Shard.default
+      @active_shards&.[](category) || Shard.default
     end
   end
 end

data/lib/switchman/test_helper.rb

@@ -24,7 +24,7 @@ module Switchman
         else
           server1 = Shard.default.database_server
         end
-        server2 = DatabaseServer.create(Shard.default.database_server.config)
+        server2 = DatabaseServer.create(Shard.default.database_server.config.merge(server2: true))
 
         if server1 == Shard.default.database_server && server1.config[:shard1] && server1.config[:shard2]
           # look for the shards in the db already

data/lib/switchman/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Switchman
-  VERSION = "2.0.5"
+  VERSION = "2.0.10"
 end

metadata

@@ -1,16 +1,16 @@
 --- !ruby/object:Gem::Specification
 name: switchman
 version: !ruby/object:Gem::Version
-  version: 2.0.5
+  version: 2.0.10
 platform: ruby
 authors:
 - Cody Cutrer
 - James Williams
 - Jacob Fugal
-autorequire: 
+autorequire:
 bindir: bin
 cert_chain: []
-date: 2021-03-02 00:00:00.000000000 Z
+date: 2021-06-09 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: railties
@@ -257,7 +257,7 @@ homepage: http://www.instructure.com/
 licenses:
 - MIT
 metadata: {}
-post_install_message: 
+post_install_message:
 rdoc_options: []
 require_paths:
 - lib
@@ -272,8 +272,8 @@ required_rubygems_version: !ruby/object:Gem::Requirement
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.0.3
-signing_key: 
+rubygems_version: 3.2.15
+signing_key:
 specification_version: 4
 summary: Rails sharding magic
 test_files: []