switchman 2.0.13 → 3.0.0

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,73 +4,32 @@ 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 id, not database name
       def generate_migrator_advisory_lock_id
-        shard_name_hash = Zlib.crc32(Shard.current.name)
+        shard_name_hash = Zlib.crc32("#{Shard.current.id}:#{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
       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,166 +37,92 @@ module Switchman
       end
 
       def current_schemas
-        select_values("SELECT * FROM unnest(current_schemas(false))")
+        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
-          name.instance_variable_set(:@schema, shard.name)
-        end
+        name.instance_variable_set(:@schema, shard.name) if string && !name.schema
         [name.schema, name.identifier]
       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
-          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
+      def view_exists?(name)
+        name = ::ActiveRecord::ConnectionAdapters::PostgreSQL::Utils.extract_schema_qualified_name(name.to_s)
+        return false unless name.identifier
+
+        name.instance_variable_set(:@schema, shard.name) unless name.schema
+
+        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
 
-      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}'
+      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')] # rubocop:disable Style/HashConversion
+          SELECT a.attnum, a.attname
+          FROM pg_attribute a
+          WHERE a.attrelid = #{oid}
+          AND a.attnum IN (#{indkey.join(',')})
           SQL
-        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
+          column_names = columns.stringify_keys.values_at(*indkey).compact
 
-          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
+          next if column_names.empty?
 
-        def indexes(table_name)
-          result = query(<<-SQL, 'SCHEMA')
-            SELECT distinct i.relname, d.indisunique, d.indkey, pg_get_indexdef(d.indexrelid), t.oid
+          # 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] }] : {} # rubocop:disable Style/HashConversion
+          where = inddef.scan(/WHERE (.+)$/).flatten[0]
+          using = inddef.scan(/USING (.+?) /).flatten[0].to_sym
+
+          ::ActiveRecord::ConnectionAdapters::IndexDefinition.new(table_name, index_name, unique, column_names,
+                                                                  orders: orders, where: where, using: using)
+        end.compact
+      end
+
+      def index_name_exists?(table_name, index_name, _default = nil)
+        exec_query(<<-SQL, 'SCHEMA').rows.first[0].to_i.positive?
+            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 d.indisprimary = 'f'
+              AND i.relname = '#{index_name}'
               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
-            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
+        SQL
       end
 
       def quote_local_table_name(name)
@@ -208,18 +134,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_global_table_name(&block)
-        with_local_table_name(false, &block)
-      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
@@ -227,10 +148,45 @@ 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)
+          to_table = to_table_qualified_name.identifier if to_table_qualified_name.schema == shard.name
+
+          ::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)