fix(integrations): Add RPC method to start grace period for given organization + provider (#96613)

Author: Christinarlong

src/sentry/integrations/services/integration/impl.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import logging
+from collections import defaultdict
 from collections.abc import Iterable
 from typing import TYPE_CHECKING, Any
 
@@ -12,7 +13,7 @@
     AlertRuleUiComponentWebhookSentEvent,
 )
 from sentry.api.paginator import OffsetPaginator
-from sentry.constants import SentryAppInstallationStatus
+from sentry.constants import ObjectStatus, SentryAppInstallationStatus
 from sentry.hybridcloud.rpc.pagination import RpcPaginationArgs, RpcPaginationResult
 from sentry.incidents.models.incident import INCIDENT_STATUS, IncidentStatus
 from sentry.integrations.messaging.metrics import (
@@ -358,6 +359,60 @@ def update_organization_integration(
         )
         return ois[0] if len(ois) > 0 else None
 
+    def start_grace_period_for_provider(
+        self,
+        *,
+        organization_id: int,
+        provider: str,
+        grace_period_end: datetime,
+        status: int | None = ObjectStatus.ACTIVE,
+        skip_oldest: bool = False,
+    ) -> list[RpcOrganizationIntegration]:
+        filter_kwargs = {
+            "organization_id": organization_id,
+            "integration__provider": provider,
+        }
+        all_ois_filter_kwargs: dict[str, Any] = {
+            "integration__provider": provider,
+        }
+
+        if status is not None:
+            filter_kwargs["status"] = status
+            all_ois_filter_kwargs["status"] = status
+
+        current_org_ois = OrganizationIntegration.objects.filter(**filter_kwargs)
+        ois_to_update = list(current_org_ois.values_list("id", flat=True))
+
+        if skip_oldest:
+            all_ois_filter_kwargs["integration__in"] = current_org_ois.values_list(
+                "integration_id", flat=True
+            )
+
+            # Get all associated OrganizationIntegrations for the Integrations used by this org
+            all_ois = (
+                OrganizationIntegration.objects.filter(**all_ois_filter_kwargs)
+                .order_by("id")
+                .distinct()
+            )
+
+            # Create mapping of integration_id to list of OrganizationIntegrations
+            integration_to_ois: dict[int, list[OrganizationIntegration]] = defaultdict(list)
+            for oi in all_ois:
+                integration_to_ois[oi.integration_id].append(oi)
+
+            for integration, ois in integration_to_ois.items():
+                # Check if the oldest OrganizationIntegration for the Integration belongs to THIS organization
+                # If not we want to start the grace period for this org's OrganizationIntegration
+                if integration_to_ois[integration][0].organization_id == organization_id:
+                    ois_to_update.remove(integration_to_ois[integration][0].id)
+
+        updated_ois = self.update_organization_integrations(
+            org_integration_ids=ois_to_update,
+            grace_period_end=grace_period_end,
+        )
+
+        return updated_ois
+
     def add_organization(self, *, integration_id: int, org_ids: list[int]) -> RpcIntegration | None:
         try:
             integration = Integration.objects.get(id=integration_id)

src/sentry/integrations/services/integration/service.py

@@ -7,6 +7,7 @@
 from datetime import datetime
 from typing import Any
 
+from sentry.constants import ObjectStatus
 from sentry.hybridcloud.rpc.pagination import RpcPaginationArgs, RpcPaginationResult
 from sentry.hybridcloud.rpc.service import RpcService, rpc_method
 from sentry.integrations.services.integration import RpcIntegration, RpcOrganizationIntegration
@@ -119,6 +120,31 @@ def get_organization_integration(
         )
         return ois[0] if len(ois) > 0 else None
 
+    @rpc_method
+    @abstractmethod
+    def start_grace_period_for_provider(
+        self,
+        *,
+        organization_id: int,
+        provider: str,
+        grace_period_end: datetime,
+        status: int | None = ObjectStatus.ACTIVE,
+        skip_oldest: bool = False,
+    ) -> list[RpcOrganizationIntegration]:
+        """
+        Start grace period for all OrganizationIntegrations of a given provider for an organization
+
+        Args:
+            organization_id (int): The Organization whose OrganizationIntegrations will be grace perioded
+            provider (str): The provider key - e.g. "github"
+            grace_period_end (datetime): The grace period end date
+            status (int, optional): The status of the OrganizationIntegrations. Defaults to ObjectStatus.ACTIVE. Put None to include all statuses.
+            skip_oldest (bool, optional): Flag for if we want to skip grace period for the oldest OrganizationIntegration per Integration. Defaults to False.
+
+        Returns:
+            list[RpcOrganizationIntegration]: The updated OrganizationIntegrations
+        """
+
     @rpc_method
     @abstractmethod
     def organization_context(

tests/sentry/integrations/services/test_integration.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from datetime import datetime, timezone
+from datetime import datetime, timedelta, timezone
 from unittest.mock import MagicMock, patch
 
 from sentry.constants import ObjectStatus
@@ -354,3 +354,234 @@ def test_send_incident_alert_missing_sentryapp(self, mock_record: MagicMock) ->
         assert_count_of_metric(
             mock_record=mock_record, outcome=EventLifecycleOutcome.FAILURE, outcome_count=1
         )
+
+
+@all_silo_test
+class StartGracePeriodForProviderTest(TestCase):
+    def setUp(self) -> None:
+        super().setUp()
+        with assume_test_silo_mode(SiloMode.REGION):
+            self.org1 = self.organization
+            self.org2 = self.create_organization(name="Test Org 2")
+            self.org3 = self.create_organization(name="Test Org 3")
+
+        with assume_test_silo_mode(SiloMode.CONTROL):
+
+            self.github_integration_1 = self.create_integration(
+                organization=self.org2,
+                name="GitHub Integration 1",
+                provider="github",
+                external_id="github:repo1",
+                status=ObjectStatus.ACTIVE,
+            )
+
+            self.github_integration_2 = self.create_integration(
+                organization=self.org1,
+                name="GitHub Integration 2",
+                provider="github",
+                external_id="github:repo2",
+                status=ObjectStatus.ACTIVE,
+            )
+
+            now = datetime.now(timezone.utc)
+            # Add the same GitHub integrations to other orgs (multi-org scenario)
+            # Org2 uses github_integration_1 and is older than org1's so org1 should get the grace period
+            self.org1_github_oi_1 = self.create_organization_integration(
+                integration=self.github_integration_1,
+                organization_id=self.org1.id,
+                status=ObjectStatus.ACTIVE,
+                date_added=now,  # Older than org 2
+            )
+
+            # Org3 uses github_integration_2 and is newer than org1's
+            self.org3_github_oi_2 = self.create_organization_integration(
+                integration=self.github_integration_2,
+                organization_id=self.org3.id,
+                status=ObjectStatus.ACTIVE,
+                date_added=now + timedelta(days=5),  # Newer than org1
+            )
+
+            # Get org2's OrganizationIntegration for github_integration_1
+            self.org2_github_oi_1 = OrganizationIntegration.objects.get(
+                organization_id=self.org2.id, integration_id=self.github_integration_1.id
+            )
+            # Get org1's OrganizationIntegration for github_integration_2
+            self.org1_github_oi_2 = OrganizationIntegration.objects.get(
+                organization_id=self.org1.id, integration_id=self.github_integration_2.id
+            )
+
+    @freeze_time()
+    def test_start_grace_period_for_provider_github_with_skip_oldest(self) -> None:
+        grace_period_end = datetime.now(timezone.utc) + timedelta(days=7)
+        with assume_test_silo_mode(SiloMode.REGION):
+            grace_perioded_ois = integration_service.start_grace_period_for_provider(
+                organization_id=self.org1.id,
+                provider="github",
+                grace_period_end=grace_period_end,
+                status=ObjectStatus.ACTIVE,
+                skip_oldest=True,
+            )
+
+        # Expected behavior with skip_oldest=True:
+        # - org1_github_oi_1 should get grace perioded since org 2's OI is older
+        # - org1_github_oi_2 should not get grace perioded (since it's the oldest OI for this integration)
+        # - org3_github_oi_2 should not get grace perioded (OIs that aren't from the downgrading org should be untouched)
+
+        lof_grace_perioded_ois = [oi.id for oi in grace_perioded_ois]
+
+        with assume_test_silo_mode(SiloMode.CONTROL):
+            self.org1_github_oi_1.refresh_from_db()
+            self.org1_github_oi_2.refresh_from_db()
+            self.org3_github_oi_2.refresh_from_db()
+            self.org2_github_oi_1.refresh_from_db()
+
+        # Assertions for github_integration_1
+        assert self.org1_github_oi_1.id in lof_grace_perioded_ois
+        assert self.org1_github_oi_1.grace_period_end == grace_period_end
+        assert self.org2_github_oi_1.grace_period_end is None
+        assert self.org2_github_oi_1.id not in lof_grace_perioded_ois
+
+        # Assertions for github_integration_2
+        assert self.org1_github_oi_2.id not in lof_grace_perioded_ois
+        assert self.org1_github_oi_2.grace_period_end is None
+        assert self.org3_github_oi_2.id not in lof_grace_perioded_ois
+        assert self.org3_github_oi_2.grace_period_end is None
+
+        assert len(lof_grace_perioded_ois) == 1, "Should only include org1_github_oi_1"
+
+    def test_start_grace_period_for_provider_github_without_skip_oldest(self) -> None:
+        grace_period_end = datetime.now(timezone.utc) + timedelta(days=7)
+
+        grace_perioded_ois = integration_service.start_grace_period_for_provider(
+            organization_id=self.org1.id,
+            provider="github",
+            grace_period_end=grace_period_end,
+            status=ObjectStatus.ACTIVE,
+            skip_oldest=False,
+        )
+
+        # Expected behavior with skip_oldest=False:
+        # - Both org1's GitHub OrganizationIntegrations should get grace perioded
+        # - Other orgs' OIs should NOT be included
+
+        lof_grace_perioded_ois = [oi.id for oi in grace_perioded_ois]
+
+        with assume_test_silo_mode(SiloMode.CONTROL):
+            self.org1_github_oi_1.refresh_from_db()
+            self.org1_github_oi_2.refresh_from_db()
+
+            ois_grace_perioded = OrganizationIntegration.objects.filter(
+                organization_id=self.org1.id,
+                grace_period_end__isnull=False,
+            )
+
+        for oi in ois_grace_perioded:
+            assert oi.id in lof_grace_perioded_ois
+            assert oi.grace_period_end == grace_period_end
+
+        assert self.org1_github_oi_1.id in lof_grace_perioded_ois
+        assert self.org1_github_oi_2.id in lof_grace_perioded_ois
+        assert self.org1_github_oi_1.grace_period_end == grace_period_end
+        assert self.org1_github_oi_2.grace_period_end == grace_period_end
+
+        assert len(lof_grace_perioded_ois) == 2, "Both org1's GitHub OIs should be grace perioded"
+
+    def test_start_grace_period_for_provider_github_for_all_statuses(self) -> None:
+        grace_period_end = datetime.now(timezone.utc) + timedelta(days=7)
+        self.github_integration_3 = self.create_integration(
+            organization=self.org1,
+            name="GitHub Integration 3",
+            provider="github",
+            external_id="github:repo3",
+            status=None,
+        )
+        with assume_test_silo_mode(SiloMode.CONTROL):
+            self.org1_github_oi_3 = OrganizationIntegration.objects.get(
+                organization_id=self.org1.id, integration_id=self.github_integration_3.id
+            )
+            self.org1_github_oi_3.status = ObjectStatus.HIDDEN
+            self.org1_github_oi_3.save()
+
+        grace_perioded_ois = integration_service.start_grace_period_for_provider(
+            organization_id=self.org1.id,
+            provider="github",
+            grace_period_end=grace_period_end,
+            status=None,
+            skip_oldest=False,
+        )
+
+        # Expected behavior with skip_oldest=False:
+        # - Both org1's GitHub OrganizationIntegrations should get grace perioded
+        # - Other orgs' OIs should NOT be included
+
+        lof_grace_perioded_ois = [oi.id for oi in grace_perioded_ois]
+
+        with assume_test_silo_mode(SiloMode.CONTROL):
+            self.org1_github_oi_1.refresh_from_db()
+            self.org1_github_oi_2.refresh_from_db()
+            self.org1_github_oi_3.refresh_from_db()
+
+            ois_grace_perioded = OrganizationIntegration.objects.filter(
+                organization_id=self.org1.id,
+                grace_period_end__isnull=False,
+            )
+
+        for oi in ois_grace_perioded:
+            assert oi.id in lof_grace_perioded_ois
+            assert oi.grace_period_end == grace_period_end
+
+        assert len(lof_grace_perioded_ois) == 3, "All org1's GitHub OIs should be grace perioded"
+
+    def test_start_grace_period_for_provider_github_for_all_statuses_with_skip_oldest(self) -> None:
+        grace_period_end = datetime.now(timezone.utc) + timedelta(days=7)
+        self.github_integration_3 = self.create_integration(
+            organization=self.org1,
+            name="GitHub Integration 3",
+            provider="github",
+            external_id="github:repo3",
+            status=None,
+        )
+        with assume_test_silo_mode(SiloMode.CONTROL):
+            self.org1_github_oi_3 = OrganizationIntegration.objects.get(
+                organization_id=self.org1.id, integration_id=self.github_integration_3.id
+            )
+            self.org1_github_oi_3.status = ObjectStatus.HIDDEN
+            self.org1_github_oi_3.save()
+
+        grace_perioded_ois = integration_service.start_grace_period_for_provider(
+            organization_id=self.org1.id,
+            provider="github",
+            grace_period_end=grace_period_end,
+            status=None,
+            skip_oldest=True,
+        )
+
+        # Expected behavior with skip_oldest=False:
+        # - Both org1's GitHub OrganizationIntegrations should get grace perioded
+        # - Other orgs' OIs should NOT be included
+
+        lof_grace_perioded_ois = [oi.id for oi in grace_perioded_ois]
+
+        with assume_test_silo_mode(SiloMode.CONTROL):
+            self.org1_github_oi_1.refresh_from_db()
+            self.org1_github_oi_2.refresh_from_db()
+            self.org1_github_oi_3.refresh_from_db()
+            ois_grace_perioded = OrganizationIntegration.objects.filter(
+                organization_id=self.org1.id,
+                grace_period_end__isnull=False,
+            )
+
+        for oi in ois_grace_perioded:
+            assert oi.id in lof_grace_perioded_ois
+            assert oi.grace_period_end == grace_period_end
+
+        assert (
+            len(lof_grace_perioded_ois) == 1
+        ), "Only org1_github_oi_1 should be grace perioded since it's NOT the oldest OI for its integration"
+
+        assert self.org1_github_oi_1.id in lof_grace_perioded_ois
+        assert self.org1_github_oi_1.grace_period_end == grace_period_end
+        assert self.org1_github_oi_2.id not in lof_grace_perioded_ois
+        assert self.org1_github_oi_2.grace_period_end is None
+        assert self.org1_github_oi_3.id not in lof_grace_perioded_ois
+        assert self.org1_github_oi_3.grace_period_end is None