diff --git a/go/internal/store/backfill_empty_tenant_pgtest_test.go b/go/internal/store/backfill_empty_tenant_pgtest_test.go new file mode 100644 index 000000000..3dd279ebe --- /dev/null +++ b/go/internal/store/backfill_empty_tenant_pgtest_test.go @@ -0,0 +1,97 @@ +//go:build pgtest + +package store + +import ( + "context" + "testing" +) + +// TestBackfillEmptyTenantMigration replays migration 0002 over rows a system-role +// wake stamped with an empty tenant_id. Each row must land in its agent's tenant; +// an empty-tenant binding whose tenant already binds that agent must be dropped. +func TestBackfillEmptyTenantMigration(t *testing.T) { + s := newTestStore(t) + ctx := context.Background() + tenantB := seedTenant(t, s, "tenant-b") + ctxB := WithTenant(ctx, tenantB) + lone := seedTenantAgent(t, ctxB, s, "lone") + twice := seedTenantAgent(t, ctxB, s, "twice") + + // The live state: twice is bound in tenant B and again under ''. + if _, err := s.RecordSessionBinding(ctxB, "sess-twice-live", twice.ID, "runner-1"); err != nil { + t.Fatalf("RecordSessionBinding: %v", err) + } + execAsSystem(t, s, "INSERT INTO agent_sessions (session_id, agent_account_id, tenant_id) VALUES ('sess-lone', $1, ''), ('sess-twice-old', $2, '')", + string(lone.ID), string(twice.ID)) + execAsSystem(t, s, "INSERT INTO session_bindings (agent_account_id, session_id, runner_id, tenant_id) VALUES ($1, 'sess-lone', 'runner-1', ''), ($2, 'sess-twice-old', 'runner-1', '')", + string(lone.ID), string(twice.ID)) + + execAsSystem(t, s, migrationSQL(t, 2)) + + if n := countAsSystem(t, s, "SELECT count(*) FROM agent_sessions WHERE tenant_id = ''") + + countAsSystem(t, s, "SELECT count(*) FROM session_bindings WHERE tenant_id = ''"); n != 0 { + t.Fatalf("%d rows still carry tenant_id ''", n) + } + if got, ok, err := s.LatestSessionForAccount(ctxB, lone.ID); err != nil || !ok || got != "sess-lone" { + t.Fatalf("tenant B session for lone = %q,%v,%v; want sess-lone", got, ok, err) + } + if got, err := s.SessionBindingTenant(WithSystemRole(ctx), "sess-lone", "runner-1"); err != nil || got != tenantB { + t.Fatalf("lone binding tenant = %q,%v; want %q", got, err, tenantB) + } + if got, err := s.SessionBindingTenant(WithSystemRole(ctx), "sess-twice-live", "runner-1"); err != nil || got != tenantB { + t.Fatalf("twice live binding tenant = %q,%v; want %q (no ambiguity)", got, err, tenantB) + } + if n := countAsSystem(t, s, "SELECT count(*) FROM session_bindings WHERE session_id = 'sess-twice-old'"); n != 0 { + t.Fatalf("stale '' binding for twice survived (%d rows)", n) + } +} + +func migrationSQL(t *testing.T, version int) string { + t.Helper() + migs, err := loadMigrations() + if err != nil { + t.Fatalf("loadMigrations: %v", err) + } + return migs[version-1].sql +} + +func execAsSystem(t *testing.T, s *Store, sql string, args ...any) { + t.Helper() + ctx := context.Background() + tx, err := s.pool.Begin(ctx) + if err != nil { + t.Fatalf("begin: %v", err) + } + defer func() { _ = tx.Rollback(ctx) }() + if _, err := tx.Exec(ctx, "SET LOCAL ROLE "+systemRole); err != nil { + t.Fatalf("set role: %v", err) + } + if _, err := tx.Exec(ctx, sql, args...); err != nil { + t.Fatalf("exec: %v", err) + } + if err := tx.Commit(ctx); err != nil { + t.Fatalf("commit: %v", err) + } +} + +func countAsSystem(t *testing.T, s *Store, sql string) int { + t.Helper() + var n int + execTx := func() error { + ctx := context.Background() + tx, err := s.pool.Begin(ctx) + if err != nil { + return err + } + defer func() { _ = tx.Rollback(ctx) }() + if _, err := tx.Exec(ctx, "SET LOCAL ROLE "+systemRole); err != nil { + return err + } + return tx.QueryRow(ctx, sql).Scan(&n) + } + if err := execTx(); err != nil { + t.Fatalf("count: %v", err) + } + return n +} diff --git a/go/internal/store/migrations/0002_backfill_empty_tenant.sql b/go/internal/store/migrations/0002_backfill_empty_tenant.sql new file mode 100644 index 000000000..8152253cf --- /dev/null +++ b/go/internal/store/migrations/0002_backfill_empty_tenant.sql @@ -0,0 +1,41 @@ +-- Repairs session rows that system-role wakes stamped with tenant_id '' before +-- WakeAgent ran under the agent's tenant. RLS hides '' rows from every tenant, +-- and a '' binding next to a real one makes SessionBindingTenant ambiguous. +-- Runs as compass_system: under FORCE RLS a non-superuser owner sees no rows. +SET LOCAL ROLE compass_system; + +UPDATE agent_sessions s + SET tenant_id = a.tenant_id + FROM agent_accounts a + WHERE s.tenant_id = '' AND a.account_id = s.agent_account_id; + +UPDATE agent_placements p + SET tenant_id = a.tenant_id + FROM agent_accounts a + WHERE p.tenant_id = '' AND a.account_id = p.agent_account_id; + +UPDATE agent_session_transcript_entries e + SET tenant_id = s.tenant_id + FROM agent_sessions s + WHERE e.tenant_id = '' AND s.session_id = e.session_id; + +UPDATE agent_session_archive_segments g + SET tenant_id = s.tenant_id + FROM agent_sessions s + WHERE g.tenant_id = '' AND s.session_id = g.session_id; + +-- A binding is live routing state. Re-home a '' binding only when its tenant +-- holds no binding for that account or session; drop it otherwise. +UPDATE session_bindings b + SET tenant_id = (SELECT a.tenant_id FROM agent_accounts a WHERE a.account_id = b.agent_account_id) + WHERE b.tenant_id = '' + AND NOT EXISTS ( + SELECT 1 + FROM session_bindings o + JOIN agent_accounts a ON a.account_id = b.agent_account_id + WHERE o.tenant_id = a.tenant_id + AND (o.agent_account_id = b.agent_account_id OR o.session_id = b.session_id)); + +DELETE FROM session_bindings WHERE tenant_id = ''; + +RESET ROLE;