pgsync 0.5.1 → 0.6.0

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,40 +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 PG::ConnectionBad => e
-          raise PgSync::Error, e.message
         rescue URI::InvalidURIError
-          raise PgSync::Error, "Invalid connection string"
+          raise Error, "Invalid connection string. Make sure it works with `psql`"
         end
       end
     end
@@ -110,79 +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
-      psql_version = `psql --version`.lines[0].chomp.split(" ")[-1].split(/[^\d.]/)[0]
-      if_exists = Gem::Version.new(psql_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 table_set
-      @table_set ||= Set.new(tables)
+    def server_version_num
+      @server_version_num ||= execute("SHOW server_version_num").first["server_version_num"].to_i
     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 PgSync::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

@@ -0,0 +1,61 @@
+module PgSync
+  class Init
+    include Utils
+
+    def initialize(arguments, options)
+      @arguments = arguments
+      @options = options
+    end
+
+    def perform
+      if @arguments.size > 1
+        raise Error, "Usage:\n    pgsync --init [db]"
+      end
+
+      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?
+            <<~EOS
+              exclude:
+                - schema_migrations
+                - ar_internal_metadata
+            EOS
+          else
+            <<~EOS
+              # exclude:
+              #   - 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
+
+    def heroku?
+      `git remote -v 2>&1`.include?("git.heroku.com") rescue false
+    end
+
+    def rails?
+      File.exist?("bin/rails")
+    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/sync.rb

@@ -0,0 +1,162 @@
+module PgSync
+  class Sync
+    include Utils
+
+    def initialize(arguments, options)
+      @arguments = arguments
+      @options = options
+    end
+
+    def perform
+      started_at = Time.now
+
+      args = @arguments
+      opts = @options
+
+      # 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]"
+      end
+
+      raise Error, "No source" unless source.exists?
+      raise Error, "No destination" unless destination.exists?
+
+      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
+
+      if (opts[:preserve] || opts[:overwrite]) && destination.server_version_num < 90500
+        raise Error, "Postgres 9.5+ is required for --preserve and --overwrite"
+      end
+
+      print_description("From", source)
+      print_description("To", 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
+
+      if opts[:in_batches] && tasks.size > 1
+        raise Error, "Cannot use --in-batches with multiple tables"
+      end
+
+      confirm_tables_exist(source, tasks, "source")
+
+      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]
+          SchemaSync.new(source: source, destination: destination, tasks: tasks, args: args, opts: opts).perform
+        end
+
+        unless opts[:schema_only]
+          TableSync.new(source: source, destination: destination, tasks: tasks, opts: opts, resolver: resolver).perform
+        end
+
+        log_completed(started_at)
+      end
+    end
+
+    private
+
+    def config
+      @config ||= begin
+        file = config_file
+        if file
+          begin
+            YAML.load_file(file) || {}
+          rescue Psych::SyntaxError => e
+            raise Error, e.message
+          rescue Errno::ENOENT
+            raise Error, "Config file not found: #{file}"
+          end
+        else
+          {}
+        end
+      end
+    end
+
+    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
+
+    def search_tree(file)
+      return file if File.exist?(file)
+
+      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 print_description(prefix, source)
+      location = " on #{source.host}:#{source.port}" if source.host
+      log "#{prefix}: #{source.dbname}#{location}"
+    end
+
+    def log_completed(started_at)
+      time = Time.now - started_at
+      message = "Completed in #{time.round(1)}s"
+      log colorize(message, :green)
+    end
+
+    def source
+      @source ||= data_source(@options[:from])
+    end
+
+    def destination
+      @destination ||= data_source(@options[:to])
+    end
+
+    def data_source(url)
+      ds = DataSource.new(url)
+      ObjectSpace.define_finalizer(self, self.class.finalize(ds))
+      ds
+    end
+
+    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 }
+    end
+  end
+end