hc(ref): Ensure OrganizationMember write paths have outboxes. (#49647)

Based on https://github.com/getsentry/sentry/pull/49386 with minor
adjustments. Thanks @dashed !

Gets all organization member write paths completely covered, validated
with the conftest blacklist.

---------

Co-authored-by: Leander Rodrigues <leander.rodrigues@sentry.io>

src/sentry/api/endpoints/organization_member/details.py

@@ -295,9 +295,8 @@ class OrganizationMemberDetailsEndpoint(OrganizationMemberEndpoint):
                 organizationmember=member, role__in=lesser_team_roles
             ).update(role=None)
 
-            member.update(role=role)
-            region_outbox = member.save_outbox_for_update()
-        if region_outbox:
+            member.role = role
+            region_outbox = member.save()
             region_outbox.drain_shard(max_updates_to_drain=10)
         if omt_update_count > 0:
             metrics.incr(

src/sentry/models/organization.py

@@ -406,7 +406,9 @@ class Organization(Model, SnowflakeIdMixin):
                     organization=to_org, user_id=from_member.user.id
                 )
             except OrganizationMember.DoesNotExist:
-                from_member.update(organization=to_org)
+                with transaction.atomic():
+                    from_member.organization = to_org
+                    from_member.save()
                 to_member = from_member
             else:
                 qs = OrganizationMemberTeam.objects.filter(

src/sentry/models/organizationmember.py

@@ -238,22 +238,23 @@ class OrganizationMember(Model):
             self.save_outbox_for_update()
             return super().delete(*args, **kwds)
 
-    @transaction.atomic
     def save(self, *args, **kwargs):
         assert (self.user_id is None and self.email) or (
             self.user_id and self.email is None
         ), "Must set either user or email"
-        if self.token and not self.token_expires_at:
-            self.refresh_expires_at()
-        is_new = not bool(self.id)
-        super().save(*args, **kwargs)
-        region_outbox = None
-        if is_new:
-            region_outbox = self.save_outbox_for_create()
-        else:
-            region_outbox = self.save_outbox_for_update()
-        self.__org_roles_from_teams = None
-        return region_outbox
+
+        with transaction.atomic(), in_test_psql_role_override("postgres"):
+            if self.token and not self.token_expires_at:
+                self.refresh_expires_at()
+            is_new = not bool(self.id)
+            super().save(*args, **kwargs)
+            region_outbox = None
+            if is_new:
+                region_outbox = self.save_outbox_for_create()
+            else:
+                region_outbox = self.save_outbox_for_update()
+            self.__org_roles_from_teams = None
+            return region_outbox
 
     def refresh_from_db(self, *args, **kwargs):
         super().refresh_from_db(*args, **kwargs)

src/sentry/models/user.py

@@ -27,7 +27,9 @@ from sentry.models.authenticator import Authenticator
 from sentry.models.avatars import UserAvatar
 from sentry.models.lostpasswordhash import LostPasswordHash
 from sentry.models.outbox import ControlOutbox, OutboxCategory, OutboxScope
+from sentry.services.hybrid_cloud.organization import organization_service
 from sentry.services.hybrid_cloud.user import RpcUser
+from sentry.silo import SiloMode
 from sentry.types.integrations import EXTERNAL_PROVIDERS, ExternalProviders
 from sentry.types.region import find_regions_for_user
 from sentry.utils.http import absolute_uri
@@ -291,20 +293,13 @@ class User(BaseModel, AbstractBaseUser):
 
     def merge_to(from_user, to_user):
         # TODO: we could discover relations automatically and make this useful
-        from sentry import roles
         from sentry.models import (
-            Activity,
             AuditLogEntry,
             Authenticator,
             AuthIdentity,
-            GroupAssignee,
-            GroupBookmark,
-            GroupSeen,
-            GroupShare,
-            GroupSubscription,
             Identity,
             OrganizationMember,
-            OrganizationMemberTeam,
+            OrganizationMemberMapping,
             UserAvatar,
             UserEmail,
             UserOption,
@@ -314,33 +309,20 @@ class User(BaseModel, AbstractBaseUser):
             "user.merge", extra={"from_user_id": from_user.id, "to_user_id": to_user.id}
         )
 
-        for obj in OrganizationMember.objects.filter(user_id=from_user.id):
-            try:
-                with transaction.atomic():
-                    obj.update(user_id=to_user.id)
-            # this will error if both users are members of obj.org
-            except IntegrityError:
-                pass
-
-            # identify the highest priority membership
-            # only applies if both users are members of obj.org
-            # if roles are different, grants combined user the higher of the two
-            to_member = OrganizationMember.objects.get(
-                organization=obj.organization_id, user_id=to_user.id
+        organization_ids: List[int]
+        if SiloMode.get_current_mode() == SiloMode.MONOLITH:
+            organization_ids = OrganizationMember.objects.filter(user_id=from_user.id).values_list(
+                "organization_id", flat=True
+            )
+        else:
+            organization_ids = OrganizationMemberMapping.objects.filter(
+                user_id=from_user.id
+            ).values_list("organization_id", flat=True)
+
+        for organization_id in organization_ids:
+            organization_service.merge_users(
+                organization_id=organization_id, from_user_id=from_user.id, to_user_id=to_user.id
             )
-            if roles.get(obj.role).priority > roles.get(to_member.role).priority:
-                to_member.update(role=obj.role)
-
-            for team in obj.teams.all():
-                try:
-                    with transaction.atomic():
-                        OrganizationMemberTeam.objects.create(
-                            organizationmember=to_member, team=team
-                        )
-                # this will error if both users are on the same team in obj.org,
-                # in which case, no need to update anything
-                except IntegrityError:
-                    pass
 
         model_list = (
             Authenticator,
@@ -348,11 +330,6 @@ class User(BaseModel, AbstractBaseUser):
             UserAvatar,
             UserEmail,
             UserOption,
-            GroupAssignee,
-            GroupBookmark,
-            GroupSeen,
-            GroupShare,
-            GroupSubscription,
         )
 
         for model in model_list:
@@ -363,7 +340,6 @@ class User(BaseModel, AbstractBaseUser):
                 except IntegrityError:
                     pass
 
-        Activity.objects.filter(user_id=from_user.id).update(user_id=to_user.id)
         # users can be either the subject or the object of actions which get logged
         AuditLogEntry.objects.filter(actor=from_user).update(actor=to_user)
         AuditLogEntry.objects.filter(target_user=from_user).update(target_user=to_user)

src/sentry/services/hybrid_cloud/organization/impl.py

@@ -2,11 +2,17 @@ from __future__ import annotations
 
 from typing import Iterable, List, Optional, Set, cast
 
-from django.db import models, transaction
+from django.db import IntegrityError, models, transaction
 
 from sentry import roles
 from sentry.db.postgres.roles import in_test_psql_role_override
 from sentry.models import (
+    Activity,
+    GroupAssignee,
+    GroupBookmark,
+    GroupSeen,
+    GroupShare,
+    GroupSubscription,
     Organization,
     OrganizationMember,
     OrganizationMemberTeam,
@@ -378,12 +384,67 @@ class DatabaseBackedOrganizationService(OrganizationService):
             region_outbox.drain_shard(max_updates_to_drain=10)
         return serialize_member(org_member)
 
-    def reset_idp_flags(self, *, organization_id: int) -> None:
-        OrganizationMember.objects.filter(
-            organization_id=organization_id,
-            flags=models.F("flags").bitor(OrganizationMember.flags["idp:provisioned"]),
-        ).update(
-            flags=models.F("flags")
-            .bitand(~OrganizationMember.flags["idp:provisioned"])
-            .bitand(~OrganizationMember.flags["idp:role-restricted"])
+    def merge_users(self, *, organization_id: int, from_user_id: int, to_user_id: int) -> None:
+        to_member: Optional[OrganizationMember] = OrganizationMember.objects.filter(
+            organization_id=organization_id, user_id=to_user_id
+        ).first()
+
+        from_member: Optional[OrganizationMember] = OrganizationMember.objects.filter(
+            organization_id=organization_id, user_id=from_user_id
+        ).first()
+
+        if from_member is None:
+            return
+
+        if to_member is None:
+            to_member = OrganizationMember.objects.create(
+                organization_id=organization_id,
+                user_id=to_user_id,
+                role=from_member.role,
+                flags=from_member.flags,
+                has_global_access=from_member.has_global_access,
+            )
+        else:
+            if roles.get(from_member.role).priority > roles.get(to_member.role).priority:
+                to_member.role = from_member.role
+            to_member.save()
+
+        assert to_member
+
+        for team in from_member.teams.all():
+            try:
+                with transaction.atomic():
+                    OrganizationMemberTeam.objects.create(organizationmember=to_member, team=team)
+            except IntegrityError:
+                pass
+
+        model_list = (
+            GroupAssignee,
+            GroupBookmark,
+            GroupSeen,
+            GroupShare,
+            GroupSubscription,
+            Activity,
         )
+
+        for model in model_list:
+            for obj in model.objects.filter(
+                user_id=from_user_id, project__organization_id=organization_id
+            ):
+                try:
+                    with transaction.atomic():
+                        obj.update(user_id=to_user_id)
+                except IntegrityError:
+                    pass
+
+    def reset_idp_flags(self, *, organization_id: int) -> None:
+        with in_test_psql_role_override("postgres"):
+            # Flags are not replicated -- these updates are safe without outbox application.
+            OrganizationMember.objects.filter(
+                organization_id=organization_id,
+                flags=models.F("flags").bitor(OrganizationMember.flags["idp:provisioned"]),
+            ).update(
+                flags=models.F("flags")
+                .bitand(~OrganizationMember.flags["idp:provisioned"])
+                .bitand(~OrganizationMember.flags["idp:role-restricted"])
+            )

src/sentry/services/hybrid_cloud/organization/model.py

@@ -84,6 +84,12 @@ class RpcOrganizationMemberFlags(RpcModel):
         item = escape_flag_name(item)
         return bool(getattr(self, item))
 
+    def __setattr__(self, item: str, value: bool) -> None:
+        from sentry.services.hybrid_cloud.organization.serial import escape_flag_name
+
+        item = escape_flag_name(item)
+        super().__setattr__(item, value)
+
     def __getitem__(self, item: str) -> bool:
         return bool(getattr(self, item))
 

src/sentry/services/hybrid_cloud/organization/service.py

@@ -205,6 +205,11 @@ class OrganizationService(RpcService):
     def update_membership_flags(self, *, organization_member: RpcOrganizationMember) -> None:
         pass
 
+    @regional_rpc_method(resolve=ByOrganizationId())
+    @abstractmethod
+    def merge_users(self, *, organization_id: int, from_user_id: int, to_user_id: int) -> None:
+        pass
+
     @regional_rpc_method(resolve=ByOrganizationIdAttribute("organization_member"))
     @abstractmethod
     def get_all_org_roles(

src/sentry/tasks/check_auth.py

@@ -4,7 +4,9 @@ from datetime import timedelta
 from django.utils import timezone
 
 from sentry.auth.exceptions import IdentityNotValid
-from sentry.models import AuthIdentity, OrganizationMember
+from sentry.db.postgres.roles import in_test_psql_role_override
+from sentry.models import AuthIdentity
+from sentry.services.hybrid_cloud.organization import RpcOrganizationMember, organization_service
 from sentry.silo.base import SiloMode
 from sentry.tasks.base import instrumented_task
 from sentry.utils import metrics
@@ -51,13 +53,10 @@ def check_auth_identity(auth_identity_id, **kwargs):
 
     auth_provider = auth_identity.auth_provider
 
-    try:
-        # TODO(hybridcloud) We either need to use orgmembermapping or make this use
-        # RPC services to work on AuthIdentity and AuthProvider.
-        om = OrganizationMember.objects.get(
-            user_id=auth_identity.user.id, organization=auth_provider.organization_id
-        )
-    except OrganizationMember.DoesNotExist:
+    om: RpcOrganizationMember = organization_service.check_membership_by_id(
+        organization_id=auth_provider.organization_id, user_id=auth_identity.user_id
+    )
+    if om is None:
         logger.warning(
             "Removing invalid AuthIdentity(id=%s) due to no organization access", auth_identity_id
         )
@@ -96,9 +95,11 @@ def check_auth_identity(auth_identity_id, **kwargs):
         is_valid = True
 
     if getattr(om.flags, "sso:linked") != is_linked:
-        setattr(om.flags, "sso:linked", is_linked)
-        setattr(om.flags, "sso:invalid", not is_valid)
-        om.update(flags=om.flags)
+        with in_test_psql_role_override("postgres"):
+            # flags are not replicated, so it's ok not to create outboxes here.
+            setattr(om.flags, "sso:linked", is_linked)
+            setattr(om.flags, "sso:invalid", not is_valid)
+            organization_service.update_membership_flags(organization_member=om)
 
     now = timezone.now()
     auth_identity.update(last_verified=now, last_synced=now)

src/sentry/testutils/factories.py

@@ -21,6 +21,7 @@ from django.utils.encoding import force_text
 from django.utils.text import slugify
 
 from sentry.constants import SentryAppInstallationStatus, SentryAppStatus
+from sentry.db.postgres.roles import in_test_psql_role_override
 from sentry.event_manager import EventManager
 from sentry.incidents.logic import (
     create_alert_rule,
@@ -287,6 +288,7 @@ class Factories:
 
     @staticmethod
     @exempt_from_silo_limits()
+    @in_test_psql_role_override("postgres")
     def create_member(teams=None, team_roles=None, **kwargs):
         kwargs.setdefault("role", "member")
         teamRole = kwargs.pop("teamRole", None)
@@ -326,11 +328,6 @@ class Factories:
                 organization=team.organization,
                 defaults={"role": "member"},
             )
-            if created:
-                member.outbox_for_create().drain_shard(max_updates_to_drain=10)
-                organizationmember_mapping_service.create_with_organization_member(
-                    org_member=member
-                )
 
         return OrganizationMemberTeam.objects.create(
             team=team, organizationmember=member, is_active=True, role=role

src/sentry/web/frontend/organization_auth_settings.py

@@ -12,6 +12,7 @@ from rest_framework.request import Request
 from sentry import audit_log, features, roles
 from sentry.auth import manager
 from sentry.auth.helper import AuthHelper
+from sentry.db.postgres.roles import in_test_psql_role_override
 from sentry.models import AuthProvider, Organization, OrganizationMember, User
 from sentry.plugins.base import Response
 from sentry.services.hybrid_cloud.auth import RpcAuthProvider
@@ -85,11 +86,13 @@ class OrganizationAuthSettingsView(OrganizationView):
             data=auth_provider.get_audit_log_data(),
         )
 
-        OrganizationMember.objects.filter(organization=organization).update(
-            flags=F("flags")
-            .bitand(~OrganizationMember.flags["sso:linked"])
-            .bitand(~OrganizationMember.flags["sso:invalid"])
-        )
+        # This is safe -- we're not syncing flags to the org member mapping table.
+        with in_test_psql_role_override("postgres"):
+            OrganizationMember.objects.filter(organization=organization).update(
+                flags=F("flags")
+                .bitand(~OrganizationMember.flags["sso:linked"])
+                .bitand(~OrganizationMember.flags["sso:invalid"])
+            )
 
         user_ids = OrganizationMember.objects.filter(organization=organization).values("user")
         User.objects.filter(id__in=user_ids).update(is_managed=False)