From 715078b900627ddf1028d41b27e12ac6fa5e1756 Mon Sep 17 00:00:00 2001
From: Andy Yang <yang@cockroachlabs.com>
Date: Thu, 30 Mar 2023 05:07:18 +0000
Subject: [PATCH] upgrades: update system.web_sessions migration to handle
 orphaned rows

This patch updates the system.web_sessions user ID migration to delete
orphaned rows (i.e. rows corresponding to users that have been dropped)
after backfilling user IDs. They need to be deleted since they block
the addition of the NOT NULL constraint on the column.

Release note: None
---
 .../web_sessions_table_user_id_migration.go   | 37 ++++++++++++++-----
 ...b_sessions_table_user_id_migration_test.go |  6 +++
 2 files changed, 33 insertions(+), 10 deletions(-)

diff --git a/pkg/upgrade/upgrades/web_sessions_table_user_id_migration.go b/pkg/upgrade/upgrades/web_sessions_table_user_id_migration.go
index 8ec67bc38bb3..a4be8c7b0b94 100644
--- a/pkg/upgrade/upgrades/web_sessions_table_user_id_migration.go
+++ b/pkg/upgrade/upgrades/web_sessions_table_user_id_migration.go
@@ -53,6 +53,12 @@ WHERE w.user_id IS NULL AND w.username = u.username
 LIMIT 1000
 `
 
+const deleteRowsForDroppedUsersWebSessionsTableStmt = `
+DELETE FROM system.web_sessions
+WHERE user_id IS NULL
+LIMIT 1000
+`
+
 const setUserIDColumnToNotNullWebSessionsTableStmt = `
 ALTER TABLE system.web_sessions
 ALTER COLUMN user_id SET NOT NULL
@@ -62,16 +68,27 @@ func backfillWebSessionsTableUserIDColumn(
 	ctx context.Context, cs clusterversion.ClusterVersion, d upgrade.TenantDeps,
 ) error {
 	ie := d.DB.Executor()
-	for {
-		rowsAffected, err := ie.ExecEx(ctx, "backfill-user-id-col-web-sessions-table", nil, /* txn */
-			sessiondata.NodeUserSessionDataOverride,
-			backfillUserIDColumnWebSessionsTableStmt,
-		)
-		if err != nil {
-			return err
-		}
-		if rowsAffected == 0 {
-			break
+	for _, op := range []struct {
+		name  string
+		query string
+	}{
+		{
+			name:  "backfill-user-id-col-web-sessions-table",
+			query: backfillUserIDColumnWebSessionsTableStmt,
+		},
+		{
+			name:  "delete-rows-for-dropped-users-web-sessions-table",
+			query: deleteRowsForDroppedUsersWebSessionsTableStmt,
+		},
+	} {
+		for {
+			rowsAffected, err := ie.ExecEx(ctx, op.name, nil /* txn */, sessiondata.NodeUserSessionDataOverride, op.query)
+			if err != nil {
+				return err
+			}
+			if rowsAffected == 0 {
+				break
+			}
 		}
 	}
 
diff --git a/pkg/upgrade/upgrades/web_sessions_table_user_id_migration_test.go b/pkg/upgrade/upgrades/web_sessions_table_user_id_migration_test.go
index d877a0160137..4fc2fdaba096 100644
--- a/pkg/upgrade/upgrades/web_sessions_table_user_id_migration_test.go
+++ b/pkg/upgrade/upgrades/web_sessions_table_user_id_migration_test.go
@@ -101,6 +101,12 @@ VALUES (
 	})
 	tdb.CheckQueryResults(t, "SELECT count(*) FROM system.web_sessions", [][]string{{strconv.Itoa(numUsers)}})
 
+	// Drop a user to test that migration deletes orphaned rows.
+	if numUsers > 0 {
+		tdb.Exec(t, "DROP USER testuser0")
+		numUsers -= 1
+	}
+
 	// Run migrations.
 	_, err := tc.Conns[0].ExecContext(ctx, `SET CLUSTER SETTING version = $1`,
 		clusterversion.ByKey(clusterversion.V23_1WebSessionsTableHasUserIDColumn).String())