ref(HC): Revokes org and org mapping write permissions, adds outbox handling to org model (#51544)

src/sentry/models/organization.py

@@ -11,6 +11,7 @@ from django.db.models import QuerySet
 from django.urls import NoReverseMatch, reverse
 from django.utils import timezone
 from django.utils.functional import cached_property
+from typing_extensions import override
 
 from bitfield import BitField
 from sentry import features, roles
@@ -265,6 +266,16 @@ class Organization(Model, OptionMixin, OrganizationAbsoluteUrlMixin, SnowflakeId
             with outbox_context(transaction.atomic()):
                 self.save_with_update_outbox(*args, **kwargs)
 
+    # Override for the default update method to ensure that most atomic updates
+    #  generate an outbox alongside any mutations to ensure data is replicated
+    #  properly to the control silo.
+    @override
+    def update(self, *args, **kwargs):
+        with outbox_context(transaction.atomic()):
+            results = super().update(*args, **kwargs)
+            Organization.outbox_for_update(self.id).save()
+            return results
+
     @classmethod
     def reserve_snowflake_id(cls):
         return generate_snowflake_id(cls.snowflake_redis_key)

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

@@ -30,6 +30,7 @@ from sentry.services.hybrid_cloud import OptionValue, logger
 from sentry.services.hybrid_cloud.organization import (
     OrganizationService,
     OrganizationSignalService,
+    RpcOrganization,
     RpcOrganizationFlagsUpdate,
     RpcOrganizationInvite,
     RpcOrganizationMember,
@@ -286,7 +287,9 @@ class DatabaseBackedOrganizationService(OrganizationService):
             else:
                 raise TypeError(f"Invalid value received for update_flags: {name}={value!r}")
 
-        Organization.objects.filter(id=organization_id).update(flags=updates)
+        with outbox_context(transaction.atomic()):
+            Organization.objects.filter(id=organization_id).update(flags=updates)
+            Organization.outbox_for_update(org_id=organization_id).save()
 
     @staticmethod
     def _deserialize_member_flags(flags: RpcOrganizationMemberFlags) -> int:
@@ -387,9 +390,7 @@ class DatabaseBackedOrganizationService(OrganizationService):
             )
         )
 
-    def update_default_role(
-        self, *, organization_id: int, default_role: str
-    ) -> RpcOrganizationMember:
+    def update_default_role(self, *, organization_id: int, default_role: str) -> RpcOrganization:
         org = Organization.objects.get(id=organization_id)
         org.default_role = default_role
         org.save()

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

@@ -3,7 +3,7 @@ from typing import Optional, TypedDict
 from django.db import transaction
 from django.db.models.expressions import CombinedExpression
 
-from sentry.models import Organization, OrganizationStatus
+from sentry.models import Organization, OrganizationStatus, outbox_context
 
 
 class OrganizationCreateAndUpdateOptions(TypedDict, total=False):
@@ -17,42 +17,59 @@ class OrganizationCreateAndUpdateOptions(TypedDict, total=False):
 def create_organization_with_outbox_message(
     *, create_options: OrganizationCreateAndUpdateOptions
 ) -> Organization:
-    with transaction.atomic():
-        org: Organization = Organization.objects.create(**create_options)
-        Organization.outbox_for_update(org_id=org.id).save()
+    org: Organization = Organization.objects.create(**create_options)
     return org
 
 
 def update_organization_with_outbox_message(
     *, org_id: int, update_data: OrganizationCreateAndUpdateOptions
 ) -> Organization:
-    with transaction.atomic():
+    with outbox_context(transaction.atomic()):
         org: Organization = Organization.objects.get(id=org_id)
         org.update(**update_data)
-        Organization.outbox_for_update(org_id=org.id).save()
 
-    org.refresh_from_db()
-    return org
+        org.refresh_from_db()
+        return org
 
 
 def upsert_organization_by_org_id_with_outbox_message(
     *, org_id: int, upsert_data: OrganizationCreateAndUpdateOptions
 ) -> Organization:
-    with transaction.atomic():
+    with outbox_context(transaction.atomic()):
         org, created = Organization.objects.update_or_create(id=org_id, defaults=upsert_data)
-        Organization.outbox_for_update(org_id=org_id).save()
         return org
 
 
 def mark_organization_as_pending_deletion_with_outbox_message(
     *, org_id: int
 ) -> Optional[Organization]:
-    with transaction.atomic():
-        query_result = Organization.objects.filter(
+    with outbox_context(transaction.atomic()):
+        update_count = Organization.objects.filter(
             id=org_id, status=OrganizationStatus.ACTIVE
         ).update(status=OrganizationStatus.PENDING_DELETION)
 
-        if not query_result:
+        if not update_count:
+            return None
+
+        Organization.outbox_for_update(org_id=org_id).save()
+
+        org = Organization.objects.get(id=org_id)
+        return org
+
+
+def unmark_organization_as_pending_deletion_with_outbox_message(
+    *, org_id: int
+) -> Optional[Organization]:
+    with outbox_context(transaction.atomic()):
+        update_count = Organization.objects.filter(
+            id=org_id,
+            status__in=[
+                OrganizationStatus.PENDING_DELETION,
+                OrganizationStatus.DELETION_IN_PROGRESS,
+            ],
+        ).update(status=OrganizationStatus.ACTIVE)
+
+        if not update_count:
             return None
 
         Organization.outbox_for_update(org_id=org_id).save()

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

@@ -1,5 +1,6 @@
 from typing import List, Optional
 
+from sentry.db.postgres.roles import in_test_psql_role_override
 from sentry.models.organizationmapping import OrganizationMapping
 from sentry.services.hybrid_cloud.organization_mapping import (
     OrganizationMappingService,
@@ -68,11 +69,12 @@ class DatabaseBackedOrganizationMappingService(OrganizationMappingService):
     def upsert(
         self, organization_id: int, update: RpcOrganizationMappingUpdate
     ) -> RpcOrganizationMapping:
-        org_mapping, _created = OrganizationMapping.objects.update_or_create(
-            organization_id=organization_id, defaults=update
-        )
+        with in_test_psql_role_override("postgres"):
+            org_mapping, _created = OrganizationMapping.objects.update_or_create(
+                organization_id=organization_id, defaults=update
+            )
 
-        return serialize_organization_mapping(org_mapping)
+            return serialize_organization_mapping(org_mapping)
 
     def verify_mappings(self, organization_id: int, slug: str) -> None:
         try:

src/sentry/testutils/helpers/api_gateway.py

@@ -9,6 +9,7 @@ from rest_framework.response import Response
 
 from sentry.api.base import control_silo_endpoint, region_silo_endpoint
 from sentry.api.bases.organization import OrganizationEndpoint
+from sentry.db.postgres.roles import in_test_psql_role_override
 from sentry.models.organizationmapping import OrganizationMapping
 from sentry.testutils import APITestCase
 from sentry.types.region import Region, RegionCategory, clear_global_regions
@@ -133,9 +134,11 @@ class ApiGatewayTestCase(APITestCase):
             content_type="application/json",
             adding_headers={"test": "header"},
         )
-        OrganizationMapping.objects.get(organization_id=self.organization.id).update(
-            region_name="region1"
-        )
+
+        with in_test_psql_role_override("postgres"):
+            OrganizationMapping.objects.get(organization_id=self.organization.id).update(
+                region_name="region1"
+            )
 
         # Echos the request body and header back for verification
         def return_request_body(request):

src/sentry/utils/migrations.py

@@ -1,5 +1,6 @@
 from django.db.models import F
 
+from sentry.db.postgres.roles import in_test_psql_role_override
 from sentry.utils.query import RangeQuerySetWrapperWithProgressBar
 
 
@@ -14,4 +15,5 @@ def clear_flag(Model, flag_name, flag_attr_name="flags"):
             update_kwargs = {
                 flag_attr_name: F(flag_attr_name).bitand(~getattr(Model, flag_attr_name)[flag_name])
             }
-            Model.objects.filter(id=item.id).update(**update_kwargs)
+            with in_test_psql_role_override("postgres"):
+                Model.objects.filter(id=item.id).update(**update_kwargs)

src/sentry/web/frontend/restore_organization.py

@@ -8,6 +8,9 @@ from sentry import audit_log
 from sentry.api import client
 from sentry.models import Organization, OrganizationStatus
 from sentry.services.hybrid_cloud.organization import organization_service
+from sentry.services.hybrid_cloud.organization_actions.impl import (
+    unmark_organization_as_pending_deletion_with_outbox_message,
+)
 from sentry.web.frontend.base import OrganizationView
 from sentry.web.helpers import render_to_response
 
@@ -65,9 +68,10 @@ class RestoreOrganizationView(OrganizationView):
             messages.add_message(request, messages.ERROR, ERR_MESSAGES[organization.status])
             return self.redirect(reverse("sentry"))
 
-        updated = Organization.objects.filter(
-            id=organization.id, status__in=deletion_statuses
-        ).update(status=OrganizationStatus.ACTIVE)
+        updated = unmark_organization_as_pending_deletion_with_outbox_message(
+            org_id=organization.id
+        )
+
         if updated:
             client.put(
                 f"/organizations/{organization.slug}/",

tests/conftest.py

@@ -172,7 +172,12 @@ def protect_hybrid_cloud_writes_and_deletes(request):
     create Outbox objects in the same transaction that matches what you delete.
     """
     from sentry.db.models.fields.hybrid_cloud_foreign_key import HybridCloudForeignKey
-    from sentry.models import OrganizationMember, OrganizationMemberMapping
+    from sentry.models import (
+        Organization,
+        OrganizationMapping,
+        OrganizationMember,
+        OrganizationMemberMapping,
+    )
     from sentry.testutils.silo import iter_models, reset_test_role, restrict_role
 
     try:
@@ -205,6 +210,10 @@ def protect_hybrid_cloud_writes_and_deletes(request):
     # outboxes in a transaction, and cover that transaction with `in_test_psql_role_override`
     restrict_role(role="postgres_unprivileged", model=OrganizationMember, revocation_type="INSERT")
     restrict_role(role="postgres_unprivileged", model=OrganizationMember, revocation_type="UPDATE")
+    restrict_role(role="postgres_unprivileged", model=Organization, revocation_type="INSERT")
+    restrict_role(role="postgres_unprivileged", model=Organization, revocation_type="UPDATE")
+    restrict_role(role="postgres_unprivileged", model=OrganizationMapping, revocation_type="INSERT")
+    restrict_role(role="postgres_unprivileged", model=OrganizationMapping, revocation_type="UPDATE")
     # OrganizationMember objects need to cascade, but they can't use the standard hybrid cloud foreign key because the
     # identifiers are not snowflake ids.
     restrict_role(role="postgres_unprivileged", model=OrganizationMember, revocation_type="DELETE")

tests/sentry/api/endpoints/test_accept_organization_invite.py

@@ -156,13 +156,14 @@ class AcceptInviteTest(TestCase, HybridCloudTestMixin):
                 ),
             ]
         ):
-            self.create_organization_mapping(
-                organization_id=101010,
-                slug="abcslug",
-                name="The Thing",
-                idempotency_key="",
-                region_name="some-region",
-            )
+            with in_test_psql_role_override("postgres"):
+                self.create_organization_mapping(
+                    organization_id=101010,
+                    slug="abcslug",
+                    name="The Thing",
+                    idempotency_key="",
+                    region_name="some-region",
+                )
             self._require_2fa_for_organization()
             assert not self.user.has_2fa()
 

tests/sentry/hybrid_cloud/test_organizationmapping.py

@@ -1,7 +1,7 @@
 import pytest
 from django.db import IntegrityError
 
-from sentry.models import Organization, outbox_context
+from sentry.models import outbox_context
 from sentry.models.organization import OrganizationStatus
 from sentry.models.organizationmapping import OrganizationMapping
 from sentry.services.hybrid_cloud.organization_mapping import (
@@ -9,7 +9,6 @@ from sentry.services.hybrid_cloud.organization_mapping import (
     organization_mapping_service,
 )
 from sentry.testutils import TransactionTestCase
-from sentry.testutils.factories import Factories
 from sentry.testutils.outbox import outbox_runner
 from sentry.testutils.silo import control_silo_test, exempt_from_silo_limits
 
@@ -18,10 +17,9 @@ from sentry.testutils.silo import control_silo_test, exempt_from_silo_limits
 class OrganizationMappingTest(TransactionTestCase):
     def test_create_on_organization_save(self):
         with outbox_context(flush=False), exempt_from_silo_limits():
-            organization = Organization(
+            organization = self.create_organization(
                 name="test name",
             )
-            organization.save()
 
         # Validate that organization mapping has not been created
         with pytest.raises(OrganizationMapping.DoesNotExist):
@@ -64,16 +62,11 @@ class OrganizationMappingTest(TransactionTestCase):
 
     def test_upsert__update_if_found(self):
         with exempt_from_silo_limits():
-            self.organization = Organization(
+            self.organization = self.create_organization(
                 name="test name",
                 slug="foobar",
             )
 
-            self.organization.save()
-
-        with outbox_runner():
-            pass
-
         fixture_org_mapping = OrganizationMapping.objects.get(organization_id=self.organization.id)
 
         organization_mapping_service.upsert(
@@ -89,7 +82,7 @@ class OrganizationMappingTest(TransactionTestCase):
         assert fixture_org_mapping.status == OrganizationStatus.PENDING_DELETION
 
     def test_upsert__duplicate_slug(self):
-        self.organization = Factories.create_organization(slug="alreadytaken")
+        self.organization = self.create_organization(slug="alreadytaken")
 
         with pytest.raises(IntegrityError):
             organization_mapping_service.upsert(