fix(store): backfill session rows stamped with an empty tenant (RIG-4167)

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
+}

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;