replication origin support

lib/pgsync/client.rb

@@ -101,6 +101,7 @@ def slop_options
       o.boolean "--in-batches", "sync in batches", default: false, help: false
       o.integer "--batch-size", "batch size", default: 10000, help: false
       o.float "--sleep", "time to sleep between batches", default: 0, help: false
+      o.string "--replication-origin", "replication origin", help: false
 
       o.separator ""
       o.separator "Other commands:"

lib/pgsync/data_source.rb

@@ -100,6 +100,25 @@ def triggers(table)
       execute(query, [quote_ident_full(table)])
     end
 
+    def set_replication_origin(origin)
+      query = <<~SQL
+        SELECT 
+          CASE
+            WHEN EXISTS (SELECT 1 FROM pg_replication_origin WHERE roname = $1) THEN null
+            ELSE pg_replication_origin_create($1)
+          END;
+      SQL
+      execute(query, [origin])
+      query = <<~SQL
+        SELECT
+          CASE
+            WHEN pg_replication_origin_session_is_setup() THEN null
+            ELSE pg_replication_origin_session_setup($1)
+          END;
+      SQL
+      execute(query, [origin])
+    end
+
     def conn
       @conn ||= begin
         begin

lib/pgsync/sync.rb

@@ -19,7 +19,7 @@ def perform
       end
 
       # merge other config
-      [:to_safe, :exclude, :schemas].each do |opt|
+      [:to_safe, :exclude, :schemas, :replication_origin].each do |opt|
         opts[opt] ||= config[opt.to_s]
       end
 

lib/pgsync/table_sync.rb

@@ -248,7 +248,9 @@ def run_tasks(tasks, &block)
         Parallel.each(tasks, **options) do |task|
           source.reconnect_if_needed
           destination.reconnect_if_needed
-
+          if opts[:replication_origin]
+            destination.set_replication_origin(opts[:replication_origin])
+          end
           task.perform
         end
       end

test/sync_test.rb

@@ -175,4 +175,11 @@ def test_disable_integrity_v2
     assert_equal [], conn2.exec("SELECT * FROM posts ORDER BY id").to_a
     assert_equal [{"post_id" => 1}], conn2.exec("SELECT post_id FROM comments ORDER BY post_id").to_a
   end
+
+  def test_replication_origin
+    insert(conn1, "posts", [{"id" => 1}])
+    assert_error "Sync failed for 1 table: posts", "posts", config: true
+    assert_works "posts --replication-origin=test --debug", config: true
+    assert_equal [{"exists" => 1}], conn2.exec("SELECT 1 AS exists FROM pg_replication_origin WHERE roname = 'test'").to_a
+  end
 end