highperfocused/multica
1
0
Fork 0
mirror of https://github.com/multica-ai/multica.git synced 2026-07-26 20:45:37 +02:00
Code Issues Packages Projects Releases Wiki Activity

fix(migrate): auto-backfill attribution before migration 198 to unblock self-host upgrade (MUL-4897) (#5558)

Browse Source 
Self-hosted upgrades to v0.4.3 failed closed on migration 198's VALIDATE of the strict attribution constraint, because the legacy rows migration 190 exempted were only backfilled out-of-band on cloud. Registers a pre-198 preMigrationHook that idempotently mirrors originator_user_id into accountable_user_id in batches before VALIDATE, with FOR UPDATE + repeated predicate to avoid clobbering concurrently-written rows, so a stuck-at-197 instance auto-heals on migrate up with no manual SQL. Originator-NULL rows are left untouched. Verified with unit + concurrency + end-to-end tests against real Postgres.

Fixes #5544
This commit is contained in:
Bohan Jiang
2026-07-17 14:06:03 +08:00
committed by GitHub
parent ce15300493
commit 3ce25d16e7
3 changed files with 605 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
server/cmd/migrate/main.go
View File

@@ -9,6 +9,7 @@ import (
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgxpool"
"github.com/multica-ai/multica/server/internal/attributionbackfill"
"github.com/multica-ai/multica/server/internal/logger"
"github.com/multica-ai/multica/server/internal/migrations"
"github.com/multica-ai/multica/server/internal/taskusagebackfill"
@@ -37,8 +38,19 @@ type preMigrationHook func(ctx context.Context, pool *pgxpool.Pool) error
// can advance the watermark. The hook runs the same idempotent
// monthly-slice backfill that
// `cmd/backfill_task_usage_hourly` exposes to operators.
//
// MUL-4897 / GH #5544: migration 198 VALIDATEs the strict attribution
// constraint installed by 197, which drops migration 190's
// originator_source IS NULL exemption. Self-hosted databases never ran the
// out-of-band backfill that Multica's cloud did, so their legacy rows make
// 198 fail closed and the backend refuses to start. The hook reconciles
// those rows (accountable_user_id := originator_user_id) idempotently BEFORE
// VALIDATE, so a stuck-at-197 instance auto-heals on `migrate up` with no
// manual SQL. A higher-numbered migration cannot help — the instance never
// reaches a version above the failing 198.
var preMigrationHooks = map[string]preMigrationHook{
"103_drop_legacy_daily_rollups": runTaskUsageHourlyHook,
"103_drop_legacy_daily_rollups": runTaskUsageHourlyHook,
"198_agent_task_attribution_strict_constraint_validate": runAttributionStrictHook,
}
func runTaskUsageHourlyHook(ctx context.Context, pool *pgxpool.Pool) error {
@@ -60,6 +72,22 @@ func runTaskUsageHourlyHook(ctx context.Context, pool *pgxpool.Pool) error {
return nil
}
// runAttributionStrictHook backfills accountable_user_id from
// originator_user_id before migration 198 validates the strict attribution
// constraint, so self-hosted upgrades that never ran the out-of-band
// backfill recover automatically (GH #5544 / MUL-4897).
func runAttributionStrictHook(ctx context.Context, pool *pgxpool.Pool) error {
res, err := attributionbackfill.Hook(ctx, pool, attributionbackfill.HookOptions{})
if err != nil {
return fmt.Errorf("attribution strict-constraint pre-198 hook: %w", err)
}
slog.Info("attribution backfill hook: complete",
"rows_backfilled", res.RowsBackfilled,
"batches", res.Batches,
"mismatch_normalized", res.MismatchNormalized)
return nil
}
// migrationAdvisoryLockKey is the int64 identifier used with Postgres
// pg_advisory_lock to serialize the migration loop across concurrent
// runners (multi-replica backend Deployment, scale-up, or a manual

212
server/internal/attributionbackfill/backfill.go Normal file
View File

@@ -0,0 +1,212 @@
// Package attributionbackfill reconciles agent_task_queue rows to the
// Human Attribution strict invariant BEFORE migration 198 validates it.
//
// Background (GH #5544, MUL-4302, MUL-4897). Migration 197 installs the
// strict cross-column CHECK
//
// originator_user_id IS NULL
// OR (accountable_user_id IS NOT NULL AND accountable_user_id = originator_user_id)
//
// as agent_task_queue_accountable_matches_originator_strict (NOT VALID),
// and migration 198 runs VALIDATE CONSTRAINT on it. Migration 197's
// comment assumed the violating legacy rows had "been backfilled out of
// band" — which was true only on Multica's own cloud, where the backfill
// was run manually. Self-hosted deployments never got that step, so their
// legacy rows (originator_user_id set from before migration 185, or written
// by an older rolling-deploy backend, with accountable_user_id NULL) make
// migration 198 fail closed with SQLSTATE 23514 and the backend refuses to
// start.
//
// A new higher-numbered migration cannot fix this: a stuck instance sits
// with 197 applied and 198 failing, and never reaches a migration numbered
// above 198. The fix must run at-or-before 198. cmd/migrate already exposes
// a preMigrationHook mechanism (used for migration 103, MUL-2957) that runs
// idempotent work before a specific migration's SQL; this package is that
// hook for 198.
//
// Scope. The strict constraint can ONLY be violated by rows where
// originator_user_id IS NOT NULL, and for those the design invariant is
// unambiguous: a resolved originator IS the accountable human, so
// accountable_user_id MUST equal originator_user_id (MUL-4302 §1/§11).
// This hook therefore mirrors originator_user_id into accountable_user_id
// for exactly those rows and stamps originator_source='backfill' where it
// was NULL, which is precisely and only what is required to make VALIDATE
// pass with zero manual operator steps.
//
// Rows with originator_user_id IS NULL (autopilot rule_owner / owner_fallback,
// or fully unattributed history) already satisfy the strict constraint and
// are left untouched — enriching their accountable_user_id from the trigger
// chat/agent-owner waterfall is audit-completeness work that does not block
// startup and is deliberately kept out of this fail-closed path.
package attributionbackfill
import (
"context"
"fmt"
"log/slog"
"github.com/jackc/pgx/v5/pgxpool"
)
// DefaultBatchSize bounds how many rows each UPDATE statement rewrites, so
// the reconciliation never takes one long transaction / lock on a large
// agent_task_queue. Each batch autocommits on the pool; partial progress is
// simply resumed on the next batch (or the next migrate run), because a
// reconciled row no longer matches the selection predicate.
const DefaultBatchSize = 5000
// Result reports what a single Hook run did. Exposed so cmd/migrate and
// tests can log or assert on it.
type Result struct {
// RowsBackfilled is the total number of rows whose accountable_user_id
// was mirrored from originator_user_id across all batches.
RowsBackfilled int64
// Batches counts how many UPDATE batches ran (0 when nothing needed
// fixing).
Batches int
// MismatchNormalized counts rows that had BOTH originator_user_id and
// accountable_user_id set to DIFFERENT non-NULL users before the run.
// Per the design these should not exist (migration 190's transitional
// constraint has rejected them on every source-tagged write since it
// shipped); the count is surfaced as a warning so a genuinely
// mis-attributed row is visible even though the invariant forces
// accountable to follow originator.
MismatchNormalized int64
}
// HookOptions controls Hook behaviour. Zero values are correct for
// production; tests override fields as needed.
type HookOptions struct {
// Logger receives slog records about the walk. nil = slog.Default().
Logger *slog.Logger
// BatchSize overrides DefaultBatchSize. <= 0 = DefaultBatchSize.
BatchSize int
// Table overrides the target table name. Empty = "agent_task_queue".
// Tests point this at a throwaway fixture table; production leaves it
// empty. The value is interpolated into SQL, so callers must pass a
// trusted identifier (never user input) — production passes a
// constant and tests pass their own schema-qualified fixture name.
Table string
}
// countMismatchSQL counts rows where both attribution users are set but
// disagree — the "real mis-attribution" shape worth a warning.
const countMismatchSQL = `
SELECT count(*)
FROM %s
WHERE originator_user_id IS NOT NULL
AND accountable_user_id IS NOT NULL
AND accountable_user_id <> originator_user_id`
// backfillBatchSQL mirrors originator_user_id into accountable_user_id for a
// bounded batch of violating rows and stamps originator_source='backfill'
// only where it was NULL (COALESCE preserves an existing precise source such
// as 'direct_human'). Because originator_user_id IS NOT NULL for every
// selected row, `accountable_user_id IS DISTINCT FROM originator_user_id`
// captures both the NULL-accountable legacy rows and the both-set-but-differ
// rows. Reconciled rows stop matching the predicate, so the loop terminates
// and re-running is a no-op.
//
// Concurrency safety (defense in depth — the migrate loop normally runs with
// the server down, but the hook executes on the pool and a rolling deploy or
// an operator-run migrate could overlap a live writer). Two guards close the
// stale-selection race where a row is selected as violating but a concurrent
// writer flips it to a legitimate originator-NULL fork before it is rewritten:
//
// - FOR UPDATE in the CTE locks each candidate row and, under READ
// COMMITTED, re-evaluates the predicate against the latest committed
// version once the lock is granted — a row that became originator-NULL is
// dropped from the batch instead of being blindly overwritten.
// - The outer UPDATE repeats the same predicate on q, so it never writes a
// row that no longer needs reconciliation even if it changed underneath
// the selection.
//
// Without these, `SET accountable_user_id = q.originator_user_id` on a row a
// writer just turned into (originator NULL, accountable set) would clobber the
// writer's accountable value with NULL.
const backfillBatchSQL = `
WITH batch AS (
SELECT id
FROM %s
WHERE originator_user_id IS NOT NULL
AND accountable_user_id IS DISTINCT FROM originator_user_id
LIMIT $1
FOR UPDATE
)
UPDATE %s q
SET accountable_user_id = q.originator_user_id,
originator_source = COALESCE(q.originator_source, 'backfill')
FROM batch
WHERE q.id = batch.id
AND q.originator_user_id IS NOT NULL
AND q.accountable_user_id IS DISTINCT FROM q.originator_user_id`
// Hook is the migration-time entrypoint invoked before migration 198's
// VALIDATE. It is idempotent and safe to retry: it reconciles every row
// that would violate agent_task_queue_accountable_matches_originator_strict
// and returns nil so the migration loop proceeds to VALIDATE.
//
// It returns an error only when a SQL statement itself fails (e.g. the
// database is unreachable mid-walk); that aborts the migrate run before
// schema_migrations records 198, so the same version — hook included —
// retries cleanly on the next invocation.
//
// The hook runs on the pool (autocommit per batch), NOT on the migrate
// loop's pinned lock connection, matching the taskusagebackfill hook
// contract. The loop already holds migrationAdvisoryLockKey, and migrate
// runs with the server down, so no concurrent writer competes for these
// rows.
func Hook(ctx context.Context, pool *pgxpool.Pool, opts HookOptions) (Result, error) {
log := opts.Logger
if log == nil {
log = slog.Default()
}
batchSize := opts.BatchSize
if batchSize <= 0 {
batchSize = DefaultBatchSize
}
table := opts.Table
if table == "" {
table = "agent_task_queue"
}
var res Result
// Surface any genuinely mis-attributed rows (both users set but
// different) up front. The invariant still forces accountable to follow
// originator below, but the operator/audit trail should show the count.
if err := pool.QueryRow(ctx, fmt.Sprintf(countMismatchSQL, table)).Scan(&res.MismatchNormalized); err != nil {
return res, fmt.Errorf("count attribution mismatches: %w", err)
}
if res.MismatchNormalized > 0 {
log.Warn("attribution backfill: normalizing rows where accountable_user_id disagreed with a non-NULL originator_user_id; originator is authoritative but these are worth auditing",
"mismatch_rows", res.MismatchNormalized)
}
updateSQL := fmt.Sprintf(backfillBatchSQL, table, table)
for {
tag, err := pool.Exec(ctx, updateSQL, batchSize)
if err != nil {
return res, fmt.Errorf("backfill accountable_user_id batch: %w", err)
}
n := tag.RowsAffected()
if n == 0 {
break
}
res.RowsBackfilled += n
res.Batches++
log.Info("attribution backfill: batch reconciled",
"rows", n,
"total", res.RowsBackfilled)
}
if res.RowsBackfilled == 0 {
log.Info("attribution backfill: no rows needed reconciliation before migration 198")
} else {
log.Info("attribution backfill: complete",
"rows_backfilled", res.RowsBackfilled,
"batches", res.Batches,
"mismatch_normalized", res.MismatchNormalized)
}
return res, nil
}

364
server/internal/attributionbackfill/backfill_test.go Normal file
View File

@@ -0,0 +1,364 @@
package attributionbackfill
import (
"context"
"fmt"
"math/rand/v2"
"os"
"testing"
"time"
"github.com/jackc/pgx/v5/pgxpool"
)
// These tests exercise the migration-198 backfill hook against a live
// Postgres. They connect to DATABASE_URL (default
// postgres://multica:multica@localhost:5432/multica?sslmode=disable),
// matching every other live-Postgres suite in the repo, and skip cleanly
// when no database is reachable so CI without a DB sees SKIP, not failure.
//
// Each test isolates itself in a throwaway schema so it never touches the
// real agent_task_queue, and drops the schema on cleanup.
const (
// The exact strict constraint installed by migration 197, applied to
// the fixture table so we can prove VALIDATE passes only after the hook
// reconciles the violating rows.
strictCheck = `originator_user_id IS NULL
OR (accountable_user_id IS NOT NULL AND accountable_user_id = originator_user_id)`
)
func newTestPool(t *testing.T) *pgxpool.Pool {
t.Helper()
dbURL := os.Getenv("DATABASE_URL")
if dbURL == "" {
dbURL = "postgres://multica:multica@localhost:5432/multica?sslmode=disable"
}
ctx := context.Background()
pool, err := pgxpool.New(ctx, dbURL)
if err != nil {
t.Skipf("could not connect to %s: %v", dbURL, err)
}
if err := pool.Ping(ctx); err != nil {
pool.Close()
t.Skipf("database not reachable at %s: %v", dbURL, err)
}
return pool
}
// newFixture creates a throwaway schema with a minimal agent_task_queue-like
// table carrying only the attribution columns. It returns the
// schema-qualified, quoted table identifier for HookOptions.Table.
//
// The strict constraint is deliberately NOT added here: callers seed their
// (possibly violating) rows first, then call addStrictConstraint, exactly
// mirroring the real upgrade order where legacy rows predate migration 197.
// A NOT VALID CHECK still rejects new INSERTs, so a violating row can only
// exist if it was written before the constraint was added.
func newFixture(t *testing.T, pool *pgxpool.Pool) string {
t.Helper()
ctx := context.Background()
schema := fmt.Sprintf("attr_backfill_test_%d_%d", time.Now().UnixNano(), rand.IntN(1_000_000))
if _, err := pool.Exec(ctx, fmt.Sprintf(`CREATE SCHEMA %q`, schema)); err != nil {
t.Fatalf("create schema: %v", err)
}
t.Cleanup(func() {
_, _ = pool.Exec(context.Background(), fmt.Sprintf(`DROP SCHEMA %q CASCADE`, schema))
})
table := fmt.Sprintf("%q.agent_task_queue", schema)
if _, err := pool.Exec(ctx, fmt.Sprintf(`
CREATE TABLE %s (
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
originator_user_id UUID NULL,
accountable_user_id UUID NULL,
originator_source TEXT NULL
)`, table)); err != nil {
t.Fatalf("create fixture table: %v", err)
}
return table
}
// addStrictConstraint installs migration 197's strict CHECK as NOT VALID,
// which does not scan existing rows — matching the real migration.
func addStrictConstraint(t *testing.T, pool *pgxpool.Pool, table string) {
t.Helper()
if _, err := pool.Exec(context.Background(), fmt.Sprintf(`
ALTER TABLE %s
ADD CONSTRAINT agent_task_queue_accountable_matches_originator_strict
CHECK (%s) NOT VALID`, table, strictCheck)); err != nil {
t.Fatalf("add strict constraint: %v", err)
}
}
func TestHook_ReconcilesViolatingRowsSoValidatePasses(t *testing.T) {
pool := newTestPool(t)
defer pool.Close()
ctx := context.Background()
table := newFixture(t, pool)
const (
u1 = "11111111-1111-1111-1111-111111111111" // legacy: originator set, accountable NULL, source NULL
u2 = "22222222-2222-2222-2222-222222222222" // rolling-deploy: same shape
u3 = "33333333-3333-3333-3333-333333333333" // already-good: accountable == originator, precise source
u4 = "44444444-4444-4444-4444-444444444444" // mismatch originator
u5 = "55555555-5555-5555-5555-555555555555" // mismatch accountable
u6 = "66666666-6666-6666-6666-666666666666" // legit divergent: originator NULL, accountable set
)
// id, originator, accountable, source
rows := [][4]any{
{"aaaaaaaa-0000-0000-0000-000000000001", u1, nil, nil}, // -> backfill to u1, source 'backfill'
{"aaaaaaaa-0000-0000-0000-000000000002", u2, nil, nil}, // -> backfill to u2, source 'backfill'
{"aaaaaaaa-0000-0000-0000-000000000003", u3, u3, "direct_human"}, // untouched
{"aaaaaaaa-0000-0000-0000-000000000004", u4, u5, "direct_human"}, // mismatch -> accountable u4, source kept
{"aaaaaaaa-0000-0000-0000-000000000005", nil, u6, "rule_owner"}, // legit divergent, untouched
{"aaaaaaaa-0000-0000-0000-000000000006", nil, nil, "unattributed"}, // orphan, untouched
}
for _, r := range rows {
if _, err := pool.Exec(ctx,
fmt.Sprintf(`INSERT INTO %s (id, originator_user_id, accountable_user_id, originator_source) VALUES ($1,$2,$3,$4)`, table),
r[0], r[1], r[2], r[3]); err != nil {
t.Fatalf("insert fixture row %v: %v", r[0], err)
}
}
addStrictConstraint(t, pool, table)
res, err := Hook(ctx, pool, HookOptions{Table: table})
if err != nil {
t.Fatalf("Hook: %v", err)
}
// u1, u2 (NULL accountable) + the mismatch row = 3 reconciled.
if res.RowsBackfilled != 3 {
t.Errorf("RowsBackfilled = %d, want 3", res.RowsBackfilled)
}
if res.MismatchNormalized != 1 {
t.Errorf("MismatchNormalized = %d, want 1", res.MismatchNormalized)
}
// The whole point: VALIDATE now succeeds.
if _, err := pool.Exec(ctx, fmt.Sprintf(
`ALTER TABLE %s VALIDATE CONSTRAINT agent_task_queue_accountable_matches_originator_strict`, table)); err != nil {
t.Fatalf("VALIDATE after hook should pass, got: %v", err)
}
assertRow := func(id, wantAccountable, wantSource string) {
t.Helper()
var acc, src *string
if err := pool.QueryRow(ctx, fmt.Sprintf(
`SELECT accountable_user_id::text, originator_source FROM %s WHERE id=$1`, table), id).Scan(&acc, &src); err != nil {
t.Fatalf("read row %s: %v", id, err)
}
gotAcc, gotSrc := "<nil>", "<nil>"
if acc != nil {
gotAcc = *acc
}
if src != nil {
gotSrc = *src
}
if gotAcc != wantAccountable || gotSrc != wantSource {
t.Errorf("row %s = (accountable %s, source %s), want (%s, %s)", id, gotAcc, gotSrc, wantAccountable, wantSource)
}
}
assertRow("aaaaaaaa-0000-0000-0000-000000000001", u1, "backfill") // legacy filled, source stamped
assertRow("aaaaaaaa-0000-0000-0000-000000000002", u2, "backfill") // rolling filled, source stamped
assertRow("aaaaaaaa-0000-0000-0000-000000000003", u3, "direct_human") // untouched
assertRow("aaaaaaaa-0000-0000-0000-000000000004", u4, "direct_human") // normalized to originator, source preserved
assertRow("aaaaaaaa-0000-0000-0000-000000000005", u6, "rule_owner") // legit divergent untouched
assertRow("aaaaaaaa-0000-0000-0000-000000000006", "<nil>", "unattributed")
// Idempotent: a second run reconciles nothing.
res2, err := Hook(ctx, pool, HookOptions{Table: table})
if err != nil {
t.Fatalf("second Hook: %v", err)
}
if res2.RowsBackfilled != 0 || res2.Batches != 0 {
t.Errorf("second run not a no-op: %+v", res2)
}
}
func TestHook_EmptyTableIsNoOp(t *testing.T) {
pool := newTestPool(t)
defer pool.Close()
ctx := context.Background()
table := newFixture(t, pool)
addStrictConstraint(t, pool, table)
res, err := Hook(ctx, pool, HookOptions{Table: table})
if err != nil {
t.Fatalf("Hook on empty table: %v", err)
}
if res.RowsBackfilled != 0 || res.MismatchNormalized != 0 || res.Batches != 0 {
t.Errorf("empty-table run should be a no-op, got %+v", res)
}
if _, err := pool.Exec(ctx, fmt.Sprintf(
`ALTER TABLE %s VALIDATE CONSTRAINT agent_task_queue_accountable_matches_originator_strict`, table)); err != nil {
t.Fatalf("VALIDATE on empty table should pass, got: %v", err)
}
}
// TestHook_ConcurrentForkNotClobbered reproduces the stale-selection race: a
// row is selected as violating, but a concurrent writer flips it to a
// legitimate originator-NULL fork (originator NULL, accountable set) while the
// hook is mid-flight. The FOR UPDATE lock + repeated outer predicate must make
// the hook skip that row rather than overwrite the writer's accountable value
// with NULL, while still reconciling the untouched violating rows.
func TestHook_ConcurrentForkNotClobbered(t *testing.T) {
pool := newTestPool(t)
defer pool.Close()
ctx := context.Background()
table := newFixture(t, pool)
const (
u1 = "11111111-1111-1111-1111-111111111111" // R: legacy originator, gets flipped by the writer
u2 = "22222222-2222-2222-2222-222222222222" // fork accountable the writer sets
u3 = "33333333-3333-3333-3333-333333333333" // R2: control violating row, must be backfilled
)
rID := "aaaaaaaa-0000-0000-0000-0000000000f1"
r2ID := "aaaaaaaa-0000-0000-0000-0000000000f2"
for _, r := range [][2]any{{rID, u1}, {r2ID, u3}} {
if _, err := pool.Exec(ctx, fmt.Sprintf(
`INSERT INTO %s (id, originator_user_id, accountable_user_id, originator_source) VALUES ($1,$2,NULL,NULL)`, table),
r[0], r[1]); err != nil {
t.Fatalf("seed row %v: %v", r[0], err)
}
}
addStrictConstraint(t, pool, table)
// Writer B: open a transaction that flips R to a legit originator-NULL fork
// and holds the row lock without committing yet.
writer, err := pool.Acquire(ctx)
if err != nil {
t.Fatalf("acquire writer conn: %v", err)
}
defer writer.Release()
tx, err := writer.Begin(ctx)
if err != nil {
t.Fatalf("begin writer tx: %v", err)
}
if _, err := tx.Exec(ctx, fmt.Sprintf(
`UPDATE %s SET originator_user_id = NULL, accountable_user_id = $2, originator_source = 'rule_owner' WHERE id = $1`, table),
rID, u2); err != nil {
t.Fatalf("writer flip: %v", err)
}
// Run the hook concurrently; it will block trying to FOR UPDATE-lock R.
type hookResult struct {
res Result
err error
}
done := make(chan hookResult, 1)
go func() {
res, err := Hook(ctx, pool, HookOptions{Table: table})
done <- hookResult{res, err}
}()
// Wait until the hook is actually blocked on a lock before committing the
// writer — this deterministically forces the "row changed under the
// selection" ordering rather than relying on a sleep.
waitUntilBlocked(t, pool, table)
if err := tx.Commit(ctx); err != nil {
t.Fatalf("commit writer: %v", err)
}
hr := <-done
if hr.err != nil {
t.Fatalf("hook: %v", hr.err)
}
// Only the control row R2 should have been reconciled; R was dropped.
if hr.res.RowsBackfilled != 1 {
t.Errorf("RowsBackfilled = %d, want 1 (R2 only; R must be skipped)", hr.res.RowsBackfilled)
}
// R keeps the writer's legit fork — NOT clobbered to (NULL, NULL).
var origR, accR, srcR *string
if err := pool.QueryRow(ctx, fmt.Sprintf(
`SELECT originator_user_id::text, accountable_user_id::text, originator_source FROM %s WHERE id=$1`, table), rID).
Scan(&origR, &accR, &srcR); err != nil {
t.Fatalf("read R: %v", err)
}
if origR != nil || accR == nil || *accR != u2 || srcR == nil || *srcR != "rule_owner" {
t.Errorf("R was clobbered: originator=%v accountable=%v source=%v; want (NULL, %s, rule_owner)", deref(origR), deref(accR), deref(srcR), u2)
}
// R2 was reconciled.
var accR2, srcR2 *string
if err := pool.QueryRow(ctx, fmt.Sprintf(
`SELECT accountable_user_id::text, originator_source FROM %s WHERE id=$1`, table), r2ID).Scan(&accR2, &srcR2); err != nil {
t.Fatalf("read R2: %v", err)
}
if accR2 == nil || *accR2 != u3 || srcR2 == nil || *srcR2 != "backfill" {
t.Errorf("R2 not reconciled: accountable=%v source=%v; want (%s, backfill)", deref(accR2), deref(srcR2), u3)
}
// And the constraint validates.
if _, err := pool.Exec(ctx, fmt.Sprintf(
`ALTER TABLE %s VALIDATE CONSTRAINT agent_task_queue_accountable_matches_originator_strict`, table)); err != nil {
t.Fatalf("VALIDATE after concurrent run should pass, got: %v", err)
}
}
func deref(s *string) string {
if s == nil {
return "<nil>"
}
return *s
}
// waitUntilBlocked polls pg_stat_activity until at least one backend is blocked
// on a lock while running a statement against the fixture schema, so the test
// can commit the writer at exactly the racy moment. It fails if the hook never
// blocks within the timeout, which would itself signal the FOR UPDATE guard is
// missing.
func waitUntilBlocked(t *testing.T, pool *pgxpool.Pool, table string) {
t.Helper()
ctx := context.Background()
deadline := time.Now().Add(10 * time.Second)
for time.Now().Before(deadline) {
var n int
if err := pool.QueryRow(ctx, `
SELECT count(*) FROM pg_stat_activity
WHERE wait_event_type = 'Lock'
AND state = 'active'
AND query ILIKE '%' || $1 || '%'`, table).Scan(&n); err != nil {
t.Fatalf("poll pg_stat_activity: %v", err)
}
if n > 0 {
return
}
time.Sleep(25 * time.Millisecond)
}
t.Fatal("hook never blocked on the row lock; FOR UPDATE guard likely missing")
}
func TestHook_BatchingReconcilesAll(t *testing.T) {
pool := newTestPool(t)
defer pool.Close()
ctx := context.Background()
table := newFixture(t, pool)
// Seed more violating rows than one batch so the loop must iterate.
const n = 25
for i := 0; i < n; i++ {
if _, err := pool.Exec(ctx, fmt.Sprintf(
`INSERT INTO %s (originator_user_id, accountable_user_id, originator_source) VALUES (gen_random_uuid(), NULL, NULL)`, table)); err != nil {
t.Fatalf("seed row %d: %v", i, err)
}
}
addStrictConstraint(t, pool, table)
res, err := Hook(ctx, pool, HookOptions{Table: table, BatchSize: 10})
if err != nil {
t.Fatalf("Hook: %v", err)
}
if res.RowsBackfilled != n {
t.Errorf("RowsBackfilled = %d, want %d", res.RowsBackfilled, n)
}
if res.Batches != 3 { // 10 + 10 + 5
t.Errorf("Batches = %d, want 3", res.Batches)
}
if _, err := pool.Exec(ctx, fmt.Sprintf(
`ALTER TABLE %s VALIDATE CONSTRAINT agent_task_queue_accountable_matches_originator_strict`, table)); err != nil {
t.Fatalf("VALIDATE after batched backfill should pass, got: %v", err)
}
}