pcrd 0.1.0

data/lib/pcrd/advisory_lock.rb

@@ -0,0 +1,50 @@
+# frozen_string_literal: true
+
+module Pcrd
+  # A PostgreSQL session-level advisory lock used to stop two `pcrd migrate`
+  # processes from running against the same replication slot at once — which
+  # would corrupt checkpoint/LSN progress and fight over the slot.
+  #
+  # The lock is taken on the source database (where the slot and publication
+  # live, the truly shared resource) and is keyed by the slot name. Being
+  # session-level, it is released by #release or automatically when the
+  # connection closes, so a crashed run does not leave it stuck.
+  class AdvisoryLock
+    NAMESPACE = "pcrd-migrate"
+
+    def initialize(pool:, name:)
+      @pool = pool
+      @name = name
+      @held = false
+    end
+
+    # Tries to take the lock without blocking. Returns true if acquired, false
+    # if another session already holds it.
+    def try_acquire
+      row = @pool.exec("SELECT pg_try_advisory_lock(hashtext($1)::bigint) AS locked", [key])
+      @held = (row[0]["locked"] == "t")
+      @held
+    end
+
+    # Releases the lock if held. Best-effort: a closed connection has already
+    # dropped it.
+    def release
+      return unless @held
+
+      @pool.exec("SELECT pg_advisory_unlock(hashtext($1)::bigint)", [key])
+      @held = false
+    rescue Connection::Error
+      nil
+    end
+
+    def held?
+      @held
+    end
+
+    private
+
+    def key
+      "#{NAMESPACE}:#{@name}"
+    end
+  end
+end

data/lib/pcrd/apply/engine.rb

@@ -0,0 +1,184 @@
+# frozen_string_literal: true
+
+module Pcrd
+  module Apply
+    # Applies transactions from the WAL consumer queue to the target cluster.
+    #
+    # Each transaction is a Replication::Consumer::Transaction containing a list
+    # of Insert/Update/Delete events. Events for tables not in the migration spec
+    # are silently skipped (other tables may be in the publication).
+    #
+    # INSERT events use ON CONFLICT DO UPDATE so they are safe during the
+    # backfill/streaming overlap window — if backfill already wrote a row,
+    # the WAL replay will update it to the latest version instead of failing.
+    #
+    # UPDATE events are implemented as upserts (same SQL as INSERT) because:
+    #   - The row may not yet exist on the target (if the WAL event precedes
+    #     the backfill batch that covers that key range).
+    #   - Upsert semantics are always correct here.
+    #
+    # DELETE events use the primary-key values from old_tuple (key columns).
+    class Engine
+      # Per-table execution plan built at initialisation time.
+      TablePlan = Data.define(
+        :table_name,
+        :transformer,
+        :pk_source_cols,   # Array<String>: pk column names in source schema
+        :pk_target_cols,   # Array<String>: pk column names after renames
+        :upsert_sql,       # prebuilt SQL string
+        :delete_sql        # prebuilt SQL string
+      )
+
+      def initialize(target_pool:, config:, parser:, source_schema:)
+        @target_pool   = target_pool
+        @parser        = parser
+        @plans         = build_plans(config, source_schema)
+      end
+
+      # Applies one complete transaction to the target inside a single DB transaction.
+      # Returns the commit LSN string.
+      def apply(txn)
+        @target_pool.transaction do
+          txn.events.each { |event| apply_event(event) }
+        end
+        txn.commit_lsn
+      end
+
+      private
+
+      def apply_event(event)
+        rel  = @parser.relation(event.relation_id)
+        return unless rel
+
+        # Route by schema-qualified relation name. Keying on the bare table name
+        # would mis-route events when two schemas hold a same-named table that
+        # the publication happens to include.
+        plan = @plans[relation_key(rel.namespace, rel.name)]
+        return unless plan
+
+        case event
+        when Replication::Pgoutput::Messages::Insert
+          # INSERT always carries every column ('t'/'n', never 'u'), so the
+          # full-row upsert is always safe here.
+          apply_upsert(plan, event.new_tuple)
+        when Replication::Pgoutput::Messages::Update
+          apply_update(plan, event.new_tuple)
+        when Replication::Pgoutput::Messages::Delete
+          apply_delete(plan, event.old_tuple)
+        end
+      end
+
+      def apply_upsert(plan, tuple)
+        transformed = plan.transformer.transform(tuple)
+        @target_pool.exec(plan.upsert_sql, transformed.values)
+      end
+
+      # An UPDATE's new tuple may contain :toast sentinels for TOASTed columns
+      # whose value did not change — PostgreSQL does not re-send those values.
+      # Writing the sentinel through the upsert would corrupt the column with a
+      # literal "toast". When no column is unchanged-TOAST, the full-row upsert
+      # is correct and idempotent; otherwise we emit a partial UPDATE that sets
+      # only the changed columns, leaving the existing target value in place.
+      def apply_update(plan, tuple)
+        transformed = plan.transformer.transform(tuple)
+        if transformed.value?(:toast)
+          apply_partial_update(plan, transformed)
+        else
+          @target_pool.exec(plan.upsert_sql, transformed.values)
+        end
+      end
+
+      # Builds an UPDATE that excludes unchanged-TOAST columns (and the PK) from
+      # the SET list, keyed by primary key. If the row has not been backfilled
+      # yet this updates zero rows, which is fine: backfill reads live rows and
+      # will copy the current value later, and replayed upserts are idempotent,
+      # so the target still converges.
+      def apply_partial_update(plan, transformed)
+        set_cols = transformed.reject do |col, val|
+          val == :toast || plan.pk_target_cols.include?(col)
+        end
+        return if set_cols.empty? # only PK and unchanged-TOAST columns present
+
+        assignments = set_cols.keys.each_with_index
+                              .map { |c, i| "#{Sql.quote_ident(c)} = $#{i + 1}" }
+                              .join(", ")
+        where = plan.pk_target_cols.each_with_index
+                    .map { |c, i| "#{Sql.quote_ident(c)} = $#{set_cols.size + i + 1}" }
+                    .join(" AND ")
+        sql      = "UPDATE #{Sql.quote_table(plan.table_name)} SET #{assignments} WHERE #{where}"
+        pk_vals  = plan.pk_target_cols.map { |c| transformed[c] }
+        @target_pool.exec(sql, set_cols.values + pk_vals)
+      end
+
+      def apply_delete(plan, tuple)
+        pk_values = plan.pk_source_cols.map { tuple[_1] }
+        @target_pool.exec(plan.delete_sql, pk_values)
+      end
+
+      # ── plan building ────────────────────────────────────────────────────
+
+      # Tables configured today are all in the public schema (there is no
+      # per-table schema field yet). When that lands, key off table_config.schema.
+      DEFAULT_SCHEMA = "public"
+
+      def relation_key(namespace, name)
+        "#{namespace}.#{name}"
+      end
+
+      def build_plans(config, source_schema)
+        (config.migrate&.tables || []).each_with_object({}) do |table_config, plans|
+          schema = source_schema[table_config.name]
+          next unless schema
+
+          source_cols = schema[:columns]
+          pk_source   = schema[:pk_columns]
+          transformer = Transform::RowTransformer.new(table_config, source_cols)
+          pk_target   = map_pk_to_target(pk_source, table_config)
+          target_cols = transformer.target_column_names
+
+          plans[relation_key(DEFAULT_SCHEMA, table_config.name)] = TablePlan.new(
+            table_name:     table_config.name,
+            transformer:    transformer,
+            pk_source_cols: pk_source,
+            pk_target_cols: pk_target,
+            upsert_sql:     build_upsert_sql(table_config.name, target_cols, pk_target),
+            delete_sql:     build_delete_sql(table_config.name, pk_target)
+          )
+        end
+      end
+
+      def map_pk_to_target(pk_source_cols, table_config)
+        pk_source_cols.map do |src|
+          spec = table_config.columns&.[](src) || table_config.columns&.[](src.to_sym)
+          spec&.rename || src
+        end
+      end
+
+      def build_upsert_sql(table_name, target_cols, pk_target_cols)
+        tbl  = Sql.quote_table(table_name)
+        cols = Sql.quote_columns(target_cols)
+        phs  = target_cols.each_index.map { "$#{_1 + 1}" }.join(", ")
+        pk   = Sql.quote_columns(pk_target_cols)
+
+        set_pairs = target_cols
+          .reject { |c| pk_target_cols.include?(c) }
+          .map    { |c| "#{Sql.quote_ident(c)} = EXCLUDED.#{Sql.quote_ident(c)}" }
+          .join(", ")
+
+        if set_pairs.empty?
+          "INSERT INTO #{tbl} (#{cols}) VALUES (#{phs}) ON CONFLICT (#{pk}) DO NOTHING"
+        else
+          "INSERT INTO #{tbl} (#{cols}) VALUES (#{phs}) ON CONFLICT (#{pk}) DO UPDATE SET #{set_pairs}"
+        end
+      end
+
+      def build_delete_sql(table_name, pk_target_cols)
+        tbl  = Sql.quote_table(table_name)
+        cond = pk_target_cols.each_with_index
+                             .map { |c, i| "#{Sql.quote_ident(c)} = $#{i + 1}" }
+                             .join(" AND ")
+        "DELETE FROM #{tbl} WHERE #{cond}"
+      end
+    end
+  end
+end

data/lib/pcrd/apply/worker.rb

@@ -0,0 +1,97 @@
+# frozen_string_literal: true
+
+module Pcrd
+  module Apply
+    # Drains buffered transactions from the WAL consumer queue and applies them
+    # to the target, in its own thread, concurrently with backfill.
+    #
+    # This is what makes streaming run *alongside* the bulk copy instead of
+    # after it: the consumer fills a bounded queue, this worker empties it, and
+    # the source slot can keep advancing so WAL is not retained for the whole
+    # backfill.
+    #
+    # Threading contract:
+    #   - The Apply::Engine here MUST use a target connection that is not shared
+    #     with backfill — a Connection::Client wraps a single PG connection and is
+    #     not safe to use from two threads at once.
+    #   - on_committed is invoked (on this thread) after each transaction is
+    #     durably applied, with the commit LSN. Wire it to checkpoint + the
+    #     consumer's LSN acknowledgement so WAL is only released after apply.
+    #
+    # Lifecycle:
+    #   start            — launch the background thread
+    #   stop             — drain whatever is already queued, then exit and join
+    #   failed?/error    — surface a fatal apply error to the supervising thread
+    #   last_applied_lsn — most recent commit LSN handed to on_committed
+    class Worker
+      POLL_INTERVAL = 0.05 # seconds to wait when the queue is momentarily empty
+
+      def initialize(engine:, queue:, on_committed: nil)
+        @engine       = engine
+        @queue        = queue
+        @on_committed = on_committed
+        @stop         = false
+        @mutex        = Mutex.new
+        @error        = nil
+        @last_lsn     = nil
+        @thread       = nil
+      end
+
+      def start
+        @thread = Thread.new { run_loop }
+        self
+      end
+
+      # Signals the worker to finish: it keeps applying until the queue is
+      # empty, then exits. Joins the thread before returning.
+      def stop
+        @mutex.synchronize { @stop = true }
+        @thread&.join
+      end
+
+      def failed?
+        @mutex.synchronize { !@error.nil? }
+      end
+
+      def last_applied_lsn
+        @mutex.synchronize { @last_lsn }
+      end
+
+      attr_reader :error
+
+      private
+
+      def run_loop
+        loop do
+          txn = pop_nonblocking
+
+          if txn
+            process(txn)
+          elsif stopped?
+            break # stop requested and nothing left to drain
+          else
+            sleep POLL_INTERVAL
+          end
+        end
+      rescue => e
+        @mutex.synchronize { @error = e }
+      end
+
+      def process(txn)
+        @engine.apply(txn)
+        @on_committed&.call(txn.commit_lsn)
+        @mutex.synchronize { @last_lsn = txn.commit_lsn }
+      end
+
+      def pop_nonblocking
+        @queue.pop(true)
+      rescue ThreadError
+        nil
+      end
+
+      def stopped?
+        @mutex.synchronize { @stop }
+      end
+    end
+  end
+end

data/lib/pcrd/backfill/batch.rb

@@ -0,0 +1,158 @@
+# frozen_string_literal: true
+
+module Pcrd
+  module Backfill
+    # Executes one backfill batch: SELECT a page of rows from source,
+    # transform them, and COPY to target.
+    #
+    # Returns a result hash with row_count, duration_ms, start_key, end_key.
+    # Returns nil when the source page is empty (signals end-of-table to Engine).
+    class Batch
+      # PostgreSQL COPY TEXT format: tab-delimited, \N for NULL.
+      # Using text format avoids CSV quoting edge cases and is marginally faster.
+      NULL_MARKER = "\\N"
+      DELIMITER   = "\t"
+
+      def initialize(source_pool:, target_pool:, transformer:, table_name:,
+                     pk_columns:, batch_size:, schema_name: "public")
+        @source_pool = source_pool
+        @target_pool = target_pool
+        @transformer  = transformer
+        @table_name   = table_name
+        @pk_columns   = pk_columns
+        @batch_size   = batch_size
+        @schema_name  = schema_name
+        @quoted_table = "#{source_pool.quote_ident(schema_name)}.#{source_pool.quote_ident(table_name)}"
+      end
+
+      # Copies one page starting after `after_key`.
+      # after_key: nil (first page), a scalar, or an Array for composite PKs.
+      #
+      # Returns Hash or nil.
+      def execute(after_key:)
+        t0   = monotonic_ms
+        rows = fetch_source_rows(after_key)
+        return nil if rows.empty?
+
+        transformed = rows.map { |r| @transformer.transform(r) }
+        copy_to_target(transformed)
+
+        duration_ms = monotonic_ms - t0
+
+        {
+          row_count:   rows.size,
+          duration_ms: duration_ms,
+          start_key:   extract_key(rows.first),
+          end_key:     extract_key(rows.last)
+        }
+      end
+
+      private
+
+      # ── source SELECT ────────────────────────────────────────────────────
+
+      def fetch_source_rows(after_key)
+        src_cols  = @transformer.source_column_names_kept
+        col_list  = src_cols.map { |c| @source_pool.quote_ident(c) }.join(", ")
+        pk_quoted = @pk_columns.map { |c| @source_pool.quote_ident(c) }.join(", ")
+
+        if after_key.nil?
+          sql    = "SELECT #{col_list} FROM #{@quoted_table} ORDER BY #{pk_quoted} LIMIT $1"
+          params = [@batch_size]
+        elsif @pk_columns.size == 1
+          sql    = "SELECT #{col_list} FROM #{@quoted_table} " \
+                   "WHERE #{@source_pool.quote_ident(@pk_columns.first)} > $1 " \
+                   "ORDER BY #{pk_quoted} LIMIT $2"
+          params = [after_key, @batch_size]
+        else
+          # Composite PK: row-value comparison
+          pk_placeholders = @pk_columns.each_with_index.map { |_, i| "$#{i + 1}" }.join(", ")
+          sql    = "SELECT #{col_list} FROM #{@quoted_table} " \
+                   "WHERE (#{pk_quoted}) > (#{pk_placeholders}) " \
+                   "ORDER BY #{pk_quoted} LIMIT $#{@pk_columns.size + 1}"
+          params = Array(after_key) + [@batch_size]
+        end
+
+        result = @source_pool.exec(sql, params)
+        result.to_a
+      end
+
+      # ── target COPY ──────────────────────────────────────────────────────
+
+      # COPY the batch into a session-local staging table, then merge into the
+      # real target with ON CONFLICT DO NOTHING.
+      #
+      # COPY itself has no conflict handling, so copying straight into the
+      # PK-constrained target would abort the moment the apply worker has
+      # already written a row in this key range during the backfill/streaming
+      # overlap. Merging via staging makes the bulk load idempotent and safe to
+      # run concurrently with apply: any row the worker already wrote (an
+      # insert/update replayed for a post-slot change) is left untouched. WAL
+      # replay is authoritative for changes after slot creation; backfill only
+      # fills the rows it has not seen.
+      def copy_to_target(transformed_rows)
+        target_cols = @transformer.target_column_names
+        col_list    = target_cols.map { |c| @target_pool.quote_ident(c) }.join(", ")
+
+        ensure_stage_table
+        @target_pool.exec_sql("TRUNCATE #{stage_ident}")
+
+        copy_sql = "COPY #{stage_ident} (#{col_list}) FROM STDIN WITH (FORMAT text)"
+        @target_pool.copy_data(copy_sql) do |conn|
+          transformed_rows.each do |row|
+            values = target_cols.map { |col| encode_copy_value(row[col]) }
+            conn.put_copy_data(values.join(DELIMITER) + "\n")
+          end
+        end
+
+        @target_pool.exec_sql(
+          "INSERT INTO #{quoted_target} (#{col_list}) " \
+          "SELECT #{col_list} FROM #{stage_ident} ON CONFLICT DO NOTHING"
+        )
+      end
+
+      def quoted_target
+        @quoted_target ||=
+          "#{@target_pool.quote_ident(@schema_name)}.#{@target_pool.quote_ident(@table_name)}"
+      end
+
+      # Session-local TEMP table (pg_temp resolves before the search_path, so it
+      # needs no schema qualifier). Created once per Batch and reused across this
+      # table's batches; TRUNCATEd before each load.
+      def stage_ident
+        @stage_ident ||= @target_pool.quote_ident("pcrd_stage_#{@table_name}")
+      end
+
+      def ensure_stage_table
+        return if @stage_ready
+
+        @target_pool.exec_sql(
+          "CREATE TEMP TABLE IF NOT EXISTS #{stage_ident} " \
+          "(LIKE #{quoted_target} INCLUDING DEFAULTS)"
+        )
+        @stage_ready = true
+      end
+
+      # ── helpers ──────────────────────────────────────────────────────────
+
+      def encode_copy_value(val)
+        return NULL_MARKER if val.nil?
+
+        val.to_s
+           .gsub("\\", "\\\\")
+           .gsub("\t", "\\t")
+           .gsub("\n", "\\n")
+           .gsub("\r", "\\r")
+      end
+
+      def extract_key(pg_row)
+        keys = @pk_columns.map { |col| pg_row[col] }
+        keys.size == 1 ? keys.first : keys
+      end
+
+      def monotonic_ms
+        Process.clock_gettime(Process::CLOCK_MONOTONIC, :millisecond)
+      end
+    end
+  end
+end

data/lib/pcrd/backfill/engine.rb

@@ -0,0 +1,153 @@
+# frozen_string_literal: true
+
+module Pcrd
+  module Backfill
+    # Drives the full backfill loop for all tables in the migration spec.
+    #
+    # For each table:
+    #   1. Reads last_completed_key from the checkpoint store (nil = fresh start)
+    #   2. Loops: execute one Batch, record it in checkpoint, call on_batch
+    #   3. Stops when the batch returns no rows (end of table) or stop! is called
+    #
+    # Thread safety: stop! can be called from any thread; the engine checks
+    # @stop between batches and exits cleanly after the current batch finishes.
+    class Engine
+      Result = Data.define(:table_name, :rows_copied, :batch_count, :duration_ms, :stopped_early)
+
+      # Seconds to pause after a batch to hold the average copy rate at or below
+      # `cap` rows/sec. Returns 0 when unthrottled or already slower than the cap.
+      def self.throttle_delay(row_count, duration_ms, cap)
+        return 0.0 unless cap && cap.positive?
+
+        needed  = row_count.to_f / cap
+        elapsed = duration_ms / 1000.0
+        [needed - elapsed, 0.0].max
+      end
+
+      def initialize(source_pool:, target_pool:, config:, checkpoint:, source_schema: {})
+        @source_pool   = source_pool
+        @target_pool   = target_pool
+        @config        = config
+        @checkpoint    = checkpoint
+        @source_schema = source_schema  # Hash<table_name, { columns:, pk_columns: }>
+        @stop          = false
+        @mutex         = Mutex.new
+      end
+
+      # Runs backfill for all configured tables.
+      #
+      # on_batch: optional Proc called after each batch with a stats Hash:
+      #   { table:, batch_num:, row_count:, rows_so_far:, duration_ms:, last_key: }
+      #
+      # Returns Array<Result>.
+      def run(on_batch: nil)
+        @checkpoint.set_phase(:backfill)
+        @checkpoint.set_started_at(Time.now.iso8601)
+
+        @config.migrate.tables.map do |table_config|
+          run_table(table_config, on_batch: on_batch)
+        end
+      end
+
+      # Signal the engine to stop cleanly after the current batch.
+      def stop!
+        @mutex.synchronize { @stop = true }
+      end
+
+      def stopped?
+        @mutex.synchronize { @stop }
+      end
+
+      private
+
+      def run_table(table_config, on_batch:)
+        table_name   = table_config.name
+        schema_info  = @source_schema[table_name] ||
+                       fetch_schema(table_name)
+        source_cols  = schema_info[:columns]
+        pk_cols      = schema_info[:pk_columns]
+
+        transformer = Transform::RowTransformer.new(table_config, source_cols)
+
+        batch_runner = Batch.new(
+          source_pool: @source_pool,
+          target_pool: @target_pool,
+          transformer: transformer,
+          table_name:  table_name,
+          pk_columns:  pk_cols,
+          batch_size:  @config.migrate.batch_size
+        )
+
+        last_key    = @checkpoint.last_completed_key(table: table_name)
+        batch_num   = @checkpoint.batch_stats(table: table_name)[:batch_count]
+        rows_so_far = @checkpoint.total_rows_copied(table: table_name)
+        t_start     = monotonic_ms
+
+        loop do
+          break if stopped?
+
+          result = batch_runner.execute(after_key: last_key)
+          break unless result  # empty page — end of table
+
+          batch_num   += 1
+          rows_so_far += result[:row_count]
+          last_key     = result[:end_key]
+
+          @checkpoint.record_batch(
+            table:       table_name,
+            start_key:   result[:start_key],
+            end_key:     result[:end_key],
+            row_count:   result[:row_count],
+            duration_ms: result[:duration_ms]
+          )
+
+          on_batch&.call(
+            table:       table_name,
+            batch_num:   batch_num,
+            row_count:   result[:row_count],
+            rows_so_far: rows_so_far,
+            duration_ms: result[:duration_ms],
+            last_key:    last_key
+          )
+
+          delay = self.class.throttle_delay(
+            result[:row_count], result[:duration_ms], @config.migrate.max_rows_per_second
+          )
+          interruptible_sleep(delay) if delay.positive?
+        end
+
+        Result.new(
+          table_name:    table_name,
+          rows_copied:   rows_so_far,
+          batch_count:   batch_num,
+          duration_ms:   monotonic_ms - t_start,
+          stopped_early: stopped?
+        )
+      end
+
+      def fetch_schema(table_name)
+        reader = Schema::Reader.new(@source_pool)
+        {
+          columns:    reader.read(table_name),
+          pk_columns: reader.primary_key_columns(table_name)
+        }
+      end
+
+      def monotonic_ms
+        Process.clock_gettime(Process::CLOCK_MONOTONIC, :millisecond)
+      end
+
+      # Sleeps for `seconds`, in small slices, so a throttle pause still honors
+      # a stop request promptly instead of blocking until the full delay elapses.
+      def interruptible_sleep(seconds)
+        deadline = monotonic_ms + (seconds * 1000)
+        while monotonic_ms < deadline
+          break if stopped?
+
+          remaining = (deadline - monotonic_ms) / 1000.0
+          sleep [0.1, remaining].min
+        end
+      end
+    end
+  end
+end