sql: allow creating or dropping roles while membership cache is being…

… used

When we set `allow_role_memberships_to_change_during_transaction` we should be
able to create and drop users even when there are contending transactions on the
`system.users` and `system.role_options`. These code changes release the locks on
those system tables when `allow_role_memberships_to_change_during_transaction` is set.

Fixes: #137710
Release note (bug fix): When we set `allow_role_memberships_to_change_during_transaction`
are able to create and drop users quickly even when there are contending transactions
on the `system.users` and `system.role_options`.

pkg/sql/rolemembershipcache/cache.go

@@ -117,27 +117,58 @@ func (m *MembershipCache) GetRolesForMember(
 		return nil, errors.AssertionFailedf("cannot use MembershipCache without a txn")
 	}
 
-	// Lookup table version.
-	tableDesc, err := txn.Descriptors().ByIDWithLeased(txn.KV()).Get().Table(ctx, keys.RoleMembersTableID)
+	// Lookup table versions.
+	roleMembersTableDesc, err := txn.Descriptors().ByIDWithLeased(txn.KV()).Get().Table(ctx, keys.RoleMembersTableID)
 	if err != nil {
 		return nil, err
 	}
 
-	tableVersion := tableDesc.GetVersion()
-	if tableDesc.IsUncommittedVersion() {
+	tableVersion := roleMembersTableDesc.GetVersion()
+	if roleMembersTableDesc.IsUncommittedVersion() {
 		return resolveRolesForMember(ctx, txn, member)
 	}
+
 	if txn.SessionData().AllowRoleMembershipsToChangeDuringTransaction {
+		systemUsersTableDesc, err := txn.Descriptors().ByIDWithLeased(txn.KV()).Get().Table(ctx, keys.UsersTableID)
+		if err != nil {
+			return nil, err
+		}
+
+		roleOptionsTableDesc, err := txn.Descriptors().ByIDWithLeased(txn.KV()).Get().Table(ctx, keys.RoleOptionsTableID)
+		if err != nil {
+			return nil, err
+		}
+
+		systemUsersTableVersion := systemUsersTableDesc.GetVersion()
+		if systemUsersTableDesc.IsUncommittedVersion() {
+			return resolveRolesForMember(ctx, txn, member)
+		}
+
+		roleOptionsTableVersion := roleOptionsTableDesc.GetVersion()
+		if roleOptionsTableDesc.IsUncommittedVersion() {
+			return resolveRolesForMember(ctx, txn, member)
+		}
+
 		defer func() {
 			if retErr != nil {
 				return
 			}
 			txn.Descriptors().ReleaseSpecifiedLeases(ctx, []lease.IDVersion{
 				{
-					Name:    tableDesc.GetName(),
-					ID:      tableDesc.GetID(),
+					Name:    roleMembersTableDesc.GetName(),
+					ID:      roleMembersTableDesc.GetID(),
 					Version: tableVersion,
 				},
+				{
+					Name:    systemUsersTableDesc.GetName(),
+					ID:      systemUsersTableDesc.GetID(),
+					Version: systemUsersTableVersion,
+				},
+				{
+					Name:    roleOptionsTableDesc.GetName(),
+					ID:      roleOptionsTableDesc.GetID(),
+					Version: roleOptionsTableVersion,
+				},
 			})
 		}()
 	}

pkg/sql/tests/allow_role_memberships_to_change_during_transaction_test.go

@@ -82,11 +82,12 @@ func TestAllowRoleMembershipsToChangeDuringTransaction(t *testing.T) {
 		require.NoError(t, fooTx.Commit())
 		require.NoError(t, <-errCh)
 	})
+
 	// In this test we ensure that we can perform role grant and revoke
 	// operations while the transaction which uses the relevant roles
 	// remains open. We ensure that the transaction still succeeds and
 	// that the operations occur in a timely manner.
-	t.Run("session variable prevents waiting", func(t *testing.T) {
+	t.Run("session variable prevents waiting during GRANT and REVOKE", func(t *testing.T) {
 		fooConn, err := fooDB.Conn(ctx)
 		require.NoError(t, err)
 		defer func() { _ = fooConn.Close() }()
@@ -116,4 +117,36 @@ func TestAllowRoleMembershipsToChangeDuringTransaction(t *testing.T) {
 		// Ensure the transaction we held open commits without issue.
 		require.NoError(t, fooTx.Commit())
 	})
+
+	t.Run("session variable prevents waiting during CREATE and DROP role", func(t *testing.T) {
+		fooConn, err := fooDB.Conn(ctx)
+		require.NoError(t, err)
+		defer func() { _ = fooConn.Close() }()
+		_, err = fooConn.ExecContext(ctx, "SET allow_role_memberships_to_change_during_transaction = true;")
+		require.NoError(t, err)
+		fooTx, err := fooConn.BeginTx(ctx, nil)
+		require.NoError(t, err)
+		// We need to use show roles because that access the system.users table.
+		_, err = fooTx.Exec("SHOW ROLES")
+		require.NoError(t, err)
+
+		conn, err := sqlDB.Conn(ctx)
+		require.NoError(t, err)
+		defer func() { _ = conn.Close() }()
+		// Set a timeout on the SQL operations to ensure that they both
+		// happen in a timely manner.
+		grantRevokeTimeout, cancel := context.WithTimeout(
+			ctx, testutils.DefaultSucceedsSoonDuration,
+		)
+		defer cancel()
+
+		_, err = conn.ExecContext(grantRevokeTimeout, "CREATE ROLE new_role;")
+		require.NoError(t, err)
+		_, err = conn.ExecContext(grantRevokeTimeout, "DROP ROLE new_role;")
+		require.NoError(t, err)
+
+		// Ensure the transaction we held open commits without issue.
+		require.NoError(t, fooTx.Commit())
+	})
+
 }