pgsync 0.5.5 → 0.6.4

data/lib/pgsync/data_source.rb

@@ -1,9 +1,13 @@
 module PgSync
   class DataSource
+    include Utils
+
     attr_reader :url
 
-    def initialize(source)
-      @url = resolve_url(source)
+    def initialize(url, name:, debug:)
+      @url = url
+      @name = name
+      @debug = debug
     end
 
     def exists?
@@ -29,8 +33,18 @@ module PgSync
     # gets visible tables
     def tables
       @tables ||= begin
-        query = "SELECT table_schema, table_name FROM information_schema.tables WHERE table_type = 'BASE TABLE' AND table_schema NOT IN ('information_schema', 'pg_catalog') ORDER BY 1, 2"
-        execute(query).map { |row| "#{row["table_schema"]}.#{row["table_name"]}" }
+        query = <<~SQL
+          SELECT
+            table_schema AS schema,
+            table_name AS table
+          FROM
+            information_schema.tables
+          WHERE
+            table_type = 'BASE TABLE' AND
+            table_schema NOT IN ('information_schema', 'pg_catalog')
+          ORDER BY 1, 2
+        SQL
+        execute(query).map { |row| Table.new(row["schema"], row["table"]) }
       end
     end
 
@@ -38,54 +52,24 @@ module PgSync
       table_set.include?(table)
     end
 
-    def columns(table)
-      query = "SELECT column_name FROM information_schema.columns WHERE table_schema = $1 AND table_name = $2"
-      execute(query, table.split(".", 2)).map { |row| row["column_name"] }
-    end
-
-    def sequences(table, columns)
-      execute("SELECT #{columns.map { |f| "pg_get_serial_sequence(#{escape("#{quote_ident_full(table)}")}, #{escape(f)}) AS #{quote_ident(f)}" }.join(", ")}")[0].values.compact
-    end
-
     def max_id(table, primary_key, sql_clause = nil)
-      execute("SELECT MAX(#{quote_ident(primary_key)}) FROM #{quote_ident_full(table)}#{sql_clause}")[0]["max"].to_i
+      execute("SELECT MAX(#{quote_ident(primary_key)}) FROM #{quote_ident_full(table)}#{sql_clause}").first["max"].to_i
     end
 
     def min_id(table, primary_key, sql_clause = nil)
-      execute("SELECT MIN(#{quote_ident(primary_key)}) FROM #{quote_ident_full(table)}#{sql_clause}")[0]["min"].to_i
+      execute("SELECT MIN(#{quote_ident(primary_key)}) FROM #{quote_ident_full(table)}#{sql_clause}").first["min"].to_i
     end
 
     def last_value(seq)
-      execute("select last_value from #{seq}")[0]["last_value"]
+      execute("SELECT last_value FROM #{quote_ident_full(seq)}").first["last_value"]
     end
 
     def truncate(table)
       execute("TRUNCATE #{quote_ident_full(table)} CASCADE")
     end
 
-    # https://stackoverflow.com/a/20537829
-    def primary_key(table)
-      query = <<-SQL
-        SELECT
-          pg_attribute.attname,
-          format_type(pg_attribute.atttypid, pg_attribute.atttypmod)
-        FROM
-          pg_index, pg_class, pg_attribute, pg_namespace
-        WHERE
-          pg_class.oid = $2::regclass AND
-          indrelid = pg_class.oid AND
-          nspname = $1 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
-      row = execute(query, [table.split(".", 2)[0], quote_ident_full(table)])[0]
-      row && row["attname"]
-    end
-
     def triggers(table)
-      query = <<-SQL
+      query = <<~SQL
         SELECT
           tgname AS name,
           tgisinternal AS internal,
@@ -108,9 +92,10 @@ module PgSync
           else
             config = {dbname: @url}
           end
+          @concurrent_id = concurrent_id
           PG::Connection.new(config)
         rescue URI::InvalidURIError
-          raise Error, "Invalid connection string"
+          raise Error, "Invalid connection string. Make sure it works with `psql`"
         end
       end
     end
@@ -122,55 +107,57 @@ module PgSync
       end
     end
 
-    def reconnect
-      @conn.reset
-    end
-
-    def dump_command(tables)
-      tables = tables ? tables.keys.map { |t| "-t #{Shellwords.escape(quote_ident_full(t))}" }.join(" ") : ""
-      "pg_dump -Fc --verbose --schema-only --no-owner --no-acl #{tables} -d #{@url}"
+    # reconnect for new thread or process
+    def reconnect_if_needed
+      reconnect if @concurrent_id != concurrent_id
     end
 
-    def restore_command
-      if_exists = Gem::Version.new(pg_restore_version) >= Gem::Version.new("9.4.0")
-      "pg_restore --verbose --no-owner --no-acl --clean #{if_exists ? "--if-exists" : nil} -d #{@url}"
-    end
-
-    def fully_resolve_tables(tables)
-      no_schema_tables = {}
-      search_path_index = Hash[search_path.map.with_index.to_a]
-      self.tables.group_by { |t| t.split(".", 2)[-1] }.each do |group, t2|
-        no_schema_tables[group] = t2.sort_by { |t| [search_path_index[t.split(".", 2)[0]] || 1000000, t] }[0]
-      end
-
-      Hash[tables.map { |k, v| [no_schema_tables[k] || k, v] }]
+    def search_path
+      @search_path ||= execute("SELECT unnest(current_schemas(true)) AS schema").map { |r| r["schema"] }
     end
 
-    def search_path
-      @search_path ||= execute("SELECT current_schemas(true)")[0]["current_schemas"][1..-2].split(",")
+    def server_version_num
+      @server_version_num ||= execute("SHOW server_version_num").first["server_version_num"].to_i
     end
 
     def execute(query, params = [])
+      log_sql query, params
       conn.exec_params(query, params).to_a
     end
 
     def transaction
       if conn.transaction_status == 0
         # not currently in transaction
-        conn.transaction do
-          yield
-        end
+        log_sql "BEGIN"
+        result =
+          conn.transaction do
+            yield
+          end
+        log_sql "COMMIT"
+        result
       else
         yield
       end
     end
 
+    # TODO log time for each statement
+    def log_sql(query, params = {})
+      if @debug
+        message = "#{colorize("[#{@name}]", :cyan)} #{query.gsub(/\s+/, " ").strip}"
+        message = "#{message} #{params.inspect}" if params.any?
+        log message
+      end
+    end
+
     private
 
-    def pg_restore_version
-      `pg_restore --version`.lines[0].chomp.split(" ")[-1].split(/[^\d.]/)[0]
-    rescue Errno::ENOENT
-      raise Error, "pg_restore not found"
+    def concurrent_id
+      [Process.pid, Thread.current.object_id]
+    end
+
+    def reconnect
+      @conn.reset
+      @concurrent_id = concurrent_id
     end
 
     def table_set
@@ -185,41 +172,5 @@ module PgSync
         conn.conninfo_hash
       end
     end
-
-    def quote_ident_full(ident)
-      ident.split(".", 2).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
-    end
-
-    # activerecord
-    def quote_string(s)
-      s.gsub(/\\/, '\&\&').gsub(/'/, "''")
-    end
-
-    def resolve_url(source)
-      if source
-        source = source.dup
-        source.gsub!(/\$\([^)]+\)/) do |m|
-          command = m[2..-2]
-          result = `#{command}`.chomp
-          unless $?.success?
-            raise Error, "Command exited with non-zero status:\n#{command}"
-          end
-          result
-        end
-      end
-      source
-    end
   end
 end

data/lib/pgsync/init.rb

@@ -2,11 +2,26 @@ module PgSync
   class Init
     include Utils
 
-    def perform(opts)
-      # needed for config_file method
-      @options = opts.to_hash
+    def initialize(arguments, options)
+      @arguments = arguments
+      @options = options
+    end
+
+    def perform
+      if @arguments.size > 1
+        raise Error, "Usage:\n    pgsync --init [db]"
+      end
 
-      file = db_config_file(opts.arguments[0]) || config_file || ".pgsync.yml"
+      file =
+        if @options[:config]
+          @options[:config]
+        elsif @arguments.any?
+          db_config_file(@arguments.first)
+        elsif @options[:db]
+          db_config_file(@options[:db])
+        else
+          ".pgsync.yml"
+        end
 
       if File.exist?(file)
         raise Error, "#{file} exists."
@@ -15,8 +30,19 @@ module PgSync
           if rails?
             <<~EOS
               exclude:
-                - schema_migrations
                 - ar_internal_metadata
+                - schema_migrations
+            EOS
+          elsif django?
+            # TODO exclude other tables?
+            <<~EOS
+              exclude:
+                - django_migrations
+            EOS
+          elsif laravel?
+            <<~EOS
+              exclude:
+                - migrations
             EOS
           else
             <<~EOS
@@ -35,12 +61,30 @@ module PgSync
       end
     end
 
+    def django?
+      file_exists?("manage.py", /django/i)
+    end
+
     def heroku?
       `git remote -v 2>&1`.include?("git.heroku.com") rescue false
     end
 
+    def laravel?
+      file_exists?("artisan")
+    end
+
     def rails?
-      File.exist?("bin/rails")
+      file_exists?("bin/rails")
+    end
+
+    def file_exists?(path, contents = nil)
+      if contents
+        File.read(path).match(contents)
+      else
+        File.exist?(path)
+      end
+    rescue
+      false
     end
   end
 end

data/lib/pgsync/schema_sync.rb

@@ -0,0 +1,83 @@
+module PgSync
+  class SchemaSync
+    include Utils
+
+    attr_reader :args, :opts
+
+    def initialize(source:, destination:, tasks:, args:, opts:)
+      @source = source
+      @destination = destination
+      @tasks = tasks
+      @args = args
+      @opts = opts
+    end
+
+    def perform
+      if opts[:preserve]
+        raise Error, "Cannot use --preserve with --schema-first or --schema-only"
+      end
+
+      show_spinner = output.tty? && !opts[:debug]
+
+      if show_spinner
+        spinner = TTY::Spinner.new(":spinner Syncing schema", format: :dots)
+        spinner.auto_spin
+      end
+
+      # if spinner, capture lines to show on error
+      lines = []
+      success =
+        run_command("#{dump_command} | #{restore_command}") do |line|
+          if show_spinner
+            lines << line
+          else
+            log line
+          end
+        end
+
+      if show_spinner
+        if success
+          spinner.success
+        else
+          spinner.error
+          log lines.join
+        end
+      end
+
+      raise Error, "Schema sync returned non-zero exit code" unless success
+    end
+
+    private
+
+    def run_command(command)
+      Open3.popen2e(command) do |stdin, stdout, wait_thr|
+        stdout.each do |line|
+          yield line
+        end
+        wait_thr.value.success?
+      end
+    end
+
+    def pg_restore_version
+      `pg_restore --version`.lines[0].chomp.split(" ")[-1].split(/[^\d.]/)[0]
+    rescue Errno::ENOENT
+      raise Error, "pg_restore not found"
+    end
+
+    def dump_command
+      tables =
+        if !opts[:all_schemas] || opts[:tables] || opts[:groups] || args[0] || opts[:exclude] || opts[:schemas]
+          @tasks.map { |task| "-t #{Shellwords.escape(task.quoted_table)}" }
+        else
+          []
+        end
+
+      "pg_dump -Fc --verbose --schema-only --no-owner --no-acl #{tables.join(" ")} -d #{@source.url}"
+    end
+
+    def restore_command
+      if_exists = Gem::Version.new(pg_restore_version) >= Gem::Version.new("9.4.0")
+      "pg_restore --verbose --no-owner --no-acl --clean #{if_exists ? "--if-exists" : nil} -d #{@destination.url}"
+    end
+  end
+end

data/lib/pgsync/sequence.rb

@@ -0,0 +1,29 @@
+# minimal class to keep schema and sequence name separate
+module PgSync
+  class Sequence
+    attr_reader :schema, :name, :column
+
+    def initialize(schema, name, column:)
+      @schema = schema
+      @name = name
+      @column = column
+    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/sync.rb

@@ -2,21 +2,26 @@ module PgSync
   class Sync
     include Utils
 
-    def perform(options)
-      args = options.arguments
-      opts = options.to_hash
-      @options = opts
+    def initialize(arguments, options)
+      @arguments = arguments
+      @options = options
+    end
 
-      # merge config
-      [:to, :from, :to_safe, :exclude, :schemas].each do |opt|
-        opts[opt] ||= config[opt.to_s]
-      end
+    def perform
+      started_at = Time.now
 
-      # TODO remove deprecations in 0.6.0
-      map_deprecations(args, opts)
+      args = @arguments
+      opts = @options
 
-      # start
-      start_time = Time.now
+      # only resolve commands from config, not CLI arguments
+      [:to, :from].each do |opt|
+        opts[opt] ||= resolve_source(config[opt.to_s])
+      end
+
+      # merge other config
+      [:to_safe, :exclude, :schemas].each do |opt|
+        opts[opt] ||= config[opt.to_s]
+      end
 
       if args.size > 2
         raise Error, "Usage:\n    pgsync [options]"
@@ -32,89 +37,41 @@ module PgSync
       print_description("From", source)
       print_description("To", destination)
 
-      tables = TableList.new(args, opts, source, config).tables
-      confirm_tables_exist(source, tables, "source")
+      if (opts[:preserve] || opts[:overwrite]) && destination.server_version_num < 90500
+        raise Error, "Postgres 9.5+ is required for --preserve and --overwrite"
+      end
 
-      if opts[:list]
-        confirm_tables_exist(destination, tables, "destination")
+      resolver = TaskResolver.new(args: args, opts: opts, source: source, destination: destination, config: config, first_schema: first_schema)
+      tasks =
+        resolver.tasks.map do |task|
+          Task.new(source: source, destination: destination, config: config, table: task[:table], opts: opts.merge(sql: task[:sql]))
+        end
 
-        list_items =
-          if args[0] == "groups"
-            (config["groups"] || {}).keys
-          else
-            tables.keys
-          end
+      if opts[:in_batches] && tasks.size > 1
+        raise Error, "Cannot use --in-batches with multiple tables"
+      end
+
+      confirm_tables_exist(source, tasks, "source")
 
-        pretty_list list_items
+      if opts[:list]
+        confirm_tables_exist(destination, tasks, "destination")
+        tasks.each do |task|
+          log task_name(task)
+        end
       else
         if opts[:schema_first] || opts[:schema_only]
-          if opts[:preserve]
-            raise Error, "Cannot use --preserve with --schema-first or --schema-only"
-          end
-
-          log "* Dumping schema"
-          schema_tables = tables if !opts[:all_schemas] || opts[:tables] || opts[:groups] || args[0] || opts[:exclude]
-          sync_schema(source, destination, schema_tables)
+          SchemaSync.new(source: source, destination: destination, tasks: tasks, args: args, opts: opts).perform
         end
 
         unless opts[:schema_only]
-          confirm_tables_exist(destination, tables, "destination")
-
-          in_parallel(tables, first_schema: source.search_path.find { |sp| sp != "pg_catalog" }) do |table, table_opts, source, destination|
-            TableSync.new(source: source, destination: destination).sync(config, table, opts.merge(table_opts))
-          end
+          TableSync.new(source: source, destination: destination, tasks: tasks, opts: opts, resolver: resolver).perform
         end
 
-        log_completed(start_time)
+        log_completed(started_at)
       end
     end
 
-    def confirm_tables_exist(data_source, tables, description)
-      tables.keys.each do |table|
-        unless data_source.table_exists?(table)
-          raise Error, "Table does not exist in #{description}: #{table}"
-        end
-      end
-    end
-
-    def map_deprecations(args, opts)
-      command = args[0]
-
-      case command
-      when "schema"
-        args.shift
-        opts[:schema_only] = true
-        deprecated "Use `psync --schema-only` instead"
-      when "tables"
-        args.shift
-        opts[:tables] = args.shift
-        deprecated "Use `pgsync #{opts[:tables]}` instead"
-      when "groups"
-        args.shift
-        opts[:groups] = args.shift
-        deprecated "Use `pgsync #{opts[:groups]}` instead"
-      end
-
-      if opts[:where]
-        opts[:sql] ||= String.new
-        opts[:sql] << " WHERE #{opts[:where]}"
-        deprecated "Use `\"WHERE #{opts[:where]}\"` instead"
-      end
-
-      if opts[:limit]
-        opts[:sql] ||= String.new
-        opts[:sql] << " LIMIT #{opts[:limit]}"
-        deprecated "Use `\"LIMIT #{opts[:limit]}\"` instead"
-      end
-    end
-
-    def sync_schema(source, destination, tables = nil)
-      dump_command = source.dump_command(tables)
-      restore_command = destination.restore_command
-      unless system("#{dump_command} | #{restore_command}")
-        raise Error, "Schema sync returned non-zero exit code"
-      end
-    end
+    private
 
     def config
       @config ||= begin
@@ -124,6 +81,8 @@ module PgSync
             YAML.load_file(file) || {}
           rescue Psych::SyntaxError => e
             raise Error, e.message
+          rescue Errno::ENOENT
+            raise Error, "Config file not found: #{file}"
           end
         else
           {}
@@ -131,147 +90,71 @@ module PgSync
       end
     end
 
-    def print_description(prefix, source)
-      location = " on #{source.host}:#{source.port}" if source.host
-      log "#{prefix}: #{source.dbname}#{location}"
-    end
-
-    def in_parallel(tables, first_schema:, &block)
-      spinners = TTY::Spinner::Multi.new(format: :dots, output: output)
-      item_spinners = {}
-
-      start = lambda do |item, i|
-        table, opts = item
-        message = String.new(":spinner ")
-        message << table.sub("#{first_schema}.", "")
-        message << " #{opts[:sql]}" if opts[:sql]
-        spinner = spinners.register(message)
-        if @options[:in_batches]
-          # log instead of spin for non-tty
-          log message.sub(":spinner", "⠋")
-        else
-          spinner.auto_spin
-        end
-        item_spinners[item] = spinner
-      end
-
-      failed_tables = []
-
-      finish = lambda do |item, i, result|
-        spinner = item_spinners[item]
-        table_name = item.first.sub("#{first_schema}.", "")
-
-        if result[:status] == "success"
-          spinner.success(display_message(result))
-        else
-          # TODO add option to fail fast
-          spinner.error(display_message(result))
-          failed_tables << table_name
-          fail_sync(failed_tables) if @options[:fail_fast]
-        end
-
-        unless spinner.send(:tty?)
-          status = result[:status] == "success" ? "✔" : "✖"
-          log [status, table_name, item.last[:sql], display_message(result)].compact.join(" ")
-        end
-      end
-
-      options = {start: start, finish: finish}
-
-      jobs = @options[:jobs]
-      if @options[:debug] || @options[:in_batches] || @options[:defer_constraints]
-        warn "--jobs ignored" if jobs
-        jobs = 0
-      end
-
-      if windows?
-        options[:in_threads] = jobs || 4
-      else
-        options[:in_processes] = jobs if jobs
-      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(tables, **options) do |table, table_opts|
-          # must reconnect for new thread or process
-          # TODO only reconnect first time
-          unless options[:in_processes] == 0
-            source.reconnect
-            destination.reconnect
-          end
-
-          # TODO warn if there are non-deferrable constraints on the table
-
-          yield table, table_opts, source, destination
-        end
-      end
-
-      fail_sync(failed_tables) if failed_tables.any?
-    end
-
-    def maybe_defer_constraints
-      if @options[: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
+    def config_file
+      if @options[:config]
+        @options[:config]
+      elsif @options[:db]
+        file = db_config_file(@options[:db])
+        search_tree(file) || file
       else
-        yield
+        search_tree(".pgsync.yml")
       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_message(result)
-      messages = []
-      messages << "- #{result[:time]}s" if result[:time]
-      messages << "(#{result[:message].lines.first.to_s.strip})" if result[:message]
-      messages.join(" ")
-    end
+    def search_tree(file)
+      return file if File.exist?(file)
 
-    def pretty_list(items)
-      items.each do |item|
-        log item
+      path = Dir.pwd
+      # prevent infinite loop
+      20.times do
+        absolute_file = File.join(path, file)
+        break absolute_file if File.exist?(absolute_file)
+        path = File.dirname(path)
+        break if path == "/"
       end
     end
 
-    def deprecated(message)
-      log colorize("[DEPRECATED] #{message}", :yellow)
+    def print_description(prefix, source)
+      location = " on #{source.host}:#{source.port}" if source.host
+      log "#{prefix}: #{source.dbname}#{location}"
     end
 
-    def log_completed(start_time)
-      time = Time.now - start_time
+    def log_completed(started_at)
+      time = Time.now - started_at
       message = "Completed in #{time.round(1)}s"
       log colorize(message, :green)
     end
 
-    def windows?
-      Gem.win_platform?
-    end
-
     def source
-      @source ||= data_source(@options[:from])
+      @source ||= data_source(@options[:from], "from")
     end
 
     def destination
-      @destination ||= data_source(@options[:to])
+      @destination ||= data_source(@options[:to], "to")
     end
 
-    def data_source(url)
-      ds = DataSource.new(url)
+    def data_source(url, name)
+      ds = DataSource.new(url, name: name, debug: @options[:debug])
       ObjectSpace.define_finalizer(self, self.class.finalize(ds))
       ds
     end
 
+    # ideally aliases would work, but haven't found a nice way to do this
+    def resolve_source(source)
+      if source
+        source = source.dup
+        source.gsub!(/\$\([^)]+\)/) do |m|
+          command = m[2..-2]
+          result = `#{command}`.chomp
+          unless $?.success?
+            raise Error, "Command exited with non-zero status:\n#{command}"
+          end
+          result
+        end
+      end
+      source
+    end
+
     def self.finalize(ds)
       # must use proc instead of stabby lambda
       proc { ds.close }