Add origin and timestamp preservation to table repair.

db/queries/queries.go

@@ -898,13 +898,13 @@ func GetSpockNodeNames(ctx context.Context, db DBQuerier) (map[string]string, er
 	return names, nil
 }
 
-func GetSpockOriginLSNForNode(ctx context.Context, db DBQuerier, failedNode, survivor string) (*string, error) {
+func GetSpockOriginLSNForNode(ctx context.Context, db DBQuerier, originNodeName string) (*string, error) {
 	sql, err := RenderSQL(SQLTemplates.GetSpockOriginLSNForNode, nil)
 	if err != nil {
 		return nil, err
 	}
 	var lsn *string
-	if err := db.QueryRow(ctx, sql, failedNode, survivor).Scan(&lsn); err != nil {
+	if err := db.QueryRow(ctx, sql, originNodeName).Scan(&lsn); err != nil {
 		if errors.Is(err, pgx.ErrNoRows) {
 			return nil, nil
 		}
@@ -3044,3 +3044,96 @@ func RemoveTableFromCDCMetadata(ctx context.Context, db DBQuerier, tableName, pu
 
 	return nil
 }
+
+func GetReplicationOriginByName(ctx context.Context, db DBQuerier, originName string) (*uint32, error) {
+	sql, err := RenderSQL(SQLTemplates.GetReplicationOriginByName, nil)
+	if err != nil {
+		return nil, fmt.Errorf("failed to render GetReplicationOriginByName SQL: %w", err)
+	}
+
+	var originID uint32
+	err = db.QueryRow(ctx, sql, originName).Scan(&originID)
+	if err != nil {
+		if errors.Is(err, pgx.ErrNoRows) {
+			return nil, nil
+		}
+		return nil, fmt.Errorf("query to get replication origin by name '%s' failed: %w", originName, err)
+	}
+
+	return &originID, nil
+}
+
+func CreateReplicationOrigin(ctx context.Context, db DBQuerier, originName string) (uint32, error) {
+	sql, err := RenderSQL(SQLTemplates.CreateReplicationOrigin, nil)
+	if err != nil {
+		return 0, fmt.Errorf("failed to render CreateReplicationOrigin SQL: %w", err)
+	}
+
+	var originID uint32
+	err = db.QueryRow(ctx, sql, originName).Scan(&originID)
+	if err != nil {
+		return 0, fmt.Errorf("query to create replication origin '%s' failed: %w", originName, err)
+	}
+
+	return originID, nil
+}
+
+func SetupReplicationOriginSession(ctx context.Context, db DBQuerier, originName string) error {
+	sql, err := RenderSQL(SQLTemplates.SetupReplicationOriginSession, nil)
+	if err != nil {
+		return fmt.Errorf("failed to render SetupReplicationOriginSession SQL: %w", err)
+	}
+
+	_, err = db.Exec(ctx, sql, originName)
+	if err != nil {
+		return fmt.Errorf("query to setup replication origin session for origin '%s' failed: %w", originName, err)
+	}
+
+	return nil
+}
+
+func ResetReplicationOriginSession(ctx context.Context, db DBQuerier) error {
+	sql, err := RenderSQL(SQLTemplates.ResetReplicationOriginSession, nil)
+	if err != nil {
+		return fmt.Errorf("failed to render ResetReplicationOriginSession SQL: %w", err)
+	}
+
+	_, err = db.Exec(ctx, sql)
+	if err != nil {
+		return fmt.Errorf("query to reset replication origin session failed: %w", err)
+	}
+
+	return nil
+}
+
+func SetupReplicationOriginXact(ctx context.Context, db DBQuerier, originLSN string, originTimestamp *time.Time) error {
+	if originTimestamp == nil {
+		return fmt.Errorf("origin timestamp is required for pg_replication_origin_xact_setup (LSN %s)", originLSN)
+	}
+
+	sql, err := RenderSQL(SQLTemplates.SetupReplicationOriginXact, nil)
+	if err != nil {
+		return fmt.Errorf("failed to render SetupReplicationOriginXact SQL: %w", err)
+	}
+
+	_, err = db.Exec(ctx, sql, originLSN, *originTimestamp)
+	if err != nil {
+		return fmt.Errorf("query to setup replication origin xact with LSN %s failed: %w", originLSN, err)
+	}
+
+	return nil
+}
+
+func ResetReplicationOriginXact(ctx context.Context, db DBQuerier) error {
+	sql, err := RenderSQL(SQLTemplates.ResetReplicationOriginXact, nil)
+	if err != nil {
+		return fmt.Errorf("failed to render ResetReplicationOriginXact SQL: %w", err)
+	}
+
+	_, err = db.Exec(ctx, sql)
+	if err != nil {
+		return fmt.Errorf("query to reset replication origin xact failed: %w", err)
+	}
+
+	return nil
+}

db/queries/templates.go

@@ -124,6 +124,12 @@ type Templates struct {
 	GetHashVersion                   *template.Template
 	MarkAllLeavesDirty               *template.Template
 	UpdateHashVersion                *template.Template
+	GetReplicationOriginByName       *template.Template
+	CreateReplicationOrigin          *template.Template
+	SetupReplicationOriginSession    *template.Template
+	ResetReplicationOriginSession    *template.Template
+	SetupReplicationOriginXact       *template.Template
+	ResetReplicationOriginXact       *template.Template
 }
 
 var SQLTemplates = Templates{
@@ -1537,7 +1543,9 @@ var SQLTemplates = Templates{
 		SELECT ros.remote_lsn::text
 		FROM pg_catalog.pg_replication_origin_status ros
 		JOIN pg_catalog.pg_replication_origin ro ON ro.roident = ros.local_id
-		WHERE ro.roname LIKE 'spk_%_' || $1 || '_sub_' || $1 || '_' || $2
+		JOIN spock.subscription s ON ro.roname LIKE '%' || s.sub_name
+		JOIN spock.node o ON o.node_id = s.sub_origin
+		WHERE o.node_name = $1
 			AND ros.remote_lsn IS NOT NULL
 		LIMIT 1
 	`)),
@@ -1572,4 +1580,22 @@ var SQLTemplates = Templates{
 		SET hash_version = $1, last_updated = current_timestamp
 		WHERE schema_name = $2 AND table_name = $3
 	`)),
+	GetReplicationOriginByName: template.Must(template.New("getReplicationOriginByName").Parse(`
+		SELECT roident FROM pg_replication_origin WHERE roname = $1
+	`)),
+	CreateReplicationOrigin: template.Must(template.New("createReplicationOrigin").Parse(`
+		SELECT pg_replication_origin_create($1)
+	`)),
+	SetupReplicationOriginSession: template.Must(template.New("setupReplicationOriginSession").Parse(`
+		SELECT pg_replication_origin_session_setup($1)
+	`)),
+	ResetReplicationOriginSession: template.Must(template.New("resetReplicationOriginSession").Parse(`
+		SELECT pg_replication_origin_session_reset()
+	`)),
+	SetupReplicationOriginXact: template.Must(template.New("setupReplicationOriginXact").Parse(`
+		SELECT pg_replication_origin_xact_setup($1, $2)
+	`)),
+	ResetReplicationOriginXact: template.Must(template.New("resetReplicationOriginXact").Parse(`
+		SELECT pg_replication_origin_xact_reset()
+	`)),
 }

docs/api.md

@@ -78,6 +78,7 @@ Repairs table inconsistencies using a diff file.
 | `--fix-nulls`       | `-X`  | Fill NULL columns on each node using non-NULL values from its peers  | false   |
 | `--bidirectional`   | `-Z`  | Perform insert-only repairs in both directions                       | false   |
 | `--fire-triggers`   | `-t`  | Fire triggers during repairs                                         | false   |
+| `--preserve-origin` |       | Preserve replication origin with per-row timestamp accuracy          | false   |
 | `--quiet`           | `-q`  | Suppress output                                                      | false   |
 | `--debug`           | `-v`  | Enable debug logging                                                 | false   |
 

docs/commands/repair/table-repair.md

@@ -30,6 +30,7 @@ Performs repairs on tables of divergent nodes based on the diff report generated
 | `--bidirectional` | `-Z` | Perform insert-only repairs in both directions | `false` |
 | `--fire-triggers` | `-t` | Execute triggers (otherwise runs with `session_replication_role='replica'`) | `false` |
 | `--recovery-mode` |  | Enable recovery-mode repair when the diff was generated with `--against-origin`; can auto-select a source of truth using Spock LSNs | `false` |
+| `--preserve-origin` |  | Preserve replication origin node ID and LSN with per-row timestamp accuracy for repaired rows. When enabled, repaired rows will have commits with the original node's origin ID and exact commit timestamp (microsecond precision) instead of the local node ID. Requires LSN to be available from a survivor node. | `false` |
 | `--quiet` | `-q` | Suppress non-essential logging | `false` |
 | `--debug` | `-v` | Enable verbose logging | `false` |
 
@@ -70,3 +71,53 @@ Diff reports share the same prefix generated by `table-diff` (for example `publi
 ## Fixing null-only drifts (`--fix-nulls`)
 
 Replication hiccups can leave some columns NULL on one node while populated on another. The `--fix-nulls` mode cross-fills those NULLs in both directions using values from the paired node(s); it does **not** require a source-of-truth. Use it when the diff shows only NULL/NOT NULL mismatches and you want to reconcile columns without preferring a single node.
+
+## Preserving replication origin (`--preserve-origin`)
+
+When `--preserve-origin` is enabled, repaired rows maintain the correct replication origin node ID and LSN from the original transaction, along with precise per-row timestamp preservation. This is particularly important in recovery scenarios where:
+
+- A node fails and rows are repaired from a survivor
+- The failed node may come back online
+- Without origin tracking, the repaired rows would have the local node's origin ID, which could cause conflicts when the original node resumes replication
+- Temporal ordering and conflict resolution depend on accurate commit timestamps
+
+### How it works
+
+1. **Origin extraction**: ACE extracts the `node_origin` and `commit_ts` from the diff file metadata for each row being repaired.
+
+2. **LSN retrieval**: For each origin node, ACE queries a survivor node to obtain the origin LSN. If the LSN is not available, the affected rows fall back to a regular repair without origin preservation (a warning will be logged).
+
+3. **Replication origin session**: Before executing repairs for each origin group, ACE:
+   - Gets or creates a replication origin for the origin node
+   - Sets up a replication origin session
+   - Configures the session with the origin LSN and timestamp
+   - Executes the repairs
+   - Resets the session
+
+4. **Transaction management**: When preserve-origin is active, repairs use a two-phase commit pattern (not a single atomic transaction):
+   - **Phase 1**: All deletes are executed and committed in a single transaction.
+   - **Phase 2**: Upserts are grouped by (origin node, LSN, timestamp) tuples. Each group runs in its own transaction with its replication origin xact configured (LSN and commit timestamp), then commits. This is required because `pg_replication_origin_xact_setup` is transaction-scoped.
+   - **Failure/rollback boundaries**: If a failure occurs during phase 2, deletes remain committed and any upsert batches that already committed remain committed; only the current batch rolls back. The repair is not all-or-nothing.
+
+5. **Timestamp precision**: Timestamps are stored in RFC3339Nano format (e.g., `2026-01-15T14:23:45.123456Z`) to preserve microsecond-level accuracy. This precision is essential when:
+   - Multiple transactions occurred within the same second on the origin node
+   - Conflict resolution depends on precise temporal ordering
+   - Recovery scenarios require exact timestamp matching for conflict-free reintegration
+
+### Requirements and limitations
+
+- **LSN availability**: The origin LSN should be available from at least one survivor node. If not available, the affected rows fall back to a regular repair without origin preservation.
+- **Survivor nodes**: At least one survivor node must be accessible to fetch the origin LSN.
+- **Privileges**: Replication origin functions require superuser or replication privileges on the target database.
+- **Missing metadata**: If origin metadata is missing from the diff file for some rows, those rows will be repaired without origin tracking (a warning will be logged).
+
+### When to use
+
+Enable `--preserve-origin` in recovery scenarios where:
+- You want to prevent replication conflicts when the origin node returns
+- You need to maintain the original transaction timestamps and origin metadata
+
+Origin preservation is off by default. Consider leaving it disabled if:
+- You're certain the origin node will not come back online
+- You've permanently removed the origin node from the cluster
+- You want repaired rows to be treated as local writes

docs/http-api.md

@@ -153,11 +153,13 @@ Request body:
 | `generate_report` | bool | no | Write a JSON report. |
 | `fix_nulls` | bool | no | Fill NULLs using peer values. |
 | `bidirectional` | bool | no | Insert-only in both directions. |
+| `preserve_origin` | bool | no | Preserve replication origin node ID and LSN with per-row timestamp accuracy. Default: `false` |
 
 Notes:
 
 - Recovery-mode is not exposed via HTTP; origin-only diff files will be rejected.
 - The client certificate CN must map to a DB role that can run `SET ROLE` and perform required DML.
+- Preserve-origin maintains microsecond-precision timestamps for each repaired row, ensuring accurate temporal ordering in recovery scenarios.
 
 ### POST /api/v1/spock-diff
 

go.mod

@@ -22,7 +22,7 @@ require (
 	github.com/jackc/pglogrepl v0.0.0-20250509230407-a9884f6bd75a
 	github.com/jackc/pgtype v1.14.4
 	github.com/jackc/pgx/v5 v5.7.5
-	github.com/mattn/go-sqlite3 v1.14.22
+	github.com/mattn/go-sqlite3 v1.14.33
 	github.com/stretchr/testify v1.11.1
 	github.com/testcontainers/testcontainers-go v0.37.0
 	github.com/testcontainers/testcontainers-go/modules/compose v0.37.0

go.sum

@@ -420,8 +420,8 @@ github.com/mattn/go-runewidth v0.0.16/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh
 github.com/mattn/go-shellwords v1.0.12 h1:M2zGm7EW6UQJvDeQxo4T51eKPurbeFbe8WtebGE2xrk=
 github.com/mattn/go-shellwords v1.0.12/go.mod h1:EZzvwXDESEeg03EKmM+RmDnNOPKG4lLtQsUlTZDWQ8Y=
 github.com/mattn/go-sqlite3 v1.6.0/go.mod h1:FPy6KqzDD04eiIsT53CuJW3U88zkxoIYsOqkbpncsNc=
-github.com/mattn/go-sqlite3 v1.14.22 h1:2gZY6PC6kBnID23Tichd1K+Z0oS6nE/XwU+Vz/5o4kU=
-github.com/mattn/go-sqlite3 v1.14.22/go.mod h1:Uh1q+B4BYcTPb+yiD3kU8Ct7aC0hY9fxUwlHK0RXw+Y=
+github.com/mattn/go-sqlite3 v1.14.33 h1:A5blZ5ulQo2AtayQ9/limgHEkFreKj1Dv226a1K73s0=
+github.com/mattn/go-sqlite3 v1.14.33/go.mod h1:Uh1q+B4BYcTPb+yiD3kU8Ct7aC0hY9fxUwlHK0RXw+Y=
 github.com/matttproud/golang_protobuf_extensions v1.0.1/go.mod h1:D8He9yQNgCq6Z5Ld7szi9bcBfOoFv/3dc6xSMkL2PC0=
 github.com/mgutz/ansi v0.0.0-20170206155736-9520e82c474b h1:j7+1HpAFS1zy5+Q4qx1fWh90gTKwiN4QCGoY9TWyyO4=
 github.com/mgutz/ansi v0.0.0-20170206155736-9520e82c474b/go.mod h1:01TrycV0kFyexm33Z7vhZRXopbI8J3TDReVlkTgMUxE=

internal/api/http/handler.go

@@ -55,6 +55,7 @@ type tableRepairRequest struct {
 	GenerateReport bool     `json:"generate_report"`
 	FixNulls       bool     `json:"fix_nulls"`
 	Bidirectional  bool     `json:"bidirectional"`
+	PreserveOrigin bool     `json:"preserve_origin"`
 }
 
 type spockDiffRequest struct {
@@ -434,6 +435,7 @@ func (s *APIServer) handleTableRepair(w http.ResponseWriter, r *http.Request) {
 	task.GenerateReport = req.GenerateReport
 	task.FixNulls = req.FixNulls
 	task.Bidirectional = req.Bidirectional
+	task.PreserveOrigin = req.PreserveOrigin
 	task.Ctx = r.Context()
 	task.ClientRole = clientInfo.role
 	task.InvokeMethod = "api"

internal/cli/cli.go

@@ -223,6 +223,11 @@ func SetupCLI() *cli.App {
 			Usage: "Enable recovery-mode repair using origin-only diffs",
 			Value: false,
 		},
+		&cli.BoolFlag{
+			Name:    "preserve-origin",
+			Usage:   "Preserve replication origin node ID and commit timestamp for repaired rows",
+			Value:   false,
+		},
 		&cli.BoolFlag{
 			Name:    "fix-nulls",
 			Aliases: []string{"X"},
@@ -1199,6 +1204,7 @@ func TableRepairCLI(ctx *cli.Context) error {
 	task.Bidirectional = ctx.Bool("bidirectional")
 	task.GenerateReport = ctx.Bool("generate-report")
 	task.RecoveryMode = ctx.Bool("recovery-mode")
+	task.PreserveOrigin = ctx.Bool("preserve-origin")
 
 	if err := task.ValidateAndPrepare(); err != nil {
 		return fmt.Errorf("validation failed: %w", err)

internal/consistency/mtree/merkle.go

@@ -1364,6 +1364,11 @@ func (m *MerkleTreeTask) BuildMtree() (err error) {
 		}
 	}
 
+	// Ensure numBlocks matches actual block range count (covers both computed and RangesFile paths).
+	if len(blockRanges) > 0 {
+		numBlocks = len(blockRanges)
+	}
+
 	for _, nodeInfo := range m.ClusterNodes {
 		logger.Info("Processing node: %s", nodeInfo["Name"])
 		pool, ok := pools[nodeInfo["Name"].(string)]