Merge pull request #707 from grafana/add-region-to-organization

Add region info to organizations

engine/apps/api/serializers/organization.py

@@ -170,4 +170,14 @@ class PluginOrganizationSerializer(serializers.ModelSerializer):
 
     class Meta:
         model = Organization
-        fields = ["pk", "stack_id", "stack_slug", "grafana_url", "org_id", "org_slug", "org_title", "grafana_token"]
+        fields = [
+            "pk",
+            "stack_id",
+            "stack_slug",
+            "grafana_url",
+            "org_id",
+            "org_slug",
+            "org_title",
+            "region_slug",
+            "grafana_token",
+        ]

engine/apps/auth_token/auth.py

@@ -11,6 +11,7 @@ from rest_framework.request import Request
 from apps.grafana_plugin.helpers.gcom import check_token
 from apps.user_management.models import User
 from apps.user_management.models.organization import Organization
+from apps.user_management.models.region import OrganizationMovedException
 from common.constants.role import Role
 
 from .constants import SCHEDULE_EXPORT_TOKEN_NAME, SLACK_AUTH_TOKEN_NAME
@@ -46,6 +47,10 @@ class ApiTokenAuthentication(BaseAuthentication):
             auth_token = self.model.validate_token_string(token)
         except InvalidToken:
             raise exceptions.AuthenticationFailed("Invalid token.")
+
+        if auth_token.organization.is_moved:
+            raise OrganizationMovedException(auth_token.organization)
+
         return auth_token.user, auth_token
 
 
@@ -167,6 +172,9 @@ class ScheduleExportAuthentication(BaseAuthentication):
         except InvalidToken:
             raise exceptions.AuthenticationFailed("Invalid token.")
 
+        if auth_token.organization.is_moved:
+            raise OrganizationMovedException(auth_token.organization)
+
         if auth_token.schedule.public_primary_key != public_primary_key:
             raise exceptions.AuthenticationFailed("Invalid schedule export token for schedule")
 
@@ -197,6 +205,9 @@ class UserScheduleExportAuthentication(BaseAuthentication):
         except InvalidToken:
             raise exceptions.AuthenticationFailed("Invalid token")
 
+        if auth_token.organization.is_moved:
+            raise OrganizationMovedException(auth_token.organization)
+
         if auth_token.user.public_primary_key != public_primary_key:
             raise exceptions.AuthenticationFailed("Invalid schedule export token for user")
 

engine/apps/grafana_plugin/helpers/client.py

@@ -143,3 +143,6 @@ class GcomAPIClient(APIClient):
 
     def post_active_users(self, body):
         return self.api_post("app-active-users", body)
+
+    def get_stack_regions(self):
+        return self.api_get("stack-regions")

engine/apps/grafana_plugin/helpers/gcom.py

@@ -57,6 +57,7 @@ def check_gcom_permission(token_string: str, context) -> Optional["GcomToken"]:
                 org_id=str(instance_info["orgId"]),
                 org_slug=instance_info["orgSlug"],
                 org_title=instance_info["orgName"],
+                region_slug=instance_info["regionSlug"],
                 gcom_token=token_string,
                 gcom_token_org_last_time_synced=timezone.now(),
             )
@@ -64,6 +65,7 @@ def check_gcom_permission(token_string: str, context) -> Optional["GcomToken"]:
         organization.stack_slug = instance_info["slug"]
         organization.org_slug = instance_info["orgSlug"]
         organization.org_title = instance_info["orgName"]
+        organization.region_slug = instance_info["regionSlug"]
         organization.grafana_url = instance_info["url"]
         organization.gcom_token = token_string
         organization.gcom_token_org_last_time_synced = timezone.now()
@@ -72,6 +74,7 @@ def check_gcom_permission(token_string: str, context) -> Optional["GcomToken"]:
                 "stack_slug",
                 "org_slug",
                 "org_title",
+                "region_slug",
                 "grafana_url",
                 "gcom_token",
                 "gcom_token_org_last_time_synced",
@@ -109,3 +112,16 @@ def get_active_instance_ids() -> Tuple[Optional[set], bool]:
 
 def get_deleted_instance_ids() -> Tuple[Optional[set], bool]:
     return get_instance_ids(GcomAPIClient.DELETED_INSTANCE_QUERY)
+
+
+def get_stack_regions() -> Tuple[Optional[set], bool]:
+    if not settings.GRAFANA_COM_API_TOKEN or settings.LICENSE != settings.CLOUD_LICENSE_NAME:
+        return None, False
+
+    client = GcomAPIClient(settings.GRAFANA_COM_API_TOKEN)
+    regions, status = client.get_stack_regions()
+
+    if not regions or "items" not in regions:
+        return None, True
+
+    return regions["items"], True

engine/apps/grafana_plugin/tasks/sync.py

@@ -5,8 +5,9 @@ from django.conf import settings
 from django.utils import timezone
 
 from apps.grafana_plugin.helpers import GcomAPIClient
-from apps.grafana_plugin.helpers.gcom import get_active_instance_ids, get_deleted_instance_ids
+from apps.grafana_plugin.helpers.gcom import get_active_instance_ids, get_deleted_instance_ids, get_stack_regions
 from apps.user_management.models import Organization
+from apps.user_management.models.region import sync_regions
 from apps.user_management.sync import cleanup_organization, sync_organization
 from common.custom_celery_tasks import shared_dedicated_queue_retry_task
 
@@ -103,3 +104,16 @@ def start_cleanup_deleted_organizations():
 @shared_dedicated_queue_retry_task(autoretry_for=(Exception,), max_retries=1)
 def cleanup_organization_async(organization_pk):
     cleanup_organization(organization_pk)
+
+
+@shared_dedicated_queue_retry_task(autoretry_for=(Exception,), max_retries=1)
+def start_sync_regions():
+    regions, is_cloud_configured = get_stack_regions()
+    if not is_cloud_configured:
+        return
+
+    if not regions:
+        logger.warning("Did not find any stack-regions!")
+        return
+
+    sync_regions(regions)

engine/apps/grafana_plugin/views/self_hosted_install.py

@@ -46,6 +46,7 @@ class SelfHostedInstallView(GrafanaHeadersMixin, APIView):
                 org_id=org_id,
                 org_slug=settings.SELF_HOSTED_SETTINGS["ORG_SLUG"],
                 org_title=settings.SELF_HOSTED_SETTINGS["ORG_TITLE"],
+                region_slug=settings.SELF_HOSTED_SETTINGS["REGION_SLUG"],
                 grafana_url=self.instance_context["grafana_url"],
                 api_token=self.instance_context["grafana_token"],
             )

engine/apps/integrations/mixins/alert_channel_defining_mixin.py

@@ -7,6 +7,8 @@ from django.core.cache import cache
 from django.core.exceptions import PermissionDenied
 from django.db import OperationalError
 
+from apps.user_management.models.region import OrganizationMovedException
+
 logger = logging.getLogger(__name__)
 
 
@@ -64,6 +66,9 @@ class AlertChannelDefiningMixin(object):
                 logger.info("Cache is empty!")
                 raise
 
+        if alert_receive_channel.organization.is_moved:
+            raise OrganizationMovedException(alert_receive_channel.organization)
+
         del kwargs["alert_channel_key"]
         kwargs["alert_receive_channel"] = alert_receive_channel
 

engine/apps/user_management/middlewares.py

@@ -0,0 +1,46 @@
+import logging
+import re
+
+import requests
+from django.http import HttpResponse
+from django.utils.deprecation import MiddlewareMixin
+from rest_framework import status
+
+from apps.user_management.models.region import OrganizationMovedException
+from common.api_helpers.utils import create_engine_url
+
+logger = logging.getLogger(__name__)
+
+
+class OrganizationMovedMiddleware(MiddlewareMixin):
+    def process_exception(self, request, exception):
+        if isinstance(exception, OrganizationMovedException):
+            region = exception.organization.migration_destination
+            if not region.oncall_backend_url:
+                return HttpResponse(
+                    "Organization migration destination undefined URL", status=status.HTTP_500_INTERNAL_SERVER_ERROR
+                )
+
+            url = create_engine_url(request.path, override_base=region.oncall_backend_url)
+            if request.META["QUERY_STRING"]:
+                url = f"{url}?{request.META['QUERY_STRING']}"
+
+            regex = re.compile("^HTTP_")
+            headers = dict(
+                (regex.sub("", header), value) for (header, value) in request.META.items() if header.startswith("HTTP_")
+            )
+
+            response = self.make_request(request.method, url, headers, request.body)
+            return HttpResponse(response.content, status=response.status_code)
+
+    def make_request(self, method, url, headers, body):
+        if method == "GET":
+            return requests.get(url, headers=headers)
+        elif method == "POST":
+            return requests.post(url, data=body, headers=headers)
+        elif method == "PUT":
+            return requests.put(url, data=body, headers=headers)
+        elif method == "DELETE":
+            return requests.delete(url, headers=headers)
+        elif method == "OPTIONS":
+            return requests.options(url, headers=headers)

engine/apps/user_management/models/__init__.py

@@ -1,3 +1,4 @@
 from .user import User  # noqa: F401, isort: skip
 from .organization import Organization  # noqa: F401
+from .region import Region  # noqa: F401
 from .team import Team  # noqa: F401

engine/apps/user_management/models/organization.py

@@ -73,6 +73,16 @@ class Organization(MaintainableObject):
     stack_slug = models.CharField(max_length=300)
     org_slug = models.CharField(max_length=300)
     org_title = models.CharField(max_length=300)
+    region_slug = models.CharField(max_length=300, null=True, default=None)
+    migration_destination = models.ForeignKey(
+        to="user_management.Region",
+        to_field="slug",
+        db_column="migration_destination_slug",
+        on_delete=models.SET_NULL,
+        related_name="regions",
+        default=None,
+        null=True,
+    )
 
     grafana_url = models.URLField()
 
@@ -293,3 +303,7 @@ class Organization(MaintainableObject):
     @property
     def insight_logs_metadata(self):
         return {}
+
+    @property
+    def is_moved(self):
+        return self.migration_destination_id is not None

engine/apps/user_management/models/region.py

@@ -0,0 +1,52 @@
+import logging
+
+from django.apps import apps
+from django.db import models
+
+from apps.user_management.models import Organization
+
+logger = logging.getLogger(__name__)
+
+
+def sync_regions(regions: list[dict]):
+    Region = apps.get_model("user_management", "Region")
+    gcom_regions = {region["slug"]: region for region in regions}
+    existing_region_slugs = set(Region.objects.all().values_list("slug", flat=True))
+
+    # create new regions
+    regions_to_create = tuple(
+        Region(
+            name=region["name"],
+            slug=region["slug"],
+            oncall_backend_url=region["oncallApiUrl"],
+        )
+        for region in gcom_regions.values()
+        if region["slug"] not in existing_region_slugs
+    )
+    Region.objects.bulk_create(regions_to_create, batch_size=5000)
+
+    # delete excess regions
+    regions_to_delete = existing_region_slugs - gcom_regions.keys()
+    Region.objects.filter(slug__in=regions_to_delete).delete()
+
+    # update existing regions
+    regions_to_update = []
+    for region in Region.objects.filter(slug__in=existing_region_slugs):
+        gcom_region = gcom_regions[region.slug]
+        if region.name != gcom_region["name"] or region.oncall_backend_url != gcom_region["oncallApiUrl"]:
+            region.name = gcom_region["name"]
+            region.oncall_backend_url = gcom_region["oncallApiUrl"]
+            regions_to_update.append(region)
+
+    Region.objects.bulk_update(regions_to_update, ["name", "oncall_backend_url"], batch_size=5000)
+
+
+class OrganizationMovedException(Exception):
+    def __init__(self, organization: Organization):
+        self.organization = organization
+
+
+class Region(models.Model):
+    name = models.CharField(max_length=300)
+    slug = models.CharField(max_length=50, unique=True)
+    oncall_backend_url = models.URLField(null=True)

engine/apps/user_management/sync.py

@@ -29,6 +29,7 @@ def sync_organization(organization):
             "stack_slug",
             "org_slug",
             "org_title",
+            "region_slug",
             "grafana_url",
             "last_time_synced",
             "api_token_status",
@@ -47,6 +48,7 @@ def sync_instance_info(organization):
         organization.stack_slug = instance_info["slug"]
         organization.org_slug = instance_info["orgSlug"]
         organization.org_title = instance_info["orgName"]
+        organization.region_slug = instance_info["regionSlug"]
         organization.grafana_url = instance_info["url"]
         organization.gcom_token_org_last_time_synced = timezone.now()
 

engine/apps/user_management/tests/factories.py

@@ -1,6 +1,6 @@
 import factory
 
-from apps.user_management.models import Organization, Team, User
+from apps.user_management.models import Organization, Region, Team, User
 from common.utils import UniqueFaker
 
 
@@ -31,3 +31,12 @@ class TeamFactory(factory.DjangoModelFactory):
 
     class Meta:
         model = Team
+
+
+class RegionFactory(factory.DjangoModelFactory):
+    name = factory.Faker("country")
+    slug = factory.Faker("slug")
+    oncall_backend_url = factory.Faker("url")
+
+    class Meta:
+        model = Region

engine/apps/user_management/tests/test_region.py

@@ -0,0 +1,219 @@
+from unittest.mock import patch
+
+import pytest
+from django.http import HttpResponse
+from django.urls import reverse
+from rest_framework import status
+from rest_framework.test import APIClient
+
+from apps.alerts.models import AlertReceiveChannel
+from apps.auth_token.auth import ApiTokenAuthentication, ScheduleExportAuthentication, UserScheduleExportAuthentication
+from apps.auth_token.models import ScheduleExportAuthToken, UserScheduleExportAuthToken
+from apps.integrations.views import AlertManagerAPIView
+from apps.schedules.models import OnCallScheduleWeb
+from apps.user_management.models.region import OrganizationMovedException
+
+
+@pytest.mark.django_db
+def test_organization_region_delete(
+    make_organization_and_region,
+):
+    organization, region = make_organization_and_region()
+    organization.save()
+
+    organization.refresh_from_db()
+    assert organization.migration_destination.slug == region.slug
+    region.delete()
+
+    organization.refresh_from_db()
+    assert organization.migration_destination is None
+
+
+@pytest.mark.django_db
+def test_integration_does_not_raise_exception_organization_moved(
+    make_organization,
+    make_alert_receive_channel,
+):
+    organization = make_organization()
+    alert_receive_channel = make_alert_receive_channel(
+        organization=organization,
+        integration=AlertReceiveChannel.INTEGRATION_ALERTMANAGER,
+    )
+
+    try:
+        am = AlertManagerAPIView()
+        am.dispatch(alert_channel_key=alert_receive_channel.token)
+        assert False
+    except OrganizationMovedException:
+        assert False
+    except Exception:
+        assert True
+
+
+@pytest.mark.django_db
+def test_integration_raises_exception_organization_moved(
+    make_organization_and_region,
+    make_alert_receive_channel,
+):
+    organization, region = make_organization_and_region()
+    organization.save()
+
+    alert_receive_channel = make_alert_receive_channel(
+        organization=organization,
+        integration=AlertReceiveChannel.INTEGRATION_ALERTMANAGER,
+    )
+
+    try:
+        am = AlertManagerAPIView()
+        am.dispatch(alert_channel_key=alert_receive_channel.token)
+        assert False
+    except OrganizationMovedException as e:
+        assert e.organization == organization
+
+
+@patch("apps.user_management.middlewares.OrganizationMovedMiddleware.make_request")
+@pytest.mark.django_db
+def test_organization_moved_middleware(
+    mocked_make_request,
+    make_organization_and_region,
+    make_alert_receive_channel,
+):
+    organization, region = make_organization_and_region()
+    organization.save()
+
+    alert_receive_channel = make_alert_receive_channel(
+        organization=organization,
+        integration=AlertReceiveChannel.INTEGRATION_ALERTMANAGER,
+    )
+
+    expected_message = bytes(f"Redirected to {region.oncall_backend_url}", "utf-8")
+    mocked_make_request.return_value = HttpResponse(expected_message, status=status.HTTP_200_OK)
+
+    client = APIClient()
+    url = reverse("integrations:alertmanager", kwargs={"alert_channel_key": alert_receive_channel.token})
+
+    data = {"value": "test"}
+    response = client.post(url, data, format="json")
+    assert mocked_make_request.called
+    assert response.content == expected_message
+    assert response.status_code == status.HTTP_200_OK
+
+
+@pytest.mark.django_db
+def test_api_token_does_not_raise_exception_organization_moved(
+    make_organization,
+    make_user_for_organization,
+    make_public_api_token,
+):
+    organization = make_organization()
+
+    admin = make_user_for_organization(organization)
+    _, token = make_public_api_token(admin, organization)
+
+    try:
+        api_auth = ApiTokenAuthentication()
+        api_auth.authenticate_credentials(token)
+        assert True
+    except OrganizationMovedException:
+        assert False
+
+
+@pytest.mark.django_db
+def test_api_token_raises_exception_organization_moved(
+    make_organization_and_region,
+    make_user_for_organization,
+    make_public_api_token,
+):
+    organization, region = make_organization_and_region()
+    organization.save()
+
+    admin = make_user_for_organization(organization)
+    _, token = make_public_api_token(admin, organization)
+
+    try:
+        api_auth = ApiTokenAuthentication()
+        api_auth.authenticate_credentials(token)
+        assert False
+    except OrganizationMovedException as e:
+        assert e.organization == organization
+
+
+@pytest.mark.django_db
+def test_schedule_export_token_does_not_raise_exception_organization_moved(
+    make_organization,
+    make_user_for_organization,
+    make_public_api_token,
+    make_schedule,
+):
+    organization = make_organization()
+    schedule = make_schedule(organization, schedule_class=OnCallScheduleWeb)
+
+    admin = make_user_for_organization(organization)
+    _, token = ScheduleExportAuthToken.create_auth_token(admin, organization, schedule)
+
+    try:
+        schedule_auth = ScheduleExportAuthentication()
+        schedule_auth.authenticate_credentials(token, schedule.public_primary_key)
+        assert True
+    except OrganizationMovedException:
+        assert False
+
+
+@pytest.mark.django_db
+def test_schedule_export_token_raises_exception_organization_moved(
+    make_organization_and_region,
+    make_user_for_organization,
+    make_public_api_token,
+    make_schedule,
+):
+    organization, region = make_organization_and_region()
+    organization.save()
+    schedule = make_schedule(organization, schedule_class=OnCallScheduleWeb)
+
+    admin = make_user_for_organization(organization)
+    _, token = ScheduleExportAuthToken.create_auth_token(admin, organization, schedule)
+
+    try:
+        schedule_auth = ScheduleExportAuthentication()
+        schedule_auth.authenticate_credentials(token, schedule.public_primary_key)
+        assert False
+    except OrganizationMovedException as e:
+        assert e.organization == organization
+
+
+@pytest.mark.django_db
+def test_user_schedule_export_token_does_not_raise_exception_organization_moved(
+    make_organization,
+    make_user_for_organization,
+    make_public_api_token,
+):
+    organization = make_organization()
+    admin = make_user_for_organization(organization)
+    _, token = UserScheduleExportAuthToken.create_auth_token(admin, organization)
+
+    try:
+        user_schedule_auth = UserScheduleExportAuthentication()
+        user_schedule_auth.authenticate_credentials(token, admin.public_primary_key)
+        assert True
+    except OrganizationMovedException:
+        assert False
+
+
+@pytest.mark.django_db
+def test_user_schedule_export_token_raises_exception_organization_moved(
+    make_organization_and_region,
+    make_user_for_organization,
+    make_public_api_token,
+):
+    organization, region = make_organization_and_region()
+    organization.save()
+
+    admin = make_user_for_organization(organization)
+    _, token = UserScheduleExportAuthToken.create_auth_token(admin, organization)
+
+    try:
+        user_schedule_auth = UserScheduleExportAuthentication()
+        user_schedule_auth.authenticate_credentials(token, admin.public_primary_key)
+        assert False
+    except OrganizationMovedException as e:
+        assert e.organization == organization

engine/conftest.py

@@ -70,7 +70,7 @@ from apps.telegram.tests.factories import (
 )
 from apps.twilioapp.tests.factories import PhoneCallFactory, SMSFactory
 from apps.user_management.models.user import User, listen_for_user_model_save
-from apps.user_management.tests.factories import OrganizationFactory, TeamFactory, UserFactory
+from apps.user_management.tests.factories import OrganizationFactory, RegionFactory, TeamFactory, UserFactory
 from common.constants.role import Role
 
 register(OrganizationFactory)
@@ -666,3 +666,23 @@ def load_slack_urls(settings):
         reload(sys.modules[urlconf])
     else:
         import_module(urlconf)
+
+
+@pytest.fixture
+def make_region():
+    def _make_region(**kwargs):
+        region = RegionFactory(**kwargs)
+        return region
+
+    return _make_region
+
+
+@pytest.fixture
+def make_organization_and_region(make_organization, make_region):
+    def _make_organization_and_region():
+        organization = make_organization()
+        region = make_region()
+        organization.migration_destination = region
+        return organization, region
+
+    return _make_organization_and_region

engine/settings/base.py

@@ -244,6 +244,7 @@ MIDDLEWARE = [
     "social_django.middleware.SocialAuthExceptionMiddleware",
     "apps.social_auth.middlewares.SocialAuthAuthCanceledExceptionMiddleware",
     "apps.integrations.middlewares.IntegrationExceptionMiddleware",
+    "apps.user_management.middlewares.OrganizationMovedMiddleware",
 ]
 
 LOG_REQUEST_ID_HEADER = "HTTP_X_CLOUD_TRACE_CONTEXT"
@@ -558,6 +559,7 @@ SELF_HOSTED_SETTINGS = {
     "ORG_ID": 100,
     "ORG_SLUG": "self_hosted_org",
     "ORG_TITLE": "Self-Hosted Organization",
+    "REGION_SLUG": "self_hosted_region",
 }
 
 GRAFANA_INCIDENT_STATIC_API_KEY = os.environ.get("GRAFANA_INCIDENT_STATIC_API_KEY", None)