pg_logical_replicator 0.1.3 → 0.1.6

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: b7db48910687f53995bb0481317730468d28b8d27c250c7e61434a462609c437
-  data.tar.gz: d41e788e6748bbf60eaf07ea45abf2a493d5ba6ad668562de06be582ff94a073
+  metadata.gz: d85998326734611fffa7a261b4da9e7aaa23cb0fc8e460c3e2e10206e8c2012d
+  data.tar.gz: 375f48b3a024bdcedb8a10c2752c60ea8d4acc5305dc0123d6ab993f6da5d185
 SHA512:
-  metadata.gz: 52dfcf1355d45c6b9321b73dbaedea50d2d71175230063caf17a408a52356833372bba9dad7fb66f35b6548a7417ab3602fb5035ffd5dfff37c7e425124c6dc6
-  data.tar.gz: 9d3e54057f4636619a62542072a4455a5be60a4fe1872f5ccbcdc339daf8b66fef3df08cc4a34c2dfcb8d9b805d5aa6f9d06185d5150b5aaeef25f73af884185
+  metadata.gz: 387f5aa4609c1d69f11eb996fe4ee64890178ced6b7fb9a1f53eaad0233f9050c941362dc51c1b73e7d1133bef7058898ad1375b6550553048a011d224a27eb9
+  data.tar.gz: a3a0fe506552ecd598d4a31f7958d24ea6f257fe3e073136143b3252ea24a16c57d7e517577a7221c6e994704a4b4947c58fb7ad6b1de979fd36e6906348cd43

data/lib/pg_logical_replicator/cli.rb

@@ -19,7 +19,7 @@ module PgLogicalReplicator
     method_option :target_rep_password, type: :string
     method_option :num_slots, type: :numeric, default: 10
     method_option :groups, type: :array
-
+    
     def setup
       source_host         = options[:source_host]
       source_port         = options[:source_port]
@@ -35,11 +35,11 @@ module PgLogicalReplicator
       target_rep_password = options[:target_rep_password] || target_password
       num_slots           = options[:num_slots]
       target_groups       = options[:groups].map(&:to_i) if options[:groups]
-
+      
       conn = PG.connect(dbname: source_database, user: source_username, password: source_password, host: source_host, port: source_port)
-
+      
       query = "SELECT tablename FROM pg_tables WHERE schemaname = 'public' order by tablename ASC;"
-
+      
       tables = begin
                  result = conn.exec(query)
                  result.map { |row| row['tablename'] }
@@ -47,31 +47,33 @@ module PgLogicalReplicator
                  conn.close if conn
                end
       puts "Total Tables: #{tables.size}"
-
+      
       tables_per_group = tables.size / num_slots
       puts "Tables per group: #{tables_per_group}"
-
+      
       dump_dir_root = "~/db-dumps-#{Time.now.strftime('%Y%m%d%H%M%S')}"
       system("mkdir -p #{dump_dir_root}")
-
+      
       puts "Target Groups: #{target_groups}"
-
+      
       tables.each_slice(tables_per_group).with_index do |table_group, group_idx|
         puts "Processing group #{group_idx + 1} size: #{table_group.size}"
-
+        
         group_number = group_idx + 1
-
+        
         next unless target_groups.nil? || target_groups.include?(group_number)
-
+        
         table_names = table_group.join(',')
-
+        
         dump_dir = "#{dump_dir_root}/#{group_number}"
-
+        
         puts "Removing directory: #{dump_dir}"
-
+        
         FileUtils.rm_rf(dump_dir)
+        
+        puts "Creating directory: #{dump_dir}"
         FileUtils.mkdir_p(dump_dir)
-
+        
         LogicalReplicationInitializer.new({
                                             slot_name:           "logical_sub_grp_#{group_number}",
                                             publication_name:    "logical_pub_grp_#{group_number}",
@@ -81,9 +83,13 @@ module PgLogicalReplicator
                                             table_names:         table_names,
                                             dump_dir:            dump_dir
                                           }).start
+      rescue StandardError => e
+        puts "Error: #{e.message}"
+      ensure
+        FileUtils.rm_rf(dump_dir)
       end
     end
-
+    
     desc 'stop_replication', 'Stop all replication'
     method_option :source_host, type: :string, required: true
     method_option :source_port, type: :numeric, default: 5432
@@ -95,7 +101,7 @@ module PgLogicalReplicator
     method_option :source_password, type: :string, required: true
     method_option :target_username, type: :string
     method_option :target_password, type: :string
-
+    
     def stop_replication
       source_config = {
         host:     options[:source_host],
@@ -104,7 +110,7 @@ module PgLogicalReplicator
         user:     options[:source_username],
         password: options[:source_password]
       }
-
+      
       target_config = {
         host:     options[:target_host],
         port:     options[:target_port],
@@ -112,10 +118,10 @@ module PgLogicalReplicator
         user:     options[:target_username] || options[:source_username],
         password: options[:target_password] || options[:source_password]
       }
-
+      
       ReplicationStopper.new(source_config, target_config).stop_replication
     end
-
+    
     desc 'transfer_schema', 'Transfer schema from source to target database'
     method_option :source_host, type: :string, required: true
     method_option :source_port, type: :numeric, default: 5432
@@ -127,7 +133,7 @@ module PgLogicalReplicator
     method_option :source_password, type: :string, required: true
     method_option :target_username, type: :string, required: true
     method_option :target_password, type: :string, required: true
-
+    
     def transfer_schema
       SchemaTransfer.new(
         source_host:     options[:source_host],

data/lib/pg_logical_replicator/logical_replication_initializer.rb

@@ -1,145 +1,206 @@
+#!/usr/bin/env ruby
+# frozen_string_literal: true
+
 require 'optparse'
-require 'pg'
 require 'logger'
+require 'pg'
 require 'fileutils'
 require 'open3'
 
-module PgLogicalReplicator
-  class LogicalReplicationInitializer
-    def initialize(args)
-      @options = args
-      setup_logger
-      validate_options
-      connect_to_databases
-    end
+class LogicalReplicationInitializer
+
+  attr :options
 
-    def setup_logger
-      @log = Logger.new(STDOUT)
-      @log.level = Logger::INFO
+  def initialize(args)
+    @options = args
+
+    setup_logger
+    validate_options
+    connect_to_databases
+  end
+
+  def setup_logger
+    @log       = Logger.new(STDOUT)
+    @log.level = Logger::INFO
+  end
+
+  def validate_options
+    [:primary_conn_str, :target_conn_str, :target_rep_conn_str, :table_names, :dump_dir].each do |opt|
+      raise OptionParser::MissingArgument, opt.to_s unless @options[opt]
     end
 
-    def validate_options
-      [:primary_conn_str, :target_conn_str, :target_rep_conn_str, :table_names, :dump_dir].each do |opt|
-        raise OptionParser::MissingArgument, opt.to_s unless @options[opt]
+    @primary_conn_str      = options[:primary_conn_str] || (raise OptionParser::MissingArgument, '--primary-conn-str is missing')
+    @secondary_conn_str    = options[:secondary_conn_str] || @primary_conn_str
+    @target_conn_str       = options[:target_conn_str] || (raise OptionParser::MissingArgument, '--target-conn-str is missing')
+    @target_rep_conn_str   = options[:target_rep_conn_str] || (raise OptionParser::MissingArgument, '--target-rep-conn-str is missing')
+    @table_names           = options[:table_names] || (raise OptionParser::MissingArgument, '--table-names is missing')
+    @subscription_conn_str = options[:subscription_conn_str] || @primary_conn_str
+    @dump_dir              = options[:dump_dir] || (raise OptionParser::MissingArgument, '--dump-dir is missing')
+    @slot                  = options[:slot_name]
+    @publication_name      = options[:publication_name]
+  end
+
+  def start
+    begin
+      @log.info("Dumping to #{@dump_dir}")
+      connect_to_databases
+
+      drop_subscription
+      drop_publication
+      create_publication
+      create_logical_replication_slot
+
+      lsn = with_snapshot do |lsn, snapshot|
+        @log.info("Creating Snapshot with lsn #{lsn}")
+
+        dump(snapshot)
+
+        lsn
       end
 
-      @options[:slot_name] ||= 'migration_sub'
-      @options[:publication_name] ||= 'migration_pub'
-      @options[:secondary_conn_str] ||= @options[:primary_conn_str]
-      @options[:subscription_conn_str] ||= @options[:primary_conn_str]
+      restore_dump
+      create_subscription
+      advance_subscription_origin(lsn)
+      enable_subscription
+    rescue StandardError => e
+      @log.error(e.message)
+      exit 1
     end
+  end
 
-    def connect_to_databases
-      @primary_conn = PG.connect(@options[:primary_conn_str])
-      @secondary_conn = PG.connect(@options[:secondary_conn_str])
-      @target_conn = PG.connect(@options[:target_conn_str])
-      @target_rep_conn = PG.connect(@options[:target_rep_conn_str])
-    end
+  def connect_to_databases
+    @primary    = PG.connect(@primary_conn_str)
+    @secondary  = PG.connect(@secondary_conn_str)
+    @target     = PG.connect(@target_conn_str)
+    @target_rep = PG.connect(@target_rep_conn_str)
+  end
 
-    def start
-      begin
-        @log.info("Dumping to #{@options[:dump_dir]}")
-
-        drop_subscription
-        drop_publication
-        create_publication
-        create_logical_replication_slot
-
-        snapshot = create_snapshot
-        dump_database(snapshot)
-        restore_dump
-        create_subscription
-        advance_subscription_origin(snapshot[:lsn])
-        enable_subscription
-      rescue => e
-        @log.error(e.message)
-        exit 1
+  def create_publication
+    @primary.exec("SELECT COUNT(*) FROM pg_publication WHERE pubname = $1", [@publication_name]) do |res|
+      cnt = res.getvalue(0, 0).to_i
+
+      if cnt.zero?
+        pub_name   = PG::Connection.quote_ident(@publication_name)
+        table_list = @table_names.split(',')
+
+        @primary.exec("CREATE PUBLICATION #{pub_name} FOR TABLE #{PG::Connection.quote_ident(table_list[0])}")
+
+        table_list[1..].each do |table|
+          begin
+            @log.info("Processing table: #{table}")
+            @primary.exec("ALTER PUBLICATION #{pub_name} ADD TABLE #{PG::Connection.quote_ident(table)}")
+          rescue StandardError => e
+            @log.error("Error adding table #{table}: #{e.message}")
+          end
+        end
+        @log.info("Created publication #{@publication_name}")
+      else
+        @log.info("Publication #{@publication_name} already exists")
       end
     end
+  end
 
-    def drop_subscription
-      @log.info("Dropping subscription #{@options[:slot_name]}")
-      @target_conn.exec("DROP SUBSCRIPTION IF EXISTS #{@options[:slot_name]}")
+  def create_logical_replication_slot
+    @primary.exec("SELECT pg_drop_replication_slot(slot_name) FROM pg_replication_slots WHERE slot_name = $1", [@slot])
+    @primary.exec("SELECT pg_create_logical_replication_slot($1, $2)", [@slot, "pgoutput"]) do |res|
+      slot_name = res.getvalue(0, 0)
+      @log.info("Created logical replication slot #{slot_name}")
     end
+  end
 
-    def drop_publication
-      @log.info("Dropping publication #{@options[:publication_name]}")
-      @primary_conn.exec("DROP PUBLICATION IF EXISTS #{@options[:publication_name]}")
-    end
+  def with_snapshot
+    @secondary.exec("BEGIN TRANSACTION ISOLATION LEVEL REPEATABLE READ")
 
-    def create_publication
-      @log.info("Creating publication #{@options[:publication_name]}")
-      table_names = @options[:table_names].split(',')
-      @primary_conn.exec("CREATE PUBLICATION #{@options[:publication_name]} FOR TABLE #{table_names.join(', ')}")
+    @secondary.exec("SELECT CASE WHEN pg_is_in_recovery() THEN pg_last_wal_replay_lsn() ELSE pg_current_wal_lsn() END, pg_export_snapshot()") do |res|
+      yield res.getvalue(0, 0), res.getvalue(0, 1)
     end
+  end
 
-    def create_logical_replication_slot
-      @log.info("Creating logical replication slot #{@options[:slot_name]}")
-      @primary_conn.exec("SELECT pg_create_logical_replication_slot('#{@options[:slot_name]}', 'pgoutput')")
+  def dump(snapshot)
+    @log.info("Dumping source DB")
+    table_list = @table_names.split(',')
+
+    cmd = [
+      'pg_dump',
+      '--no-publication',
+      '--no-subscription',
+      "--snapshot=#{snapshot}",
+      '--format=d',
+      '--data-only',
+      '--jobs=8',
+      '-f', @dump_dir,
+      @secondary_conn_str
+    ]
+
+    table_list.each do |table|
+      @log.info("Dumping table #{table}")
+      cmd << '-t' << table
+    end
+
+    @log.info("Executing DB Dump Command: #{cmd.join(' ')}")
+    stdout, stderr, status = Open3.capture3(*cmd)
+    unless status.success?
+      @log.error("Error in pg_dump: #{stderr}")
+      raise "pg_dump failed"
     end
+  end
 
-    def create_snapshot
-      @log.info("Creating snapshot")
-      @secondary_conn.transaction do |conn|
-        res = conn.exec("SELECT pg_current_wal_lsn(), pg_export_snapshot()")
-        { lsn: res.getvalue(0, 0), snapshot: res.getvalue(0, 1) }
-      end
-    end
+  def truncate_table(table_name)
+    @target.exec("TRUNCATE TABLE #{PG::Connection.quote_ident(table_name)}")
+  end
 
-    def dump_database(snapshot)
-      @log.info("Dumping source DB")
-      table_names = @options[:table_names].split(',')
-      dump_command = [
-        'pg_dump',
-        '--no-publication',
-        '--no-subscription',
-        "--snapshot=#{snapshot[:snapshot]}",
-        '--format=d',
-        '--data-only',
-        '--jobs=8',
-        '-f', @options[:dump_dir],
-        @options[:secondary_conn_str]
-      ] + table_names.flat_map { |table| ['-t', table] }
-
-      @log.info("Executing DB Dump Command: #{dump_command.join(' ')}")
-      system(*dump_command)
-    end
+  def restore_dump
+    table_names = @options[:table_names].split(',')
+    table_names.each { |table| truncate_table(table) }
 
-    def restore_dump
-      @log.info("Restoring dump to target DB")
-      table_names = @options[:table_names].split(',')
-      table_names.each { |table| truncate_table(table) }
-
-      restore_command = [
-        'pg_restore',
-        '--format=d',
-        '--jobs=8',
-        '--data-only',
-        '-d', @options[:target_rep_conn_str],
-        @options[:dump_dir]
-      ]
-
-      system(*restore_command)
+    @log.info("Restoring dump to target DB")
+    stdout, stderr, status = Open3.capture3('pg_restore', '--format=d', '--jobs=8', '--data-only', '-d', @target_rep_conn_str, @dump_dir)
+    unless status.success?
+      @log.error("Error in pg_restore: #{stderr}")
+      raise "pg_restore failed"
     end
+  end
 
-    def truncate_table(table)
-      @log.info("Truncating table #{table}")
-      @target_rep_conn.exec("TRUNCATE TABLE #{table}")
-    end
+  def create_subscription
+    @log.info("Creating subscription #{@slot}")
 
-    def create_subscription
-      @log.info("Creating subscription #{@options[:slot_name]}")
-      @target_conn.exec("CREATE SUBSCRIPTION #{@options[:slot_name]} CONNECTION '#{@options[:subscription_conn_str]}' PUBLICATION #{@options[:publication_name]} WITH (create_slot=false, slot_name='#{@options[:slot_name]}', enabled=false, copy_data=false)")
-    end
+    pub_name  = PG::Connection.quote_ident(@publication_name)
+    sub_name  = PG::Connection.quote_ident(@slot)
 
-    def advance_subscription_origin(lsn)
-      @log.info("Setting replication origin position to #{lsn}")
-      @target_conn.exec("SELECT pg_replication_origin_advance('pg_' || subid::text, '#{lsn}') FROM pg_stat_subscription WHERE subname = '#{@options[:slot_name]}'")
-    end
+    query = <<~SQL
+      CREATE SUBSCRIPTION #{sub_name}
+      CONNECTION '#{@subscription_conn_str}'
+      PUBLICATION #{pub_name}
+      WITH (create_slot=false, slot_name=#{sub_name}, enabled=false, copy_data=false)
+    SQL
 
-    def enable_subscription
-      @log.info("Enabling subscription #{@options[:slot_name]}")
-      @target_conn.exec("ALTER SUBSCRIPTION #{@options[:slot_name]} ENABLE")
-    end
+    @target.exec(query)
+  end
+
+  def advance_subscription_origin(lsn)
+    @log.info("Setting replication origin position to #{lsn}")
+
+    query    = <<~SQL
+      SELECT pg_replication_origin_advance('pg_' || subid::text, '#{lsn}')
+      FROM pg_stat_subscription
+      WHERE subname = '#{@slot}'
+    SQL
+
+    @target.exec(query)
+  end
+
+  def enable_subscription
+    @log.info("Enabling subscription #{@slot}")
+    @target.exec("ALTER SUBSCRIPTION #{PG::Connection.quote_ident(@slot)} ENABLE")
+  end
+
+  def drop_publication
+    @log.info("Dropping publication #{@publication_name}")
+    @primary.exec("DROP PUBLICATION IF EXISTS #{PG::Connection.quote_ident(@publication_name)}")
+  end
+
+  def drop_subscription
+    @log.info("Dropping subscription #{@slot}")
+    @target.exec("DROP SUBSCRIPTION IF EXISTS #{PG::Connection.quote_ident(@slot)}")
   end
-end
+end

data/lib/pg_logical_replicator/version.rb

@@ -1,3 +1,3 @@
 module PgLogicalReplicator
-  VERSION = "0.1.3"
+  VERSION = "0.1.6"
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: pg_logical_replicator
 version: !ruby/object:Gem::Version
-  version: 0.1.3
+  version: 0.1.6
 platform: ruby
 authors:
 - eni9889
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2024-07-14 00:00:00.000000000 Z
+date: 2024-07-15 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: pg
@@ -118,7 +118,7 @@ required_rubygems_version: !ruby/object:Gem::Requirement
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.0.3.1
+rubygems_version: 3.4.10
 signing_key:
 specification_version: 4
 summary: PostgreSQL logical replication setup tool