pg_logical_replicator 0.1.2 → 0.1.5

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 852eb5d62a8dc6fffbb614962fe51b7566861273c52726bf546c20c74844b6e4
-  data.tar.gz: a266b7f03ee0e9ceda86c9acaeb50d4d03d00c75b40a05fb23b73f6adef02d60
+  metadata.gz: 2aab5bea887e74a1be5a02da58eb98b525c7ce3c786c7943136332f1b953c203
+  data.tar.gz: 201bed281909128a3bf8569360eb30f2d8d85601d9696207fd31919df5a45155
 SHA512:
-  metadata.gz: f7d0119c23c3b190a28851c7c976817aba9393cbd8146624c42f3456b78f047be4afb62059109de1c1172292df0fdf9a92ef4cfd6925297bb145fbba52a750e8
-  data.tar.gz: 5de087ec30bf35dc712e0879fe32f30bea1f13ae408ab356b95cd815b290efa8eef8dec7d79516172b491cbd7e9083391a727d6a0a31e4fa1a6c43fe94c6d92f
+  metadata.gz: 0d41c1e2e0e75e51ed0dccaeffd47efdb6fd7900c720ab0203e9ac7a2c0247b3d73071b3ae91fff3d29b4b9fd3b32b1598cf7be3dfa3c59248ee75c4e685f4a7
+  data.tar.gz: 9bfa6c1923a933b4ac6d5c8c094da0de2a2230cea61b4d1957f38ecfc9ea8916afe1a4619e32f8abe4711ecd165dfdc833007546644397440a0656f04c1aba36

data/lib/pg_logical_replicator/cli.rb

@@ -71,6 +71,9 @@ module PgLogicalReplicator
 
         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}",

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.2"
+  VERSION = "0.1.5"
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: pg_logical_replicator
 version: !ruby/object:Gem::Version
-  version: 0.1.2
+  version: 0.1.5
 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