pgcp 0.0.0

checksums.yaml

@@ -0,0 +1,7 @@
+---
+SHA1:
+  metadata.gz: e0bf53cb613c6859915d664646ef18c6ef216767
+  data.tar.gz: b3e0ace6249b20249efba309f40e7385d1f4cd24
+SHA512:
+  metadata.gz: 41ae7a1c3b5c39e29b698a402f7c336d77de4a2f69c2cf92c2512624c788d6e37f61300a66b570cc5148b34100cd77e8e136332df81a64bea1395501de992a43
+  data.tar.gz: 21534f4d98783c90a57d54bcbb978864711206dc24686f275c286de031a79acbd668e1ea01ce4799a5b884821f50f07c579fb1a99dc55f6c81d26bad5713bbc2

data/bin/pgcp

@@ -0,0 +1,9 @@
+#!/usr/bin/env ruby
+require 'yaml'
+require 'thor'
+require 'active_support'
+require 'active_support/core_ext'
+require './lib/transport'
+require './lib/pgcp'
+
+PgcpRunner.start

data/lib/pgcp.rb

@@ -0,0 +1,67 @@
+require 'logger'
+
+class PgcpRunner < Thor
+  desc 'cp', 'Perform copies of tables between Postgres databases'
+  method_option :source, type: :string, aliases: '-s', desc: 'Source database', required: true
+  method_option :dest, type: :string, aliases: '-d', desc: 'Destination database', required: true
+  method_option :table, type: :string, aliases: '-t', desc: 'Table to be copied', required: true
+  method_option :config, type: :string, aliases: '-c', desc: 'Path to config file'
+  method_option :force_schema, type: :string, aliases: '-f', desc: 'Force destination schema'
+  method_option :log, type: :string, aliases: '-l', desc: 'Path to log file'
+
+  def cp
+    config = load_config_file(options['config'] || File.join(ENV['HOME'], '.pgcp.yml'))
+    if options['log']
+      Pgcp.log_file = options['log']
+    end
+
+    src = config['databases'][options['source']].symbolize_keys!
+    dest = config['databases'][options['dest']].symbolize_keys!
+
+    begin
+      tr = Transport.new(src, dest)
+      if options['table'].include? '*'
+        if (not options['table'].include? '.') or (options['table'].count('.') > 1)
+          Pgcp.logger.error 'Globbed tables must have schema name, e.g. public.test* is valid but test* is not.'
+          return
+        end
+
+        tr.copy_tables(options['table'], force_schema: options['force_schema'])
+      else
+        tr.copy_table(options['table'], nil, force_schema: options['force_schema'])
+      end
+    rescue Exception => e
+      Pgcp.logger.error(e.message)
+      return
+    end
+  end
+
+  default_task :cp
+
+  private
+  def load_config_file(path)
+    config = {}
+    if not path.nil? and File.exists?(path)
+      config = YAML::load_file(path)
+    end
+
+    config
+  end
+end
+
+
+class Pgcp
+  @@logger = nil
+
+  def self.logger
+    if not @@logger
+      @@logger = Logger.new STDOUT
+    end
+
+    @@logger
+  end
+
+  def self.log_file=(path)
+    @@logger = Logger.new(path)
+  end
+end

data/lib/postgres.rb

@@ -0,0 +1,347 @@
+require 'active_support'
+require 'active_support/core_ext'
+require 'pg'
+
+class Postgres
+  # Initialize Postgres instance
+  #
+  # @param [Hash] dbconfig Database config
+  def initialize(dbconfig)
+    @dbconfig = dbconfig
+  end
+
+  def exec(sql, val=[])
+    with_connection do |conn|
+      conn.exec sql, val
+
+    end
+  end
+
+  def list_tables(schema_name)
+    with_connection do |conn|
+      sql = <<-SQL.strip_heredoc
+        SELECT table_name
+        FROM information_schema.tables
+        WHERE table_type = 'BASE TABLE'
+          AND table_schema NOT IN ('pg_catalog', 'information_schema')
+          AND table_schema = '#{schema_name}'
+        ORDER BY 1
+      SQL
+
+      rs = conn.exec sql
+
+      rs.values.map(&:first)
+    end
+
+  end
+
+  def drop_table(schema_name, table_name)
+    with_connection do |conn|
+      if _table_exists?(conn, schema_name, table_name)
+        fq_table_name = conn.escape_string("#{schema_name}.#{table_name}")
+
+        sql = <<-SQL.strip_heredoc
+          DROP TABLE #{fq_table_name}
+        SQL
+
+        conn.exec sql
+      end
+    end
+  end
+
+  def create_table(schema_name, table_name, columns, options={})
+    with_connection do |conn|
+      if options[:temporary]
+        table_name = conn.escape_string(table_name)
+        create_sql = create_table_statement(conn, columns,
+                                            table_name,
+                                            options)
+        conn.transaction do
+          conn.exec create_sql
+          conn.exec "DROP TABLE IF EXISTS #{table_name}"
+        end
+
+        true
+      else
+        unless _table_exists?(conn, schema_name, table_name)
+          create_sql = create_table_statement(conn, columns,
+                                              "#{schema_name}.#{table_name}",
+                                              options)
+          conn.exec create_sql
+
+          true
+        end
+
+        false
+      end
+    end
+  end
+
+  def create_table_from_query(query, schema_name, table_name, columns, options={})
+    with_connection do |conn|
+      create_sql = create_table_statement(conn, columns,
+                                          "#{schema_name}.#{table_name}",
+                                          options)
+      conn.transaction do
+        conn.exec "DROP TABLE IF EXISTS #{schema_name}.#{table_name}"
+        conn.exec create_sql
+        conn.exec "INSERT INTO #{schema_name}.#{table_name}\n#{query}"
+      end
+    end
+  end
+
+  def hotswap_table(schema_name, src_table_name, dst_table_name)
+    with_connection do |conn|
+      conn.transaction do
+        schema_name = conn.escape_string(schema_name)
+        dst_table_name = conn.escape_string(dst_table_name)
+        conn.exec "DROP TABLE #{schema_name}.#{dst_table_name}" if _table_exists?(conn, schema_name, dst_table_name)
+        conn.exec "ALTER TABLE #{schema_name}.#{src_table_name} RENAME TO #{dst_table_name}"
+        #conn.exec "TRUNCATE TABLE #{schema_name}.#{dst_table_name}" if _table_exists?(conn, schema_name, dst_table_name)
+        #conn.exec "INSERT INTO #{schema_name}.#{dst_table_name}\n(SELECT * FROM #{schema_name}.#{src_table_name})"
+        #conn.exec "DROP TABLE #{schema_name}.#{src_table_name}"
+      end
+
+    end
+  end
+
+  def schema_names
+    with_connection do |conn|
+      sql = <<-SQL.strip_heredoc
+      SELECT schema_name
+      FROM information_schema.schemata
+      WHERE schema_name <> 'information_schema'
+      AND schema_name NOT LIKE 'pg_%'
+      SQL
+
+      rs = conn.exec sql
+      rs.values.map(&:first)
+    end
+  end
+
+  def copy_from_file(schema_name, table_name, csv_file, options={})
+    with_connection do |conn|
+      schema_name = conn.escape_string(schema_name)
+      table_name = conn.escape_string(table_name)
+
+      conn.copy_data "COPY #{schema_name}.#{table_name} FROM STDIN CSV #{options[:header]?'HEADER':''}" do
+        buf = ''
+        while csv_file.read(256, buf)
+          conn.put_copy_data(buf)
+        end
+      end
+
+    end
+  end
+
+  def column_definitions(schema_name, table_name)
+    with_connection do |conn|
+      sql = <<-SQL.strip_heredoc
+        SELECT 
+          c.relname, a.attname AS column_name,
+          pg_catalog.format_type(a.atttypid, a.atttypmod) as type,
+          case 
+            when a.attnotnull
+          then 'NOT NULL' 
+          else 'NULL' 
+          END as not_null 
+        FROM pg_class c,
+         pg_attribute a,
+         pg_type t,
+         pg_namespace n
+         WHERE c.relname = '#{table_name}'
+         AND n.nspname = '#{schema_name}'
+         AND a.attnum > 0
+         AND a.attrelid = c.oid
+         AND a.atttypid = t.oid
+         AND c.relnamespace = n.oid
+       ORDER BY a.attnum
+      SQL
+
+      rs = conn.exec sql
+
+      rs.values.map do |col|
+        {name: col[1], type: col[2], null: col[3]}
+      end
+    end
+  end
+
+  def index_names(schema_name, table_name)
+    with_connection do |conn|
+      sql = <<-SQL.strip_heredoc
+        SELECT
+            C.relname AS "index_name"
+        FROM pg_catalog.pg_class C,
+             pg_catalog.pg_namespace N,
+             pg_catalog.pg_index I,
+             pg_catalog.pg_class C2
+        WHERE C.relkind IN ( 'i', '' )
+          AND N.oid = C.relnamespace
+          AND N.nspname = '#{schema_name}'
+          AND I.indexrelid = C.oid
+          AND C2.oid = I.indrelid
+          AND C2.relname = '#{table_name}';
+      SQL
+
+      rs = conn.exec sql
+
+      rs.values.map(&:first)
+    end
+  end
+
+  def get_indexes(schema_name, table_name)
+    idx_names = self.index_names(schema_name, table_name)
+
+    with_connection do |conn|
+      idx_names.map do |name|
+        index_info = index_info(conn, name, schema_name)
+        index_info['name'] = name
+        index_info['columns'] = index_column_names conn, index_info['oid']
+
+        index_info
+      end
+    end
+  end
+
+  def index_info(conn, index_name, schema_name)
+    sql = <<-SQL.strip_heredoc
+    SELECT
+        C.oid,
+        I.indisunique AS "unique",
+        I.indisprimary AS "primary",
+        pg_get_expr(I.indpred, I.indrelid) AS "where"
+    FROM pg_catalog.pg_class C,
+         pg_catalog.pg_namespace N,
+         pg_catalog.pg_index I
+    WHERE C.relname = '#{index_name}'
+      AND C.relnamespace = N.oid
+      AND I.indexrelid = C.oid
+      AND N.nspname = '#{schema_name}';
+    SQL
+
+    rs = conn.exec sql
+    rs[0].tap do |info|
+      info['unique'] = info['unique'] != 'f'
+      info['primary'] = info['primary'] != 'f'
+      info['where'] = info['where'][1..-2] if info['where'].present?
+    end
+  end
+
+  def index_column_names conn, oid
+    sql = <<-SQL.strip_heredoc
+    SELECT
+         pg_catalog.pg_get_indexdef(A.attrelid, A.attnum, TRUE) AS "column_name"
+    FROM pg_catalog.pg_attribute A
+    WHERE A.attrelid = $1
+      AND A.attnum > 0
+      AND NOT A.attisdropped
+    ORDER BY A.attnum;
+    SQL
+    conn.exec(sql, [oid]).map { |row| row['column_name'] }
+  end
+
+  def create_indexes(schema_name, table_name, indexes)
+    with_connection do |conn|
+      indexes.each do |index|
+        if index['primary']
+          sql = <<-SQL.strip_heredoc
+            ALTER TABLE #{schema_name}.#{table_name} ADD PRIMARY KEY (#{index['columns'][0]})
+          SQL
+        else
+          sql = <<-SQL.strip_heredoc
+          CREATE #{index['unique'] ? 'UNIQUE': ''} INDEX #{index['name']}
+          ON #{schema_name}.#{table_name} (#{index['columns'].join(', ')})
+          #{index['where'] ? 'WHERE ' + index['where'] : ''}
+          SQL
+        end
+
+        conn.exec(sql)
+      end
+    end
+  end
+
+  def get_create_table_statement(src_schema_name, src_table_name, dest_schema_name=nil, dest_table_name=nil)
+    dest_schema_name ||= src_schema_name
+    dest_table_name ||= dest_schema_name
+
+    columns = column_definitions(src_schema_name, src_table_name)
+
+    statement = "CREATE TABLE #{dest_schema_name}.#{dest_table_name} (\n"
+    columns.each_with_index do |col, index|
+      statement << "  #{col[:name]}  #{col[:type]}  #{col[:null]}"
+      statement << ',' if index != columns.size - 1
+      statement << "\n"
+    end
+    statement << ");\n"
+
+    statement
+  end
+
+  def table_exist?(schema_name, table_name)
+    with_connection do |conn|
+      _table_exists?(conn, schema_name, table_name)
+    end
+  end
+
+  private
+
+  def with_connection(&block)
+    conn = get_connection
+
+    block.call(conn)
+  ensure
+    conn.close if not conn.nil?
+  end
+
+  def get_connection
+    PG::Connection.connect(@dbconfig)
+  end
+
+  def column_line(column)
+    name, data_type, nullable = column.symbolize_keys.values_at(:column_name, :data_type, :is_nullable)
+
+    # default type to varchar
+    data_type ||= "VARCHAR(1000)"
+
+    line_tokens = ["\"#{name}\""]
+    line_tokens << data_type
+    line_tokens << (nullable ? '' : 'NOT NULL')
+
+    line_tokens
+      .select { |token| token != '' }
+      .join " "
+  end
+
+  def create_table_statement(connection, columns, table_name, options={})
+    statement = "CREATE #{options[:temporary] ? 'TEMPORARY' : ''} TABLE #{connection.escape_string(table_name)} (\n"
+    statement << columns
+      .map { |column| column_line(column) }
+      .map(&:strip)
+      .map { |column| connection.escape_string(column) }
+      .join(",\n")
+    statement << "\n);"
+
+    statement
+  end
+
+  def _table_exists?(connection, schema_name, table_name)
+    sql = <<-SQL.strip_heredoc
+        SELECT
+          count(table_name)
+        FROM
+          information_schema.tables
+        WHERE
+          table_schema <> 'pg_catalog'
+          AND table_schema <> 'information_schema'
+          AND table_schema !~ '^pg_toast'
+          AND table_schema = '#{connection.escape_string(schema_name)}'
+          AND table_name = '#{connection.escape_string(table_name)}'
+        GROUP BY
+          table_schema,table_name;
+    SQL
+
+    res = connection.exec(sql)
+
+    res.values.size > 0
+  end
+end

data/lib/transport.rb

@@ -0,0 +1,138 @@
+require 'active_support'
+require 'active_support/core_ext'
+require './lib/postgres'
+require './lib/qualified_name'
+require './lib/pgcp'
+require 'securerandom'
+
+class Transport
+  # Initialize Transport instance
+  #
+  # @param [Hash] src_dbconfig Source database config
+  # @param [Hash] dest_dbconfig Destination database config
+  def initialize(src_dbconfig, dest_dbconfig, options={})
+    @src_dbconfig = src_dbconfig
+    @src_dbconfig[:port] ||= 5432
+    @dest_dbconfig = dest_dbconfig
+    @dest_dbconfig[:port] ||= 5432
+  end
+
+  def copy_tables(src_tablenames, options={})
+    schema_name = src_tablenames.split('.')[0]
+    table_glob = src_tablenames.split('.')[1]
+
+    dest_conn = Postgres.new(@dest_dbconfig)
+    tables = dest_conn.list_tables(schema_name)
+    tables.each do |table|
+      if File.fnmatch(table_glob, table)
+        copy_table("#{schema_name}.#{table}", options)
+      end
+    end
+  end
+
+  def copy_table(src_tablename, dest_tablename=nil, options={})
+    dest_tablename ||= src_tablename
+    options[:create_schema] = true if options[:create_schema].nil?
+    options[:skip_indexes] = false if options[:skipe_indexes].nil?
+
+    Pgcp.logger.info "Start to copy from table #{src_tablename} to table #{dest_tablename}"
+    src_table = QualifiedName.new(src_tablename)
+    dest_table = QualifiedName.new(dest_tablename)
+    dest_table.schema_name = options[:force_schema] if not options[:force_schema].nil?
+
+    src_conn = Postgres.new(@src_dbconfig)
+    dest_conn = Postgres.new(@dest_dbconfig)
+
+    dest_conn.exec "CREATE SCHEMA IF NOT EXISTS #{dest_table.schema_name};" if options[:create_schema]
+
+    src_indexes = src_conn.get_indexes(src_table.schema_name, src_table.table_name)
+    if dest_conn.table_exist?(src_table.schema_name, src_table.table_name)
+      Pgcp.logger.info "Destination table already exists, creating temporary table"
+      temp_table = QualifiedName.new("#{dest_table.schema_name}.temp_#{SecureRandom.hex}")
+      create_table_statement =
+        src_conn.get_create_table_statement(src_table.schema_name,
+                                            src_table.table_name,
+                                            temp_table.schema_name,
+                                            temp_table.table_name)
+        begin
+          dest_conn.exec(create_table_statement)
+          Pgcp.logger.info "Copying table data to temporary table. This could take a while..."
+          direct_copy(src_table.full_name, temp_table.full_name)
+          Pgcp.logger.info "Hotswapping to destination table #{dest_tablename}"
+          dest_conn.hotswap_table(dest_table.schema_name, temp_table.table_name, dest_table.table_name)
+          Pgcp.logger.info "Done copying table data."
+        rescue Exception => e
+          Pgcp.logger.error(e.message)
+          return
+        ensure
+          dest_conn.drop_table(temp_table.schema_name, temp_table.table_name)
+        end
+    else
+      Pgcp.logger.info "Destination table does not exist, creating destination table."
+      create_table_statement =
+        src_conn.get_create_table_statement(src_table.schema_name,
+                                            src_table.table_name,
+                                            dest_table.schema_name,
+                                            dest_table.table_name)
+        dest_conn.exec(create_table_statement)
+        Pgcp.logger.info "Copying table data to destination table. This could take a while..."
+        direct_copy(src_table.full_name, dest_table.full_name)
+        Pgcp.logger.info "Copying table data to destination table done."
+    end
+
+    unless options[:skip_indexes]
+      Pgcp.logger.info "Copying table indexes to destination table..."
+      dest_conn.create_indexes(dest_table.schema_name, dest_table.table_name, src_indexes)
+      Pgcp.logger.info "Done copying table indexes."
+    end
+
+  end
+
+  private
+
+  def direct_copy(src_tablename, dest_tablename)
+    sql_in = sql_copy_from_stdin(dest_tablename)
+    sql_out = sql_copy_to_stdout(src_tablename)
+    command = transfer_command(@src_dbconfig, @dest_dbconfig, sql_in, sql_out)
+    `#{command}`
+    exit_code = `echo $?`.to_i
+    if exit_code != 0
+      raise 'Failed to directly copy data'
+    end
+  end
+
+  def sql_copy_from_stdin(q_tablename)
+    <<-SQL.strip_heredoc
+      COPY #{q_tablename} FROM STDIN
+    SQL
+  end
+
+  def sql_copy_to_stdout(q_tablename)
+    <<-SQL.strip_heredoc
+      COPY (SELECT * FROM #{q_tablename}) TO STDOUT
+    SQL
+  end
+
+  def transfer_command(src_dbconfig, dest_dbconfig, sql_in, sql_out)
+    copy_to_command = %Q{
+      env PGPASSWORD="#{src_dbconfig[:password]}"
+      psql
+        -U #{src_dbconfig[:user]}
+        -h #{src_dbconfig[:host]}
+        -p #{src_dbconfig[:port]}
+        -c "#{sql_out}"
+    #{src_dbconfig[:dbname]}
+    }.gsub(/\n/, ' ')
+    copy_from_command = %Q{
+      env PGPASSWORD="#{dest_dbconfig[:password]}"
+      psql
+        -U #{dest_dbconfig[:user]}
+        -h #{dest_dbconfig[:host]}
+        -p #{dest_dbconfig[:port]}
+        -c "#{sql_in}"
+    #{dest_dbconfig[:dbname]}
+    }.gsub(/\n/, ' ')
+
+    "#{copy_to_command} | #{copy_from_command}"
+  end
+end

metadata

@@ -0,0 +1,92 @@
+--- !ruby/object:Gem::Specification
+name: pgcp
+version: !ruby/object:Gem::Version
+  version: 0.0.0
+platform: ruby
+authors:
+- Thanh Dinh Khac
+- Huy Nguyen
+autorequire: 
+bindir: bin
+cert_chain: []
+date: 2015-11-05 00:00:00.000000000 Z
+dependencies:
+- !ruby/object:Gem::Dependency
+  name: activesupport
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '4.2'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '4.2'
+- !ruby/object:Gem::Dependency
+  name: thor
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '0.19'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '0.19'
+- !ruby/object:Gem::Dependency
+  name: pg
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '0.18'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '0.18'
+description: A simple command line tool to copy tables from one Postgres database
+  to another
+email: thanh@holistics.io
+executables:
+- pgcp
+extensions: []
+extra_rdoc_files: []
+files:
+- bin/pgcp
+- lib/pgcp.rb
+- lib/postgres.rb
+- lib/transport.rb
+homepage: http://rubygems.org/gems/pgcp
+licenses:
+- GPL
+metadata: {}
+post_install_message: 
+rdoc_options: []
+require_paths:
+- lib
+required_ruby_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '0'
+required_rubygems_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '0'
+requirements: []
+rubyforge_project: 
+rubygems_version: 2.4.8
+signing_key: 
+specification_version: 4
+summary: A simple command line tool to copy tables from one Postgres database to another
+test_files: []