pgsync 0.5.2 → 0.6.1

data/lib/pgsync/data_source.rb

@@ -1,10 +1,11 @@
 module PgSync
   class DataSource
+    include Utils
+
     attr_reader :url
 
-    def initialize(source, timeout: 3)
-      @url = resolve_url(source)
-      @timeout = timeout
+    def initialize(url)
+      @url = url
     end
 
     def exists?
@@ -30,8 +31,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
 
@@ -39,25 +50,21 @@ 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
+      execute("SELECT #{columns.map { |f| "pg_get_serial_sequence(#{escape("#{quote_ident_full(table)}")}, #{escape(f)}) AS #{quote_ident(f)}" }.join(", ")}").first.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
 
+    # this value comes from pg_get_serial_sequence which is already quoted
     def last_value(seq)
-      execute("select last_value from #{seq}")[0]["last_value"]
+      execute("SELECT last_value FROM #{seq}").first["last_value"]
     end
 
     def truncate(table)
@@ -65,38 +72,57 @@ module PgSync
     end
 
     # https://stackoverflow.com/a/20537829
+    # TODO can simplify with array_position in Postgres 9.5+
     def primary_key(table)
-      query = <<-SQL
+      query = <<~SQL
         SELECT
           pg_attribute.attname,
-          format_type(pg_attribute.atttypid, pg_attribute.atttypmod)
+          format_type(pg_attribute.atttypid, pg_attribute.atttypmod),
+          pg_attribute.attnum,
+          pg_index.indkey
         FROM
           pg_index, pg_class, pg_attribute, pg_namespace
         WHERE
-          pg_class.oid = $2::regclass AND
-          indrelid = pg_class.oid AND
           nspname = $1 AND
+          relname = $2 AND
+          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
-      row = execute(query, [table.split(".", 2)[0], quote_ident_full(table)])[0]
-      row && row["attname"]
+      rows = execute(query, [table.schema, table.name])
+      rows.sort_by { |r| r["indkey"].split(" ").index(r["attnum"]) }.map { |r| r["attname"] }
+    end
+
+    def triggers(table)
+      query = <<~SQL
+        SELECT
+          tgname AS name,
+          tgisinternal AS internal,
+          tgenabled != 'D' AS enabled,
+          tgconstraint != 0 AS integrity
+        FROM
+          pg_trigger
+        WHERE
+          pg_trigger.tgrelid = $1::regclass
+      SQL
+      execute(query, [quote_ident_full(table)])
     end
 
     def conn
       @conn ||= begin
         begin
-          ENV["PGCONNECT_TIMEOUT"] ||= @timeout.to_s
+          ENV["PGCONNECT_TIMEOUT"] ||= "3"
           if @url =~ /\Apostgres(ql)?:\/\//
             config = @url
           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
@@ -108,84 +134,56 @@ module PgSync
       end
     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}"
-    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] }]
+    # reconnect for new thread or process
+    def reconnect_if_needed
+      reconnect if @concurrent_id != concurrent_id
     end
 
     def search_path
-      @search_path ||= execute("SELECT current_schemas(true)")[0]["current_schemas"][1..-2].split(",")
+      @search_path ||= execute("SELECT unnest(current_schemas(true)) AS schema").map { |r| r["schema"] }
     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 server_version_num
+      @server_version_num ||= execute("SHOW server_version_num").first["server_version_num"].to_i
     end
 
-    def table_set
-      @table_set ||= Set.new(tables)
-    end
-
-    def conninfo
-      @conninfo ||= conn.conninfo_hash
+    def execute(query, params = [])
+      conn.exec_params(query, params).to_a
     end
 
-    def quote_ident_full(ident)
-      ident.split(".", 2).map { |v| quote_ident(v) }.join(".")
+    def transaction
+      if conn.transaction_status == 0
+        # not currently in transaction
+        conn.transaction do
+          yield
+        end
+      else
+        yield
+      end
     end
 
-    def execute(query, params = [])
-      conn.exec_params(query, params).to_a
-    end
+    private
 
-    def quote_ident(value)
-      PG::Connection.quote_ident(value)
+    def concurrent_id
+      [Process.pid, Thread.current.object_id]
     end
 
-    def escape(value)
-      if value.is_a?(String)
-        "'#{quote_string(value)}'"
-      else
-        value
-      end
+    def reconnect
+      @conn.reset
+      @concurrent_id = concurrent_id
     end
 
-    # activerecord
-    def quote_string(s)
-      s.gsub(/\\/, '\&\&').gsub(/'/, "''")
+    def table_set
+      @table_set ||= Set.new(tables)
     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
+    def conninfo
+      @conninfo ||= begin
+        unless conn.respond_to?(:conninfo_hash)
+          raise Error, "libpq is too old. Upgrade it and run `gem install pg`"
         end
+        conn.conninfo_hash
       end
-      source
     end
   end
 end

data/lib/pgsync/init.rb

@@ -2,40 +2,78 @@ 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."
       else
         exclude =
-          if rails_app?
+          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
               # exclude:
-              #   - schema_migrations
-              #   - ar_internal_metadata
+              #   - table1
+              #   - table2
             EOS
           end
 
         # create file
         contents = File.read(__dir__ + "/../../config.yml")
+        contents.sub!("$(some_command)", "$(heroku config:get DATABASE_URL)") if heroku?
         File.write(file, contents % {exclude: exclude})
 
         log "#{file} created. Add your database credentials."
       end
     end
 
-    # TODO maybe check parent directories
-    def rails_app?
+    def django?
+      (File.read("manage.py") =~ /django/i) rescue false
+    end
+
+    def heroku?
+      `git remote -v 2>&1`.include?("git.heroku.com") rescue false
+    end
+
+    def laravel?
+      File.exist?("artisan")
+    end
+
+    def rails?
       File.exist?("bin/rails")
     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/sync.rb

@@ -2,147 +2,87 @@ 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
+
+      args = @arguments
+      opts = @options
 
-      # TODO remove deprecations
-      map_deprecations(args, opts)
+      # only resolve commands from config, not CLI arguments
+      [:to, :from].each do |opt|
+        opts[opt] ||= resolve_source(config[opt.to_s])
+      end
 
-      # start
-      start_time = Time.now
+      # 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]"
       end
 
-      source = DataSource.new(opts[:from])
       raise Error, "No source" unless source.exists?
-
-      destination = DataSource.new(opts[:to])
       raise Error, "No destination" unless destination.exists?
 
-      begin
-        # start connections
-        source.host
-        destination.host
-
-        unless opts[:to_safe] || destination.local?
-          raise Error, "Danger! Add `to_safe: true` to `.pgsync.yml` if the destination is not localhost or 127.0.0.1"
-        end
-
-        print_description("From", source)
-        print_description("To", destination)
-      ensure
-        source.close
-        destination.close
+      unless opts[:to_safe] || destination.local?
+        raise Error, "Danger! Add `to_safe: true` to `.pgsync.yml` if the destination is not localhost or 127.0.0.1"
       end
 
-      tables = nil
-      begin
-        tables = TableList.new(args, opts, source, config).tables
-      ensure
-        source.close
+      if (opts[:preserve] || opts[:overwrite]) && destination.server_version_num < 90500
+        raise Error, "Postgres 9.5+ is required for --preserve and --overwrite"
       end
 
-      confirm_tables_exist(source, tables, "source")
+      print_description("From", source)
+      print_description("To", destination)
 
-      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|
-            TableSync.new.sync(config, table, opts.merge(table_opts), source.url, destination.url)
-          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
-    ensure
-      data_source.close
-    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
-        if config_file
+        file = config_file
+        if file
           begin
-            YAML.load_file(config_file) || {}
+            YAML.load_file(file) || {}
           rescue Psych::SyntaxError => e
             raise Error, e.message
+          rescue Errno::ENOENT
+            raise Error, "Config file not found: #{file}"
           end
         else
           {}
@@ -150,91 +90,74 @@ 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)
-      item_spinners = {}
-
-      start = lambda do |item, i|
-        table, opts = item
-        message = String.new(":spinner ")
-        message << table.sub("#{first_schema}.", "")
-        # maybe output later
-        # message << " #{opts[:sql]}" if opts[:sql]
-        spinner = spinners.register(message)
-        spinner.auto_spin
-        item_spinners[item] = spinner
+    def config_file
+      if @options[:config]
+        @options[:config]
+      elsif @options[:db]
+        file = db_config_file(@options[:db])
+        search_tree(file) || file
+      else
+        search_tree(".pgsync.yml")
       end
+    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, display_message(result)].compact.join(" ")
-        end
-      end
+    def search_tree(file)
+      return file if File.exist?(file)
 
-      options = {start: start, finish: finish}
-      if @options[:debug] || @options[:in_batches]
-        options[:in_processes] = 0
-      else
-        options[:in_threads] = 4 if windows?
+      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
 
-      # 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, &block)
-
-      fail_sync(failed_tables) if failed_tables.any?
+    def print_description(prefix, source)
+      location = " on #{source.host}:#{source.port}" if source.host
+      log "#{prefix}: #{source.dbname}#{location}"
     end
 
-    def fail_sync(failed_tables)
-      raise Error, "Sync failed for #{failed_tables.size} table#{failed_tables.size == 1 ? nil : "s"}: #{failed_tables.join(", ")}"
+    def log_completed(started_at)
+      time = Time.now - started_at
+      message = "Completed in #{time.round(1)}s"
+      log colorize(message, :green)
     end
 
-    def display_message(result)
-      messages = []
-      messages << "- #{result[:time]}s" if result[:time]
-      messages << "(#{result[:message].gsub("\n", " ").strip})" if result[:message]
-      messages.join(" ")
+    def source
+      @source ||= data_source(@options[:from])
     end
 
-    def pretty_list(items)
-      items.each do |item|
-        log item
-      end
+    def destination
+      @destination ||= data_source(@options[:to])
     end
 
-    def deprecated(message)
-      log colorize("[DEPRECATED] #{message}", 33) # yellow
+    def data_source(url)
+      ds = DataSource.new(url)
+      ObjectSpace.define_finalizer(self, self.class.finalize(ds))
+      ds
     end
 
-    def log_completed(start_time)
-      time = Time.now - start_time
-      message = "Completed in #{time.round(1)}s"
-      log colorize(message, 32) # green
+    # 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 windows?
-      Gem.win_platform?
+    def self.finalize(ds)
+      # must use proc instead of stabby lambda
+      proc { ds.close }
     end
   end
 end