pgsync 0.5.5 → 0.6.4

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,274 +2,321 @@ module PgSync
   class TableSync
     include Utils
 
-    attr_reader :source, :destination
+    attr_reader :source, :destination, :tasks, :opts, :resolver
 
-    def initialize(source:, destination:)
+    def initialize(source:, destination:, tasks:, opts:, resolver:)
       @source = source
       @destination = destination
+      @tasks = tasks
+      @opts = opts
+      @resolver = resolver
     end
 
-    def sync(config, table, opts)
-      maybe_disable_triggers(table, opts) do
-        sync_data(config, table, opts)
-      end
-    end
+    def perform
+      confirm_tables_exist(destination, tasks, "destination")
 
-    def sync_data(config, table, opts)
-      start_time = Time.now
+      add_columns
 
-      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_primary_keys
 
-      if opts[:no_sequences]
-        from_sequences = []
-        to_sequences = []
-      else
-        from_sequences = source.sequences(table, shared_fields)
-        to_sequences = destination.sequences(table, shared_fields)
-      end
+      add_sequences unless opts[:no_sequences]
 
-      shared_sequences = to_sequences & from_sequences
-      extra_sequences = to_sequences - from_sequences
-      missing_sequences = from_sequences - to_sequences
+      show_notes
 
-      sql_clause = String.new("")
-      sql_clause << " #{opts[:sql]}" if opts[:sql]
+      # don't sync tables with no shared fields
+      # we show a warning message above
+      run_tasks(tasks.reject { |task| task.shared_fields.empty? })
+    end
 
-      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?
+    def add_columns
+      source_columns = columns(source)
+      destination_columns = columns(destination)
 
-      return {status: "success", message: "No fields to copy"} if shared_fields.empty?
+      tasks.each do |task|
+        task.from_columns = source_columns[task.table] || []
+        task.to_columns = destination_columns[task.table] || []
+      end
+    end
 
-      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(", ")
+    def add_primary_keys
+      destination_primary_keys = primary_keys(destination)
 
-      seq_values = {}
-      shared_sequences.each do |seq|
-        seq_values[seq] = source.last_value(seq)
+      tasks.each do |task|
+        task.to_primary_key = destination_primary_keys[task.table] || []
       end
+    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
+    def add_sequences
+      source_sequences = sequences(source)
+      destination_sequences = sequences(destination)
 
-        destination.truncate(table) if opts[:truncate]
+      tasks.each do |task|
+        shared_columns = Set.new(task.shared_fields)
 
-        from_max_id = source.max_id(table, primary_key)
-        to_max_id = destination.max_id(table, primary_key) + 1
+        task.from_sequences = (source_sequences[task.table] || []).select { |s| shared_columns.include?(s.column) }
+        task.to_sequences = (destination_sequences[task.table] || []).select { |s| shared_columns.include?(s.column) }
+      end
+    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 sequences(data_source)
+      query = <<~SQL
+        SELECT
+          nt.nspname as schema,
+          t.relname as table,
+          a.attname as column,
+          n.nspname as sequence_schema,
+          s.relname as sequence
+        FROM
+          pg_class s
+        INNER JOIN
+          pg_depend d ON d.objid = s.oid
+        INNER JOIN
+          pg_class t ON d.objid = s.oid AND d.refobjid = t.oid
+        INNER JOIN
+          pg_attribute a ON (d.refobjid, d.refobjsubid) = (a.attrelid, a.attnum)
+        INNER JOIN
+          pg_namespace n ON n.oid = s.relnamespace
+        INNER JOIN
+          pg_namespace nt ON nt.oid = t.relnamespace
+        WHERE
+          s.relkind = 'S'
+      SQL
+      data_source.execute(query).group_by { |r| Table.new(r["schema"], r["table"]) }.map do |k, v|
+        [k, v.map { |r| Sequence.new(r["sequence_schema"], r["sequence"], column: r["column"]) }]
+      end.to_h
+    end
 
-        starting_id = to_max_id
-        batch_size = opts[:batch_size]
+    def primary_keys(data_source)
+      # https://stackoverflow.com/a/20537829
+      # TODO can simplify with array_position in Postgres 9.5+
+      query = <<~SQL
+        SELECT
+          nspname AS schema,
+          relname AS table,
+          pg_attribute.attname AS column,
+          format_type(pg_attribute.atttypid, pg_attribute.atttypmod),
+          pg_attribute.attnum,
+          pg_index.indkey
+        FROM
+          pg_index, pg_class, pg_attribute, pg_namespace
+        WHERE
+          indrelid = pg_class.oid AND
+          pg_class.relnamespace = pg_namespace.oid AND
+          pg_attribute.attrelid = pg_class.oid AND
+          pg_attribute.attnum = any(pg_index.indkey) AND
+          indisprimary
+      SQL
+      data_source.execute(query).group_by { |r| Table.new(r["schema"], r["table"]) }.map do |k, v|
+        [k, v.sort_by { |r| r["indkey"].split(" ").index(r["attnum"]) }.map { |r| r["column"] }]
+      end.to_h
+    end
 
-        i = 1
-        batch_count = ((from_max_id - starting_id + 1) / batch_size.to_f).ceil
+    def show_notes
+      # for tables
+      resolver.notes.each do |note|
+        warning note
+      end
 
-        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}"
+      # for columns and sequences
+      tasks.each do |task|
+        task.notes.each do |note|
+          warning "#{task_name(task)}: #{note}"
+        end
+      end
 
-          # TODO be smarter for advance sql clauses
-          batch_sql_clause = " #{sql_clause.length > 0 ? "#{sql_clause} AND" : "WHERE"} #{where}"
+      # 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
 
-          batch_copy_to_command = "COPY (SELECT #{copy_fields} FROM #{quote_ident_full(table)}#{batch_sql_clause}) TO STDOUT"
-          copy(batch_copy_to_command, dest_table: table, dest_fields: fields)
+    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
 
-          starting_id += batch_size
-          i += 1
+    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
 
-          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" unless primary_key
+    def run_tasks(tasks, &block)
+      notices = []
+      failed_tables = []
+      started_at = {}
 
-        # create a temp table
-        temp_table = "pgsync_#{rand(1_000_000_000)}"
-        destination.execute("CREATE TEMPORARY TABLE #{quote_ident_full(temp_table)} AS TABLE #{quote_ident_full(table)} WITH NO DATA")
+      show_spinners = output.tty? && !opts[:in_batches] && !opts[:debug]
+      if show_spinners
+        spinners = TTY::Spinner::Multi.new(format: :dots, output: output)
+        task_spinners = {}
+      end
 
-        # load data
-        copy(copy_to_command, dest_table: temp_table, dest_fields: fields)
+      start = lambda do |task, i|
+        message = ":spinner #{display_item(task)}"
 
-        if opts[:preserve]
-          # insert into
-          destination.execute("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
-          destination.transaction do
-            destination.execute("DELETE FROM #{quote_ident_full(table)} WHERE #{quote_ident(primary_key)} IN (SELECT #{quote_ident(primary_key)} FROM #{quote_ident_full(temp_table)})")
-            destination.execute("INSERT INTO #{quote_ident_full(table)} (SELECT * FROM #{quote_ident(temp_table)})")
-          end
+        if show_spinners
+          spinner = spinners.register(message)
+          spinner.auto_spin
+          task_spinners[task] = spinner
+        elsif opts[:in_batches]
+          log message.sub(":spinner", "⠋")
         end
-      else
-        # use delete instead of truncate for foreign keys
-        if opts[:defer_constraints]
-          destination.execute("DELETE FROM #{quote_ident_full(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)})")
+
+        started_at[task] = Time.now
       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
+      finish = lambda do |task, i, result|
+        time = (Time.now - started_at[task]).round(1)
 
-      {status: "error", message: message}
-    end
+        success = result[:status] == "success"
 
-    private
-
-    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)
+        message =
+          if result[:message]
+            "(#{result[:message].lines.first.to_s.strip})"
+          else
+            "- #{time}s"
           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.split(".", 2)[-1]}.#{column}") || regex.match("#{table}.#{column}")
-    end
+        notices.concat(result[:notices])
 
-    # 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"]
+        if show_spinners
+          spinner = task_spinners[task]
+          if success
+            spinner.success(message)
+          else
+            spinner.error(message)
+          end
         else
-          raise Error, "Unknown rule #{rule.inspect} for column #{column}"
+          status = success ? "✔" : "✖"
+          log [status, display_item(task), message].join(" ")
         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}"
+
+        unless success
+          failed_tables << task_name(task)
+          fail_sync(failed_tables) if opts[:fail_fast]
         end
       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)}"
-    end
+      options = {start: start, finish: finish}
 
-    def quote_ident_full(ident)
-      ident.split(".").map { |v| quote_ident(v) }.join(".")
-    end
+      jobs = opts[:jobs]
 
-    def quote_ident(value)
-      PG::Connection.quote_ident(value)
-    end
+      # disable multiple jobs for defer constraints and disable integrity
+      # so we can use a transaction to ensure a consistent snapshot
+      if opts[:debug] || opts[:in_batches] || opts[:defer_constraints] || opts[:defer_constraints_v2] || opts[:disable_integrity] || opts[:disable_integrity_v2]
+        warning "--jobs ignored" if jobs
+        jobs = 0
+      end
 
-    def escape(value)
-      if value.is_a?(String)
-        "'#{quote_string(value)}'"
+      if windows?
+        options[:in_threads] = jobs || 4
       else
-        value
+        options[:in_processes] = jobs if jobs
       end
-    end
 
-    # activerecord
-    def quote_string(s)
-      s.gsub(/\\/, '\&\&').gsub(/'/, "''")
+      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
+
+          task.perform
+        end
+      end
+
+      notices.each do |notice|
+        warning notice
+      end
+
+      fail_sync(failed_tables) if failed_tables.any?
     end
 
-    def maybe_disable_triggers(table, opts)
-      if opts[:disable_integrity] || opts[:disable_user_triggers]
+    # TODO add option to open transaction on source when manually specifying order of tables
+    def maybe_defer_constraints
+      if opts[:disable_integrity] || opts[:disable_integrity_v2]
+        # create a transaction on the source
+        # to ensure we get a consistent snapshot
+        source.transaction do
+          yield
+        end
+      elsif opts[:defer_constraints] || opts[:defer_constraints_v2]
         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 = []
-
-          if opts[:disable_integrity]
-            integrity_triggers.each do |trigger|
-              destination.execute("ALTER TABLE #{quote_ident_full(table)} DISABLE TRIGGER #{quote_ident(trigger["name"])}")
+          if opts[:defer_constraints_v2]
+            table_constraints = non_deferrable_constraints(destination)
+            table_constraints.each do |table, constraints|
+              constraints.each do |constraint|
+                destination.execute("ALTER TABLE #{quote_ident_full(table)} ALTER CONSTRAINT #{quote_ident(constraint)} DEFERRABLE")
+              end
             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 #{quote_ident_full(table)} DISABLE TRIGGER USER")
-            restore_triggers.concat(user_triggers)
           end
 
-          result = yield
+          destination.execute("SET CONSTRAINTS ALL DEFERRED")
 
-          # restore triggers that were previously enabled
-          restore_triggers.each do |trigger|
-            destination.execute("ALTER TABLE #{quote_ident_full(table)} ENABLE TRIGGER #{quote_ident(trigger["name"])}")
+          # create a transaction on the source
+          # to ensure we get a consistent snapshot
+          source.transaction do
+            yield
           end
 
-          result
+          # set them back
+          # there are 3 modes: DEFERRABLE INITIALLY DEFERRED, DEFERRABLE INITIALLY IMMEDIATE, and NOT DEFERRABLE
+          # we only update NOT DEFERRABLE
+          # https://www.postgresql.org/docs/current/sql-set-constraints.html
+          if opts[:defer_constraints_v2]
+            destination.execute("SET CONSTRAINTS ALL IMMEDIATE")
+
+            table_constraints.each do |table, constraints|
+              constraints.each do |constraint|
+                destination.execute("ALTER TABLE #{quote_ident_full(table)} ALTER CONSTRAINT #{quote_ident(constraint)} NOT DEFERRABLE")
+              end
+            end
+          end
         end
       else
         yield
       end
     end
+
+    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 display_item(item)
+      messages = []
+      messages << task_name(item)
+      messages << item.opts[:sql] if item.opts[:sql]
+      messages.join(" ")
+    end
+
+    def windows?
+      Gem.win_platform?
+    end
   end
 end