pgsync 0.5.3 → 0.6.2

data/lib/pgsync/table.rb

@@ -0,0 +1,28 @@
+# minimal class to keep schema and table name separate
+module PgSync
+  class Table
+    attr_reader :schema, :name
+
+    def initialize(schema, name)
+      @schema = schema
+      @name = name
+    end
+
+    def full_name
+      "#{schema}.#{name}"
+    end
+
+    def eql?(other)
+      other.schema == schema && other.name == name
+    end
+
+    # override hash when overriding eql?
+    def hash
+      [schema, name].hash
+    end
+
+    def to_s
+      full_name
+    end
+  end
+end

data/lib/pgsync/table_sync.rb

@@ -2,236 +2,204 @@ module PgSync
   class TableSync
     include Utils
 
-    def sync(config, table, opts, source_url, destination_url)
-      start_time = Time.now
-
-      source = DataSource.new(source_url, timeout: 0)
-      destination = DataSource.new(destination_url, timeout: 0)
+    attr_reader :source, :destination, :tasks, :opts, :resolver
+
+    def initialize(source:, destination:, tasks:, opts:, resolver:)
+      @source = source
+      @destination = destination
+      @tasks = tasks
+      @opts = opts
+      @resolver = resolver
+    end
 
-      from_connection = source.conn
-      to_connection = destination.conn
+    def perform
+      confirm_tables_exist(destination, tasks, "destination")
 
-      from_fields = source.columns(table)
-      to_fields = destination.columns(table)
-      shared_fields = to_fields & from_fields
-      extra_fields = to_fields - from_fields
-      missing_fields = from_fields - to_fields
+      add_columns
 
-      if opts[:no_sequences]
-        from_sequences = []
-        to_sequences = []
-      else
-        from_sequences = source.sequences(table, shared_fields)
-        to_sequences = destination.sequences(table, shared_fields)
-      end
+      show_notes
 
-      shared_sequences = to_sequences & from_sequences
-      extra_sequences = to_sequences - from_sequences
-      missing_sequences = from_sequences - to_sequences
-
-      sql_clause = String.new("")
-      sql_clause << " #{opts[:sql]}" if opts[:sql]
-
-      notes = []
-      notes << "Extra columns: #{extra_fields.join(", ")}" if extra_fields.any?
-      notes << "Missing columns: #{missing_fields.join(", ")}" if missing_fields.any?
-      notes << "Extra sequences: #{extra_sequences.join(", ")}" if extra_sequences.any?
-      notes << "Missing sequences: #{missing_sequences.join(", ")}" if missing_sequences.any?
+      # don't sync tables with no shared fields
+      # we show a warning message above
+      run_tasks(tasks.reject { |task| task.shared_fields.empty? })
+    end
 
-      return {status: "success", message: "No fields to copy"} if shared_fields.empty?
+    # TODO only query specific tables
+    # TODO add sequences, primary keys, etc
+    def add_columns
+      source_columns = columns(source)
+      destination_columns = columns(destination)
 
-      bad_fields = opts[:no_rules] ? [] : config["data_rules"]
-      primary_key = destination.primary_key(table)
-      copy_fields = shared_fields.map { |f| f2 = bad_fields.to_a.find { |bf, _| rule_match?(table, f, bf) }; f2 ? "#{apply_strategy(f2[1], table, f, primary_key)} AS #{quote_ident(f)}" : "#{quote_ident_full(table)}.#{quote_ident(f)}" }.join(", ")
-      fields = shared_fields.map { |f| quote_ident(f) }.join(", ")
+      tasks.each do |task|
+        task.from_columns = source_columns[task.table] || []
+        task.to_columns = destination_columns[task.table] || []
+      end
+    end
 
-      seq_values = {}
-      shared_sequences.each do |seq|
-        seq_values[seq] = source.last_value(seq)
+    def show_notes
+      # for tables
+      resolver.notes.each do |note|
+        warning note
       end
 
-      copy_to_command = "COPY (SELECT #{copy_fields} FROM #{quote_ident_full(table)}#{sql_clause}) TO STDOUT"
-      if opts[:in_batches]
-        raise Error, "Cannot use --overwrite with --in-batches" if opts[:overwrite]
-        raise Error, "No primary key" unless primary_key
+      # for columns and sequences
+      tasks.each do |task|
+        task.notes.each do |note|
+          warning "#{task_name(task)}: #{note}"
+        end
+      end
 
-        destination.truncate(table) if opts[:truncate]
+      # for non-deferrable constraints
+      if opts[:defer_constraints]
+        constraints = non_deferrable_constraints(destination)
+        constraints = tasks.flat_map { |t| constraints[t.table] || [] }
+        warning "Non-deferrable constraints: #{constraints.join(", ")}" if constraints.any?
+      end
+    end
 
-        from_max_id = source.max_id(table, primary_key)
-        to_max_id = destination.max_id(table, primary_key) + 1
+    def columns(data_source)
+      query = <<~SQL
+        SELECT
+          table_schema AS schema,
+          table_name AS table,
+          column_name AS column,
+          data_type AS type
+        FROM
+          information_schema.columns
+        ORDER BY 1, 2, 3
+      SQL
+      data_source.execute(query).group_by { |r| Table.new(r["schema"], r["table"]) }.map do |k, v|
+        [k, v.map { |r| {name: r["column"], type: r["type"]} }]
+      end.to_h
+    end
 
-        if to_max_id == 1
-          from_min_id = source.min_id(table, primary_key)
-          to_max_id = from_min_id if from_min_id > 0
-        end
+    def non_deferrable_constraints(data_source)
+      query = <<~SQL
+        SELECT
+          table_schema AS schema,
+          table_name AS table,
+          constraint_name
+        FROM
+          information_schema.table_constraints
+        WHERE
+          constraint_type = 'FOREIGN KEY' AND
+          is_deferrable = 'NO'
+      SQL
+      data_source.execute(query).group_by { |r| Table.new(r["schema"], r["table"]) }.map do |k, v|
+        [k, v.map { |r| r["constraint_name"] }]
+      end.to_h
+    end
 
-        starting_id = to_max_id
-        batch_size = opts[:batch_size]
+    def run_tasks(tasks, &block)
+      notices = []
+      failed_tables = []
 
-        i = 1
-        batch_count = ((from_max_id - starting_id + 1) / batch_size.to_f).ceil
+      spinners = TTY::Spinner::Multi.new(format: :dots, output: output)
+      task_spinners = {}
+      started_at = {}
 
-        while starting_id <= from_max_id
-          where = "#{quote_ident(primary_key)} >= #{starting_id} AND #{quote_ident(primary_key)} < #{starting_id + batch_size}"
-          log "    #{i}/#{batch_count}: #{where}"
+      start = lambda do |task, i|
+        message = ":spinner #{display_item(task)}"
+        spinner = spinners.register(message)
+        if opts[:in_batches]
+          # log instead of spin for non-tty
+          log message.sub(":spinner", "⠋")
+        else
+          spinner.auto_spin
+        end
+        task_spinners[task] = spinner
+        started_at[task] = Time.now
+      end
 
-          # TODO be smarter for advance sql clauses
-          batch_sql_clause = " #{sql_clause.length > 0 ? "#{sql_clause} AND" : "WHERE"} #{where}"
+      finish = lambda do |task, i, result|
+        spinner = task_spinners[task]
+        time = (Time.now - started_at[task]).round(1)
 
-          batch_copy_to_command = "COPY (SELECT #{copy_fields} FROM #{quote_ident_full(table)}#{batch_sql_clause}) TO STDOUT"
-          to_connection.copy_data "COPY #{quote_ident_full(table)} (#{fields}) FROM STDIN" do
-            from_connection.copy_data batch_copy_to_command do
-              while (row = from_connection.get_copy_data)
-                to_connection.put_copy_data(row)
-              end
-            end
+        message =
+          if result[:message]
+            "(#{result[:message].lines.first.to_s.strip})"
+          else
+            "- #{time}s"
           end
 
-          starting_id += batch_size
-          i += 1
+        notices.concat(result[:notices])
 
-          if opts[:sleep] && starting_id <= from_max_id
-            sleep(opts[:sleep])
-          end
+        if result[:status] == "success"
+          spinner.success(message)
+        else
+          spinner.error(message)
+          failed_tables << task_name(task)
+          fail_sync(failed_tables) if opts[:fail_fast]
         end
 
-        log # add extra line for spinner
-      elsif !opts[:truncate] && (opts[:overwrite] || opts[:preserve] || !sql_clause.empty?)
-        raise Error, "No primary key" unless primary_key
+        unless spinner.send(:tty?)
+          status = result[:status] == "success" ? "✔" : "✖"
+          log [status, display_item(task), message].join(" ")
+        end
+      end
 
-        # create a temp table
-        temp_table = "pgsync_#{rand(1_000_000_000)}"
-        to_connection.exec("CREATE TEMPORARY TABLE #{quote_ident_full(temp_table)} AS SELECT * FROM #{quote_ident_full(table)} WITH NO DATA")
+      options = {start: start, finish: finish}
 
-        # load data
-        to_connection.copy_data "COPY #{quote_ident_full(temp_table)} (#{fields}) FROM STDIN" do
-          from_connection.copy_data copy_to_command do
-            while (row = from_connection.get_copy_data)
-              to_connection.put_copy_data(row)
-            end
-          end
-        end
+      jobs = opts[:jobs]
+      if opts[:debug] || opts[:in_batches] || opts[:defer_constraints]
+        warning "--jobs ignored" if jobs
+        jobs = 0
+      end
 
-        if opts[:preserve]
-          # insert into
-          to_connection.exec("INSERT INTO #{quote_ident_full(table)} (SELECT * FROM #{quote_ident_full(temp_table)} WHERE NOT EXISTS (SELECT 1 FROM #{quote_ident_full(table)} WHERE #{quote_ident_full(table)}.#{quote_ident(primary_key)} = #{quote_ident_full(temp_table)}.#{quote_ident(primary_key)}))")
-        else
-          to_connection.transaction do
-            to_connection.exec("DELETE FROM #{quote_ident_full(table)} WHERE #{quote_ident(primary_key)} IN (SELECT #{quote_ident(primary_key)} FROM #{quote_ident_full(temp_table)})")
-            to_connection.exec("INSERT INTO #{quote_ident_full(table)} (SELECT * FROM #{quote_ident(temp_table)})")
-          end
-        end
+      if windows?
+        options[:in_threads] = jobs || 4
       else
-        destination.truncate(table)
-        to_connection.copy_data "COPY #{quote_ident_full(table)} (#{fields}) FROM STDIN" do
-          from_connection.copy_data copy_to_command do
-            while (row = from_connection.get_copy_data)
-              to_connection.put_copy_data(row)
-            end
-          end
-        end
-      end
-      seq_values.each do |seq, value|
-        to_connection.exec("SELECT setval(#{escape(seq)}, #{escape(value)})")
+        options[:in_processes] = jobs if jobs
       end
 
-      message = nil
-      message = notes.join(", ") if notes.any?
-
-      {status: "success", message: message, time: (Time.now - start_time).round(1)}
-    rescue => e
-      message =
-        case e
-        when PG::ConnectionBad
-          # likely fine to show simplified message here
-          # the full message will be shown when first trying to connect
-          "Connection failed"
-        when PG::Error
-          e.message.sub("ERROR:  ", "")
-        when Error
-          e.message
-        else
-          "#{e.class.name}: #{e.message}"
-        end
+      maybe_defer_constraints do
+        # could try to use `raise Parallel::Kill` to fail faster with --fail-fast
+        # see `fast_faster` branch
+        # however, need to make sure connections are cleaned up properly
+        Parallel.each(tasks, **options) do |task|
+          source.reconnect_if_needed
+          destination.reconnect_if_needed
 
-      {status: "error", message: message}
-    ensure
-      source.close if source
-      destination.close if destination
-    end
+          task.perform
+        end
+      end
 
-    private
+      notices.each do |notice|
+        warning notice
+      end
 
-    # TODO better performance
-    def rule_match?(table, column, rule)
-      regex = Regexp.new('\A' + Regexp.escape(rule).gsub('\*','[^\.]*') + '\z')
-      regex.match(column) || regex.match("#{table.split(".", 2)[-1]}.#{column}") || regex.match("#{table}.#{column}")
+      fail_sync(failed_tables) if failed_tables.any?
     end
 
-    # TODO wildcard rules
-    def apply_strategy(rule, table, column, primary_key)
-      if rule.is_a?(Hash)
-        if rule.key?("value")
-          escape(rule["value"])
-        elsif rule.key?("statement")
-          rule["statement"]
-        else
-          raise Error, "Unknown rule #{rule.inspect} for column #{column}"
+    def maybe_defer_constraints
+      if opts[:defer_constraints]
+        destination.transaction do
+          destination.execute("SET CONSTRAINTS ALL DEFERRED")
+
+          # create a transaction on the source
+          # to ensure we get a consistent snapshot
+          source.transaction do
+            yield
+          end
         end
       else
-        case rule
-        when "untouched"
-          quote_ident(column)
-        when "unique_email"
-          "'email' || #{quoted_primary_key(table, primary_key, rule)}::text || '@example.org'"
-        when "unique_phone"
-          "(#{quoted_primary_key(table, primary_key, rule)}::bigint + 1000000000)::text"
-        when "unique_secret"
-          "'secret' || #{quoted_primary_key(table, primary_key, rule)}::text"
-        when "random_int", "random_number"
-          "(RANDOM() * 100)::int"
-        when "random_date"
-          "date '1970-01-01' + (RANDOM() * 10000)::int"
-        when "random_time"
-          "NOW() - (RANDOM() * 100000000)::int * INTERVAL '1 second'"
-        when "random_ip"
-          "(1 + RANDOM() * 254)::int::text || '.0.0.1'"
-        when "random_letter"
-          "chr(65 + (RANDOM() * 26)::int)"
-        when "random_string"
-          "RIGHT(MD5(RANDOM()::text), 10)"
-        when "null", nil
-          "NULL"
-        else
-          raise Error, "Unknown rule #{rule} for column #{column}"
-        end
+        yield
       end
     end
 
-    def quoted_primary_key(table, primary_key, rule)
-      raise "Primary key required for this data rule: #{rule}" unless primary_key
-      "#{quote_ident_full(table)}.#{quote_ident(primary_key)}"
+    def fail_sync(failed_tables)
+      raise Error, "Sync failed for #{failed_tables.size} table#{failed_tables.size == 1 ? nil : "s"}: #{failed_tables.join(", ")}"
     end
 
-    def quote_ident_full(ident)
-      ident.split(".").map { |v| quote_ident(v) }.join(".")
-    end
-
-    def quote_ident(value)
-      PG::Connection.quote_ident(value)
-    end
-
-    def escape(value)
-      if value.is_a?(String)
-        "'#{quote_string(value)}'"
-      else
-        value
-      end
+    def display_item(item)
+      messages = []
+      messages << task_name(item)
+      messages << item.opts[:sql] if item.opts[:sql]
+      messages.join(" ")
     end
 
-    # activerecord
-    def quote_string(s)
-      s.gsub(/\\/, '\&\&').gsub(/'/, "''")
+    def windows?
+      Gem.win_platform?
     end
   end
 end

data/lib/pgsync/task.rb

@@ -0,0 +1,329 @@
+module PgSync
+  class Task
+    include Utils
+
+    attr_reader :source, :destination, :config, :table, :opts
+    attr_accessor :from_columns, :to_columns
+
+    def initialize(source:, destination:, config:, table:, opts:)
+      @source = source
+      @destination = destination
+      @config = config
+      @table = table
+      @opts = opts
+    end
+
+    def quoted_table
+      quote_ident_full(table)
+    end
+
+    def perform
+      with_notices do
+        handle_errors do
+          maybe_disable_triggers do
+            sync_data
+          end
+        end
+      end
+    end
+
+    def from_fields
+      @from_fields ||= from_columns.map { |c| c[:name] }
+    end
+
+    def to_fields
+      @to_fields ||= to_columns.map { |c| c[:name] }
+    end
+
+    def shared_fields
+      @shared_fields ||= to_fields & from_fields
+    end
+
+    def from_sequences
+      @from_sequences ||= opts[:no_sequences] ? [] : source.sequences(table, shared_fields)
+    end
+
+    def to_sequences
+      @to_sequences ||= opts[:no_sequences] ? [] : destination.sequences(table, shared_fields)
+    end
+
+    def shared_sequences
+      @shared_sequences ||= to_sequences & from_sequences
+    end
+
+    def notes
+      notes = []
+      if shared_fields.empty?
+        notes << "No fields to copy"
+      else
+        extra_fields = to_fields - from_fields
+        notes << "Extra columns: #{extra_fields.join(", ")}" if extra_fields.any?
+
+        missing_fields = from_fields - to_fields
+        notes << "Missing columns: #{missing_fields.join(", ")}" if missing_fields.any?
+
+        extra_sequences = to_sequences - from_sequences
+        notes << "Extra sequences: #{extra_sequences.join(", ")}" if extra_sequences.any?
+
+        missing_sequences = from_sequences - to_sequences
+        notes << "Missing sequences: #{missing_sequences.join(", ")}" if missing_sequences.any?
+
+        from_types = from_columns.map { |c| [c[:name], c[:type]] }.to_h
+        to_types = to_columns.map { |c| [c[:name], c[:type]] }.to_h
+        different_types = []
+        shared_fields.each do |field|
+          if from_types[field] != to_types[field]
+            different_types << "#{field} (#{from_types[field]} -> #{to_types[field]})"
+          end
+        end
+        notes << "Different column types: #{different_types.join(", ")}" if different_types.any?
+      end
+      notes
+    end
+
+    def sync_data
+      raise Error, "This should never happen. Please file a bug." if shared_fields.empty?
+
+      sql_clause = String.new("")
+      sql_clause << " #{opts[:sql]}" if opts[:sql]
+
+      bad_fields = opts[:no_rules] ? [] : config["data_rules"]
+      primary_key = destination.primary_key(table)
+      copy_fields = shared_fields.map { |f| f2 = bad_fields.to_a.find { |bf, _| rule_match?(table, f, bf) }; f2 ? "#{apply_strategy(f2[1], table, f, primary_key)} AS #{quote_ident(f)}" : "#{quoted_table}.#{quote_ident(f)}" }.join(", ")
+      fields = shared_fields.map { |f| quote_ident(f) }.join(", ")
+
+      seq_values = {}
+      shared_sequences.each do |seq|
+        seq_values[seq] = source.last_value(seq)
+      end
+
+      copy_to_command = "COPY (SELECT #{copy_fields} FROM #{quoted_table}#{sql_clause}) TO STDOUT"
+      if opts[:in_batches]
+        raise Error, "No primary key" if primary_key.empty?
+        primary_key = primary_key.first
+
+        destination.truncate(table) if opts[:truncate]
+
+        from_max_id = source.max_id(table, primary_key)
+        to_max_id = destination.max_id(table, primary_key) + 1
+
+        if to_max_id == 1
+          from_min_id = source.min_id(table, primary_key)
+          to_max_id = from_min_id if from_min_id > 0
+        end
+
+        starting_id = to_max_id
+        batch_size = opts[:batch_size]
+
+        i = 1
+        batch_count = ((from_max_id - starting_id + 1) / batch_size.to_f).ceil
+
+        while starting_id <= from_max_id
+          where = "#{quote_ident(primary_key)} >= #{starting_id} AND #{quote_ident(primary_key)} < #{starting_id + batch_size}"
+          log "    #{i}/#{batch_count}: #{where}"
+
+          # TODO be smarter for advance sql clauses
+          batch_sql_clause = " #{sql_clause.length > 0 ? "#{sql_clause} AND" : "WHERE"} #{where}"
+
+          batch_copy_to_command = "COPY (SELECT #{copy_fields} FROM #{quoted_table}#{batch_sql_clause}) TO STDOUT"
+          copy(batch_copy_to_command, dest_table: table, dest_fields: fields)
+
+          starting_id += batch_size
+          i += 1
+
+          if opts[:sleep] && starting_id <= from_max_id
+            sleep(opts[:sleep])
+          end
+        end
+      elsif !opts[:truncate] && (opts[:overwrite] || opts[:preserve] || !sql_clause.empty?)
+        raise Error, "No primary key" if primary_key.empty?
+
+        # create a temp table
+        temp_table = "pgsync_#{rand(1_000_000_000)}"
+        destination.execute("CREATE TEMPORARY TABLE #{quote_ident_full(temp_table)} AS TABLE #{quoted_table} WITH NO DATA")
+
+        # load data
+        copy(copy_to_command, dest_table: temp_table, dest_fields: fields)
+
+        on_conflict = primary_key.map { |pk| quote_ident(pk) }.join(", ")
+        action =
+          if opts[:preserve]
+            "NOTHING"
+          else # overwrite or sql clause
+            setter = shared_fields.reject { |f| primary_key.include?(f) }.map { |f| "#{quote_ident(f)} = EXCLUDED.#{quote_ident(f)}" }
+            "UPDATE SET #{setter.join(", ")}"
+          end
+        destination.execute("INSERT INTO #{quoted_table} (SELECT * FROM #{quote_ident_full(temp_table)}) ON CONFLICT (#{on_conflict}) DO #{action}")
+      else
+        # use delete instead of truncate for foreign keys
+        if opts[:defer_constraints]
+          destination.execute("DELETE FROM #{quoted_table}")
+        else
+          destination.truncate(table)
+        end
+        copy(copy_to_command, dest_table: table, dest_fields: fields)
+      end
+      seq_values.each do |seq, value|
+        destination.execute("SELECT setval(#{escape(seq)}, #{escape(value)})")
+      end
+
+      {status: "success"}
+    end
+
+    private
+
+    def with_notices
+      notices = []
+      [source, destination].each do |data_source|
+        data_source.send(:conn).set_notice_processor do |message|
+          notices << message.strip
+        end
+      end
+      result = yield
+      result[:notices] = notices if result
+      result
+    ensure
+      # clear notice processor
+      [source, destination].each do |data_source|
+        data_source.send(:conn).set_notice_processor
+      end
+    end
+
+    # TODO add retries
+    def handle_errors
+      yield
+    rescue => e
+      raise e if opts[:debug]
+
+      message =
+        case e
+        when PG::ConnectionBad
+          # likely fine to show simplified message here
+          # the full message will be shown when first trying to connect
+          "Connection failed"
+        when PG::Error
+          e.message.sub("ERROR:  ", "")
+        when Error
+          e.message
+        else
+          "#{e.class.name}: #{e.message}"
+        end
+
+      {status: "error", message: message}
+    end
+
+    def copy(source_command, dest_table:, dest_fields:)
+      destination_command = "COPY #{quote_ident_full(dest_table)} (#{dest_fields}) FROM STDIN"
+      destination.conn.copy_data(destination_command) do
+        source.conn.copy_data(source_command) do
+          while (row = source.conn.get_copy_data)
+            destination.conn.put_copy_data(row)
+          end
+        end
+      end
+    end
+
+    # TODO better performance
+    def rule_match?(table, column, rule)
+      regex = Regexp.new('\A' + Regexp.escape(rule).gsub('\*','[^\.]*') + '\z')
+      regex.match(column) || regex.match("#{table.name}.#{column}") || regex.match("#{table.schema}.#{table.name}.#{column}")
+    end
+
+    # TODO wildcard rules
+    def apply_strategy(rule, table, column, primary_key)
+      if rule.is_a?(Hash)
+        if rule.key?("value")
+          escape(rule["value"])
+        elsif rule.key?("statement")
+          rule["statement"]
+        else
+          raise Error, "Unknown rule #{rule.inspect} for column #{column}"
+        end
+      else
+        case rule
+        when "untouched"
+          quote_ident(column)
+        when "unique_email"
+          "'email' || #{quoted_primary_key(table, primary_key, rule)}::text || '@example.org'"
+        when "unique_phone"
+          "(#{quoted_primary_key(table, primary_key, rule)}::bigint + 1000000000)::text"
+        when "unique_secret"
+          "'secret' || #{quoted_primary_key(table, primary_key, rule)}::text"
+        when "random_int", "random_number"
+          "(RANDOM() * 100)::int"
+        when "random_date"
+          "date '1970-01-01' + (RANDOM() * 10000)::int"
+        when "random_time"
+          "NOW() - (RANDOM() * 100000000)::int * INTERVAL '1 second'"
+        when "random_ip"
+          "(1 + RANDOM() * 254)::int::text || '.0.0.1'"
+        when "random_letter"
+          "chr(65 + (RANDOM() * 26)::int)"
+        when "random_string"
+          "RIGHT(MD5(RANDOM()::text), 10)"
+        when "null", nil
+          "NULL"
+        else
+          raise Error, "Unknown rule #{rule} for column #{column}"
+        end
+      end
+    end
+
+    def quoted_primary_key(table, primary_key, rule)
+      raise Error, "Single column primary key required for this data rule: #{rule}" unless primary_key.size == 1
+      "#{quoted_table}.#{quote_ident(primary_key.first)}"
+    end
+
+    def maybe_disable_triggers
+      if opts[:disable_integrity] || opts[:disable_integrity_v2] || opts[:disable_user_triggers]
+        destination.transaction do
+          triggers = destination.triggers(table)
+          triggers.select! { |t| t["enabled"] == "t" }
+          internal_triggers, user_triggers = triggers.partition { |t| t["internal"] == "t" }
+          integrity_triggers = internal_triggers.select { |t| t["integrity"] == "t" }
+          restore_triggers = []
+
+          # both --disable-integrity options require superuser privileges
+          # however, only v2 works on Amazon RDS, which added specific support for it
+          # https://aws.amazon.com/about-aws/whats-new/2014/11/10/amazon-rds-postgresql-read-replicas/
+          #
+          # session_replication_role disables more than foreign keys (like triggers and rules)
+          # this is probably fine, but keep the current default for now
+          if opts[:disable_integrity_v2] || (opts[:disable_integrity] && rds?)
+            # SET LOCAL lasts until the end of the transaction
+            # https://www.postgresql.org/docs/current/sql-set.html
+            destination.execute("SET LOCAL session_replication_role = replica")
+          elsif opts[:disable_integrity]
+            integrity_triggers.each do |trigger|
+              destination.execute("ALTER TABLE #{quoted_table} DISABLE TRIGGER #{quote_ident(trigger["name"])}")
+            end
+            restore_triggers.concat(integrity_triggers)
+          end
+
+          if opts[:disable_user_triggers]
+            # important!
+            # rely on Postgres to disable user triggers
+            # we don't want to accidentally disable non-user triggers if logic above is off
+            destination.execute("ALTER TABLE #{quoted_table} DISABLE TRIGGER USER")
+            restore_triggers.concat(user_triggers)
+          end
+
+          result = yield
+
+          # restore triggers that were previously enabled
+          restore_triggers.each do |trigger|
+            destination.execute("ALTER TABLE #{quoted_table} ENABLE TRIGGER #{quote_ident(trigger["name"])}")
+          end
+
+          result
+        end
+      else
+        yield
+      end
+    end
+
+    def rds?
+      destination.execute("SELECT name, setting FROM pg_settings WHERE name LIKE 'rds.%'").any?
+    end
+  end
+end