chore(hybrid-cloud): Update bind_org_context_from_integration for split db (#54812)

src/sentry/integrations/utils/scope.py

@@ -1,12 +1,16 @@
 from __future__ import annotations
 
 import logging
-from typing import Any, Mapping, Sequence
+from typing import Any, List, Mapping
 
 from sentry_sdk import configure_scope
 
 from sentry.models.organization import Organization
+from sentry.models.organizationmapping import OrganizationMapping
 from sentry.services.hybrid_cloud.integration import integration_service
+from sentry.services.hybrid_cloud.integration.model import RpcOrganizationIntegration
+from sentry.services.hybrid_cloud.organization import organization_service
+from sentry.silo.base import SiloMode
 from sentry.utils.sdk import (
     bind_ambiguous_org_context,
     bind_organization_context,
@@ -33,9 +37,11 @@ def clear_tags_and_context() -> None:
             logger.info("We've reset the context and tags.")
 
 
-def get_orgs_from_integration(integration_id: int) -> Sequence[Organization]:
+def get_org_integrations(
+    integration_id: int,
+) -> List[RpcOrganizationIntegration]:
     """
-    Given the id of an `Integration`, return a list of associated `Organization` objects.
+    Given the id of an `Integration`, return a list of associated `RpcOrganizationIntegration` objects.
 
     Note: An `Integration` is an instance of given provider's integration, tied to a single entity
     on the provider's end (for example, an instance of the GitHub integration tied to a particular
@@ -46,11 +52,8 @@ def get_orgs_from_integration(integration_id: int) -> Sequence[Organization]:
     _, org_integrations = integration_service.get_organization_contexts(
         integration_id=integration_id
     )
-    orgs = Organization.objects.get_many_from_cache(
-        [org_integration.organization_id for org_integration in org_integrations]
-    )
 
-    return orgs
+    return org_integrations
 
 
 def bind_org_context_from_integration(
@@ -67,9 +70,9 @@ def bind_org_context_from_integration(
     Integration or an RpcIntegration object, because corresponding ones share the same id.
     """
 
-    orgs = get_orgs_from_integration(integration_id)
+    org_integrations = get_org_integrations(integration_id)
 
-    if len(orgs) == 0:
+    if len(org_integrations) == 0:
         logger.warning(
             f"Can't bind org context - no orgs are associated with integration id={integration_id}.",
             extra=extra,
@@ -82,7 +85,29 @@ def bind_org_context_from_integration(
         # (With `add_to_scope=False`, we still log a warning - separate from the one above - on data
         # mismatch.)
         check_tag_for_scope_bleed("integration_id", integration_id, add_to_scope=False)
-    elif len(orgs) == 1:
-        bind_organization_context(orgs[0])
+    elif len(org_integrations) == 1:
+        org_integration = org_integrations[0]
+        org = organization_service.get_organization_by_id(id=org_integration.organization_id)
+        if org is not None:
+            bind_organization_context(org.organization)
+        else:
+            logger.exception(
+                f"Unable to call organization_service.get_organization_by_id with organization id={org_integration.organization_id}.",
+                extra=extra,
+            )
     else:
-        bind_ambiguous_org_context(orgs, f"integration (id={integration_id})")
+        org_ids = [org_integration.organization_id for org_integration in org_integrations]
+        org_slugs = []
+
+        if SiloMode.get_current_mode() == SiloMode.CONTROL:
+            org_slugs = [
+                org_mapping.slug
+                for org_mapping in OrganizationMapping.objects.filter(organization_id__in=org_ids)
+            ]
+        else:
+            orgs = Organization.objects.get_many_from_cache(
+                [org_integration.organization_id for org_integration in org_integrations]
+            )
+            org_slugs = [org.slug for org in orgs]
+
+        bind_ambiguous_org_context(org_slugs, f"integration (id={integration_id})")

src/sentry/utils/sdk.py

@@ -4,7 +4,7 @@ import copy
 import inspect
 import logging
 import random
-from typing import TYPE_CHECKING, Any, Mapping, Sequence
+from typing import TYPE_CHECKING, Any, List, Mapping, Sequence
 
 import sentry_sdk
 from django.conf import settings
@@ -657,7 +657,7 @@ def bind_organization_context(organization: Organization | RpcOrganization) -> N
 
 
 def bind_ambiguous_org_context(
-    orgs: Sequence[Organization | RpcOrganization], source: str | None = None
+    orgs: Sequence[Organization] | Sequence[RpcOrganization] | List[str], source: str | None = None
 ) -> None:
     """
     Add org context information to the scope in the case where the current org might be one of a
@@ -667,7 +667,12 @@ def bind_ambiguous_org_context(
 
     MULTIPLE_ORGS_TAG = "[multiple orgs]"
 
-    org_slugs = [org.slug for org in orgs]
+    def parse_org_slug(x: Organization | RpcOrganization | str) -> str:
+        if isinstance(x, str):
+            return x
+        return x.slug
+
+    org_slugs = [parse_org_slug(org) for org in orgs]
 
     # Right now there is exactly one Integration instance shared by more than 30 orgs (the generic
     # GitLab integration, at the moment shared by ~500 orgs), so 50 should be plenty for all but

tests/sentry/integrations/jira/test_uninstalled.py

@@ -6,8 +6,10 @@ import responses
 from sentry.constants import ObjectStatus
 from sentry.integrations.utils import get_query_hash
 from sentry.models import Integration
+from sentry.services.hybrid_cloud.organization.serial import serialize_rpc_organization
+from sentry.silo.base import SiloMode
 from sentry.testutils.cases import APITestCase
-from sentry.testutils.silo import control_silo_test
+from sentry.testutils.silo import assume_test_silo_mode, control_silo_test
 from sentry.utils.http import absolute_uri
 from tests.sentry.utils.test_jwt import RS256_KEY, RS256_PUB_KEY
 
@@ -65,7 +67,8 @@ class JiraUninstalledTest(APITestCase):
         integration = Integration.objects.get(id=integration.id)
 
         mock_set_tag.assert_called_with("integration_id", integration.id)
-        mock_bind_org_context.assert_called_with(org)
+        with assume_test_silo_mode(SiloMode.REGION):
+            mock_bind_org_context.assert_called_with(serialize_rpc_organization(org))
         assert integration.status == ObjectStatus.DISABLED
         assert resp.status_code == 200
 
@@ -93,6 +96,7 @@ class JiraUninstalledTest(APITestCase):
         integration = Integration.objects.get(id=integration.id)
 
         mock_set_tag.assert_called_with("integration_id", integration.id)
-        mock_bind_org_context.assert_called_with(org)
+        with assume_test_silo_mode(SiloMode.REGION):
+            mock_bind_org_context.assert_called_with(serialize_rpc_organization(org))
         assert integration.status == ObjectStatus.DISABLED
         assert resp.status_code == 200

tests/sentry/integrations/jira/test_webhooks.py

@@ -14,6 +14,7 @@ from sentry.integrations.mixins import IssueSyncMixin
 from sentry.integrations.utils import AtlassianConnectValidationError
 from sentry.models import Integration
 from sentry.services.hybrid_cloud.integration.serial import serialize_integration
+from sentry.services.hybrid_cloud.organization.serial import serialize_rpc_organization
 from sentry.shared_integrations.exceptions.base import ApiError
 from sentry.testutils.cases import APITestCase, TestCase
 
@@ -171,7 +172,7 @@ class JiraIssueUpdatedWebhookTest(APITestCase):
             self.get_success_response(**data, extra_headers=dict(HTTP_AUTHORIZATION=TOKEN))
 
             mock_set_tag.assert_called_with("integration_id", self.integration.id)
-            mock_bind_org_context.assert_called_with(self.organization)
+            mock_bind_org_context.assert_called_with(serialize_rpc_organization(self.organization))
 
     def test_missing_changelog(self):
         with patch(

tests/sentry/integrations/utils/test_scope.py

@@ -1,65 +1,83 @@
 from unittest.mock import MagicMock, patch
 
-from sentry.integrations.utils.scope import (
-    bind_org_context_from_integration,
-    get_orgs_from_integration,
-)
+from sentry.integrations.utils.scope import bind_org_context_from_integration, get_org_integrations
 from sentry.models.integrations.integration import Integration
+from sentry.models.integrations.organization_integration import OrganizationIntegration
+from sentry.services.hybrid_cloud.integration.serial import serialize_organization_integration
+from sentry.services.hybrid_cloud.organization.serial import serialize_rpc_organization
+from sentry.silo.base import SiloMode
 from sentry.testutils.cases import TestCase
+from sentry.testutils.silo import all_silo_test, assume_test_silo_mode
 
 
+@all_silo_test(stable=True)
 class GetOrgsFromIntegrationTest(TestCase):
     def test_finds_single_org(self):
         org = self.create_organization(slug="dogsaregreat")
-        integration = Integration.objects.create(name="squirrelChasers")
-        integration.add_organization(org)
+        with assume_test_silo_mode(SiloMode.CONTROL):
+            integration = Integration.objects.create(name="squirrelChasers")
+            integration.add_organization(org)
 
-        found_orgs = get_orgs_from_integration(integration.id)
+        actual = get_org_integrations(integration.id)
 
-        assert found_orgs == [org]
+        with assume_test_silo_mode(SiloMode.CONTROL):
+            ois = OrganizationIntegration.objects.filter(
+                organization_id=org.id, integration_id=integration.id
+            ).all()
+
+        assert actual == [serialize_organization_integration(oi) for oi in ois]
 
     def test_finds_multiple_orgs(self):
         maisey_org = self.create_organization(slug="themaiseymaiseydog")
         charlie_org = self.create_organization(slug="charliebear")
-        integration = Integration.objects.create(name="squirrelChasers")
-        integration.add_organization(maisey_org)
-        integration.add_organization(charlie_org)
+        with assume_test_silo_mode(SiloMode.CONTROL):
+            integration = Integration.objects.create(name="squirrelChasers")
+            integration.add_organization(maisey_org)
+            integration.add_organization(charlie_org)
 
-        found_orgs = get_orgs_from_integration(integration.id)
+        actual = get_org_integrations(integration.id)
 
-        assert found_orgs == [maisey_org, charlie_org]
+        with assume_test_silo_mode(SiloMode.CONTROL):
+            ois = OrganizationIntegration.objects.filter(integration_id=integration.id).all()
+        expected = [serialize_organization_integration(oi) for oi in ois]
+        assert actual == expected
 
     def test_finds_no_orgs_without_erroring(self):
-        integration = Integration.objects.create(name="squirrelChasers")
+        with assume_test_silo_mode(SiloMode.CONTROL):
+            integration = Integration.objects.create(name="squirrelChasers")
 
-        found_orgs = get_orgs_from_integration(integration.id)
+        actual = get_org_integrations(integration.id)
 
-        assert found_orgs == []
+        assert actual == []
 
 
+@all_silo_test(stable=True)
 class BindOrgContextFromIntegrationTest(TestCase):
     @patch("sentry.integrations.utils.scope.bind_organization_context")
     def test_binds_org_context_with_single_org(self, mock_bind_org_context: MagicMock):
         org = self.create_organization(slug="dogsaregreat")
-        integration = Integration.objects.create(name="squirrelChasers")
-        integration.add_organization(org)
+        with assume_test_silo_mode(SiloMode.CONTROL):
+            integration = Integration.objects.create(name="squirrelChasers")
+            integration.add_organization(org)
 
         bind_org_context_from_integration(integration.id)
 
-        mock_bind_org_context.assert_called_with(org)
+        with assume_test_silo_mode(SiloMode.REGION):
+            mock_bind_org_context.assert_called_with(serialize_rpc_organization(org))
 
     @patch("sentry.integrations.utils.scope.bind_ambiguous_org_context")
     def test_binds_org_context_with_multiple_orgs(self, mock_bind_ambiguous_org_context: MagicMock):
         maisey_org = self.create_organization(slug="themaiseymaiseydog")
         charlie_org = self.create_organization(slug="charliebear")
-        integration = Integration.objects.create(name="squirrelChasers")
-        integration.add_organization(maisey_org)
-        integration.add_organization(charlie_org)
+        with assume_test_silo_mode(SiloMode.CONTROL):
+            integration = Integration.objects.create(name="squirrelChasers")
+            integration.add_organization(maisey_org)
+            integration.add_organization(charlie_org)
 
         bind_org_context_from_integration(integration.id)
 
         mock_bind_ambiguous_org_context.assert_called_with(
-            [maisey_org, charlie_org], f"integration (id={integration.id})"
+            [maisey_org.slug, charlie_org.slug], f"integration (id={integration.id})"
         )
 
     @patch("sentry.integrations.utils.scope.bind_ambiguous_org_context")
@@ -73,7 +91,8 @@ class BindOrgContextFromIntegrationTest(TestCase):
         mock_bind_org_context: MagicMock,
         mock_bind_ambiguous_org_context: MagicMock,
     ):
-        integration = Integration.objects.create(name="squirrelChasers")
+        with assume_test_silo_mode(SiloMode.CONTROL):
+            integration = Integration.objects.create(name="squirrelChasers")
 
         bind_org_context_from_integration(integration.id, {"webhook": "issue_updated"})
         mock_logger_warning.assert_called_with(