chore(hybrid-cloud): mypy fixes and organization improvement (#40862)

Moves organization and project_key services into separate directories,
adds hybrid cloud to mypy.
Part of a multi step extracting of
https://github.com/getsentry/sentry/pull/40772/files into pieces.

mypy.ini

@@ -105,6 +105,7 @@ files = fixtures/mypy-stubs,
         src/sentry/sentry_metrics/,
         src/sentry/services/http.py,
         src/sentry/services/base.py,
+        src/sentry/services/hybrid_cloud/,
         src/sentry/shared_integrations/,
         src/sentry/silo/,
         src/sentry/snuba/entity_subscription.py,

src/sentry/services/hybrid_cloud/__init__.py

@@ -1,58 +1,16 @@
 import contextlib
 import logging
 from abc import ABC, abstractmethod
-from dataclasses import dataclass
-from enum import Enum
-from typing import Dict, Generic, Iterable, List, Mapping, Optional, Type, TypeVar, cast
-
-from sentry.models import Organization, OrganizationMember, OrganizationStatus
+from typing import Any, Callable, Dict, Generator, Generic, Mapping, Optional, Type, TypeVar, cast
 
 logger = logging.getLogger(__name__)
 
-from django.db.models import F
-
 from sentry.silo import SiloMode
 
 
-class ProjectKeyRole(Enum):
-    store = "store"
-    api = "api"
-
-    def as_orm_role(self):
-        from sentry.models import ProjectKey
-
-        if self == ProjectKeyRole.store:
-            return ProjectKey.roles.store
-        elif self == ProjectKeyRole.api:
-            return ProjectKey.roles.api
-        else:
-            raise ValueError("Unexpected project key role enum")
-
-
-@dataclass
-class ApiProjectKey:
-    dsn_public: str = ""
-
-
-@dataclass
-class ApiOrganizationMember:
-    # This can be null when the user is deleted.
-    user_id: Optional[int]
-    pass
-
-
-@dataclass
-class ApiOrganization:
-    slug: str = ""
-    id: int = -1
-    # exists iff the organization was queried with a user_id context, and that user_id
-    # was confirmed to be a member.
-    member: Optional[ApiOrganizationMember] = None
-
-
 class InterfaceWithLifecycle(ABC):
     @abstractmethod
-    def close(self):
+    def close(self) -> None:
         pass
 
 
@@ -70,15 +28,17 @@ class DelegatedBySiloMode(Generic[ServiceInterface]):
     service is closed, or when the backing service implementation changes.
     """
 
-    _constructors: Mapping[SiloMode, Type[ServiceInterface]]
+    _constructors: Mapping[SiloMode, Callable[[], ServiceInterface]]
     _singleton: Dict[SiloMode, ServiceInterface]
 
-    def __init__(self, mapping: Mapping[SiloMode, Type[ServiceInterface]]):
+    def __init__(self, mapping: Mapping[SiloMode, Callable[[], ServiceInterface]]):
         self._constructors = mapping
         self._singleton = {}
 
     @contextlib.contextmanager
-    def with_replacement(self, service: Optional[ServiceInterface], silo_mode: SiloMode):
+    def with_replacement(
+        self, service: Optional[ServiceInterface], silo_mode: SiloMode
+    ) -> Generator[None, None, None]:
         prev = self._singleton
         self.close()
 
@@ -90,163 +50,22 @@ class DelegatedBySiloMode(Generic[ServiceInterface]):
         self.close()
         self._singleton = prev
 
-    def __getattr__(self, item: str):
+    def __getattr__(self, item: str) -> Any:
         cur_mode = SiloMode.get_current_mode()
         if impl := self._singleton.get(cur_mode, None):
             return getattr(impl, item)
-        if Con := self._constructors.get(cur_mode, None):
+        if con := self._constructors.get(cur_mode, None):
             self.close()
-            return getattr(self._singleton.setdefault(cur_mode, Con()), item)
+            return getattr(self._singleton.setdefault(cur_mode, con()), item)
 
         raise KeyError(f"No implementation found for {cur_mode}.")
 
-    def close(self):
+    def close(self) -> None:
         for impl in self._singleton.values():
             impl.close()
         self._singleton = {}
 
 
-class ProjectKeyService(InterfaceWithLifecycle):
-    @abstractmethod
-    def get_project_key(self, project_id: str, role: ProjectKeyRole) -> Optional[ApiProjectKey]:
-        pass
-
-
-class OrganizationService(InterfaceWithLifecycle):
-    @abstractmethod
-    def get_organizations(
-        self, user_id: Optional[int], scope: Optional[str], only_visible: bool
-    ) -> List[ApiOrganization]:
-        """
-        This method is expected to follow the optionally given user_id, scope, and only_visible options to filter
-        an appropriate set.
-        :param user_id:
-        When null, this should imply the entire set of organizations, not bound by user.  Be certain to authenticate
-        users before returning this.
-        :param scope:
-        :param only_visible:
-        :return:
-        """
-        pass
-
-    @abstractmethod
-    def check_membership_by_email(
-        self, organization_id: int, email: str
-    ) -> Optional[ApiOrganizationMember]:
-        """
-        Used to look up an organization membership by an email, used in very specific edge cases.
-        """
-        pass
-
-    @abstractmethod
-    def get_organization_by_slug(
-        self, *, user_id: Optional[int], slug: str, only_visible: bool, allow_stale: bool
-    ) -> Optional[ApiOrganization]:
-        pass
-
-    def _serialize_member(self, member: OrganizationMember) -> ApiOrganizationMember:
-        return ApiOrganizationMember(user_id=member.user.id if member.user is not None else None)
-
-    def _serialize_organization(
-        self, org: Organization, membership: Iterable[OrganizationMember] = tuple()
-    ) -> ApiOrganization:
-        org = ApiOrganization(slug=org.slug, id=org.id)
-
-        for member in membership:
-            if member.organization.id == org.id:
-                org.member = self._serialize_member(member)
-                break
-
-        return org
-
-
-class DatabaseBackedOrganizationService(OrganizationService):
-    def check_membership_by_email(
-        self, organization_id: int, email: str
-    ) -> Optional[ApiOrganizationMember]:
-        try:
-            member = OrganizationMember.objects.get(organization_id=organization_id, email=email)
-        except OrganizationMember.DoesNotExist:
-            return None
-
-        return self._serialize_member(member)
-
-    def get_organization_by_slug(
-        self, *, user_id: Optional[int], slug: str, only_visible: bool, allow_stale: bool
-    ) -> Optional[ApiOrganization]:
-        membership: List[OrganizationMember]
-        if user_id is not None:
-            membership = OrganizationMember.objects.filter(user_id=user_id)
-        else:
-            membership = []
-        try:
-            if allow_stale:
-                org = Organization.objects.get_from_cache(slug=slug)
-            else:
-                org = Organization.objects.get(slug=slug)
-
-            if only_visible and org.status != OrganizationStatus.VISIBLE:
-                raise Organization.DoesNotExist
-            return self._serialize_organization(org, membership)
-        except Organization.DoesNotExist:
-            logger.info("Active organization [%s] not found", slug)
-
-        return None
-
-    def close(self):
-        pass
-
-    def get_organizations(
-        self, user_id: Optional[int], scope: Optional[str], only_visible: bool
-    ) -> List[ApiOrganization]:
-        if user_id is None:
-            return []
-        organizations = self._query_organizations(user_id, scope, only_visible)
-        membership = OrganizationMember.objects.filter(user_id=user_id)
-        return [self._serialize_organization(o, membership) for o in organizations]
-
-    def _query_organizations(
-        self, user_id: int, scope: Optional[str], only_visible: bool
-    ) -> List[Organization]:
-        from django.conf import settings
-
-        if settings.SENTRY_PUBLIC and scope is None:
-            if only_visible:
-                return list(Organization.objects.filter(status=OrganizationStatus.ACTIVE))
-            else:
-                return list(Organization.objects.filter())
-
-        qs = OrganizationMember.objects.filter(user_id=user_id)
-
-        qs = qs.select_related("organization")
-        if only_visible:
-            qs = qs.filter(organization__status=OrganizationStatus.ACTIVE)
-
-        results = list(qs)
-
-        if scope is not None:
-            return [r.organization for r in results if scope in r.get_scopes()]
-
-        return [r.organization for r in results]
-
-
-class DatabaseBackedProjectKeyService(ProjectKeyService):
-    def close(self):
-        pass
-
-    def get_project_key(self, project_id: str, role: ProjectKeyRole) -> Optional[ApiProjectKey]:
-        from sentry.models import ProjectKey
-
-        project_keys = ProjectKey.objects.filter(
-            project=project_id, roles=F("roles").bitor(role.as_orm_role())
-        )
-
-        if project_keys:
-            return ApiProjectKey(dsn_public=project_keys[0].dsn_public)
-
-        return None
-
-
 def CreateStubFromBase(base: Type[ServiceInterface]) -> Type[ServiceInterface]:
     """
     Using a concrete implementation class of a service, creates a new concrete implementation class suitable for a test
@@ -259,15 +78,14 @@ def CreateStubFromBase(base: Type[ServiceInterface]) -> Type[ServiceInterface]:
     """
     Super = base.__bases__[0]
 
-    def __init__(self, *args, **kwds):
-        Super.__init__(self, *args, **kwds)
+    def __init__(self: Any, *args: Any, **kwds: Any) -> None:
         self.backing_service = base(*args, **kwds)
 
-    def close(self):
+    def close(self: Any) -> None:
         self.backing_service.close()
 
-    def make_method(method_name: str):
-        def method(self, *args, **kwds):
+    def make_method(method_name: str) -> Any:
+        def method(self: Any, *args: Any, **kwds: Any) -> Any:
             from sentry.testutils.silo import exempt_from_silo_limits
 
             with exempt_from_silo_limits():
@@ -287,16 +105,14 @@ def CreateStubFromBase(base: Type[ServiceInterface]) -> Type[ServiceInterface]:
     return cast(Type[ServiceInterface], type(f"Stub{Super.__name__}", (Super,), methods))
 
 
-StubProjectKeyService = CreateStubFromBase(DatabaseBackedProjectKeyService)
-StubOrganizationService = CreateStubFromBase(DatabaseBackedOrganizationService)
-
-
-def silo_mode_delegation(mapping: Mapping[SiloMode, Type[ServiceInterface]]) -> ServiceInterface:
+def silo_mode_delegation(
+    mapping: Mapping[SiloMode, Callable[[], ServiceInterface]]
+) -> ServiceInterface:
     """
     Simply creates a DelegatedBySiloMode from a mapping object, but casts it as a ServiceInterface matching
     the mapping values.
     """
-    return cast(InterfaceWithLifecycle, DelegatedBySiloMode(mapping))
+    return cast(ServiceInterface, DelegatedBySiloMode(mapping))
 
 
 @contextlib.contextmanager
@@ -304,7 +120,7 @@ def service_stubbed(
     service: InterfaceWithLifecycle,
     stub: Optional[InterfaceWithLifecycle],
     silo_mode: Optional[SiloMode] = None,
-):
+) -> Generator[None, None, None]:
     """
     Replaces a service created with silo_mode_delegation with a replacement implementation while inside of the scope,
     closing the existing implementation on enter and closing the given implementation on exit.
@@ -320,52 +136,18 @@ def service_stubbed(
 
 
 @contextlib.contextmanager
-def use_real_service(service: InterfaceWithLifecycle, silo_mode: SiloMode):
+def use_real_service(
+    service: InterfaceWithLifecycle, silo_mode: SiloMode
+) -> Generator[None, None, None]:
     """
     Removes any stubbed implementations, forcing the default configured implementation.
     Important for integration tests that validate the integration of production service implementations.
     """
     from django.test import override_settings
 
-    if silo_mode is None:
-        silo_mode = SiloMode.get_current_mode()
-
     if isinstance(service, DelegatedBySiloMode):
         with override_settings(SILO_MODE=silo_mode):
             with service.with_replacement(None, silo_mode):
                 yield
     else:
         raise ValueError("Service needs to be a DelegatedBySiloMode object, but it was not!")
-
-
-class RegionClientBackedProjectKeyService(ProjectKeyService):
-    def get_project_key(self, project_id: str, role: ProjectKeyRole) -> Optional[ApiProjectKey]:
-        pass
-
-    def close(self):
-        pass
-
-
-class RegionClientBackedOrganizationService(ProjectKeyService):
-    def get_project_key(self, project_id: str, role: ProjectKeyRole) -> Optional[ApiProjectKey]:
-        pass
-
-    def close(self):
-        pass
-
-
-project_key_service: ProjectKeyService = silo_mode_delegation(
-    {
-        SiloMode.MONOLITH: DatabaseBackedProjectKeyService,
-        SiloMode.REGION: DatabaseBackedProjectKeyService,
-        SiloMode.CONTROL: StubProjectKeyService,
-    }
-)
-
-organization_service: OrganizationService = silo_mode_delegation(
-    {
-        SiloMode.MONOLITH: DatabaseBackedOrganizationService,
-        SiloMode.REGION: DatabaseBackedOrganizationService,
-        SiloMode.CONTROL: StubOrganizationService,
-    }
-)

src/sentry/services/hybrid_cloud/organization.py

@@ -0,0 +1,161 @@
+from abc import abstractmethod
+from dataclasses import dataclass
+from typing import Iterable, List, Optional
+
+from sentry.models import Organization, OrganizationMember, OrganizationStatus
+from sentry.services.hybrid_cloud import (
+    CreateStubFromBase,
+    InterfaceWithLifecycle,
+    logger,
+    silo_mode_delegation,
+)
+from sentry.silo import SiloMode
+
+
+@dataclass
+class ApiOrganizationMember:
+    # This can be null when the user is deleted.
+    user_id: Optional[int]
+    pass
+
+
+@dataclass
+class ApiOrganization:
+    slug: str = ""
+    id: int = -1
+    # exists iff the organization was queried with a user_id context, and that user_id
+    # was confirmed to be a member.
+    member: Optional[ApiOrganizationMember] = None
+
+
+class OrganizationService(InterfaceWithLifecycle):
+    @abstractmethod
+    def get_organizations(
+        self, user_id: Optional[int], scope: Optional[str], only_visible: bool
+    ) -> List[ApiOrganization]:
+        """
+        This method is expected to follow the optionally given user_id, scope, and only_visible options to filter
+        an appropriate set.
+        :param user_id:
+        When null, this should imply the entire set of organizations, not bound by user.  Be certain to authenticate
+        users before returning this.
+        :param scope:
+        :param only_visible:
+        :return:
+        """
+        pass
+
+    @abstractmethod
+    def check_membership_by_email(
+        self, organization_id: int, email: str
+    ) -> Optional[ApiOrganizationMember]:
+        """
+        Used to look up an organization membership by an email, used in very specific edge cases.
+        """
+        pass
+
+    @abstractmethod
+    def get_organization_by_slug(
+        self, *, user_id: Optional[int], slug: str, only_visible: bool, allow_stale: bool
+    ) -> Optional[ApiOrganization]:
+        pass
+
+    def _serialize_member(self, member: OrganizationMember) -> ApiOrganizationMember:
+        return ApiOrganizationMember(user_id=member.user.id if member.user is not None else None)
+
+    def _serialize_organization(
+        self, org: Organization, user_memberships: Iterable[OrganizationMember] = tuple()
+    ) -> ApiOrganization:
+        result = ApiOrganization(slug=org.slug, id=org.id)
+
+        for member in user_memberships:
+            if member.organization.id == org.id:
+                result.member = self._serialize_member(member)
+                break
+
+        return result
+
+
+class DatabaseBackedOrganizationService(OrganizationService):
+    def check_membership_by_email(
+        self, organization_id: int, email: str
+    ) -> Optional[ApiOrganizationMember]:
+        try:
+            member = OrganizationMember.objects.get(organization_id=organization_id, email=email)
+        except OrganizationMember.DoesNotExist:
+            return None
+
+        return self._serialize_member(member)
+
+    def get_organization_by_slug(
+        self, *, user_id: Optional[int], slug: str, only_visible: bool, allow_stale: bool
+    ) -> Optional[ApiOrganization]:
+        user_memberships: List[OrganizationMember] = []
+        try:
+            if allow_stale:
+                org = Organization.objects.get_from_cache(slug=slug)
+            else:
+                org = Organization.objects.get(slug=slug)
+
+            if user_id is not None:
+                try:
+                    user_memberships = [
+                        OrganizationMember.objects.get(organization_id=org.id, user_id=user_id)
+                    ]
+                except OrganizationMember.DoesNotExist:
+                    pass
+
+            if only_visible and org.status != OrganizationStatus.VISIBLE:
+                raise Organization.DoesNotExist
+            return self._serialize_organization(org, user_memberships)
+        except Organization.DoesNotExist:
+            logger.info("Active organization [%s] not found", slug)
+
+        return None
+
+    def close(self) -> None:
+        pass
+
+    def get_organizations(
+        self, user_id: Optional[int], scope: Optional[str], only_visible: bool
+    ) -> List[ApiOrganization]:
+        if user_id is None:
+            return []
+        organizations = self._query_organizations(user_id, scope, only_visible)
+        membership = OrganizationMember.objects.filter(user_id=user_id)
+        return [self._serialize_organization(o, membership) for o in organizations]
+
+    def _query_organizations(
+        self, user_id: int, scope: Optional[str], only_visible: bool
+    ) -> List[Organization]:
+        from django.conf import settings
+
+        if settings.SENTRY_PUBLIC and scope is None:
+            if only_visible:
+                return list(Organization.objects.filter(status=OrganizationStatus.ACTIVE))
+            else:
+                return list(Organization.objects.filter())
+
+        qs = OrganizationMember.objects.filter(user_id=user_id)
+
+        qs = qs.select_related("organization")
+        if only_visible:
+            qs = qs.filter(organization__status=OrganizationStatus.ACTIVE)
+
+        results = list(qs)
+
+        if scope is not None:
+            return [r.organization for r in results if scope in r.get_scopes()]
+
+        return [r.organization for r in results]
+
+
+StubOrganizationService = CreateStubFromBase(DatabaseBackedOrganizationService)
+
+organization_service: OrganizationService = silo_mode_delegation(
+    {
+        SiloMode.MONOLITH: lambda: DatabaseBackedOrganizationService(),
+        SiloMode.REGION: lambda: DatabaseBackedOrganizationService(),
+        SiloMode.CONTROL: lambda: StubOrganizationService(),
+    }
+)

src/sentry/services/hybrid_cloud/project_key.py

@@ -0,0 +1,68 @@
+from abc import abstractmethod
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any, Optional
+
+from django.db.models import F
+
+from sentry.services.hybrid_cloud import (
+    CreateStubFromBase,
+    InterfaceWithLifecycle,
+    silo_mode_delegation,
+)
+from sentry.silo import SiloMode
+
+
+class ProjectKeyRole(Enum):
+    store = "store"
+    api = "api"
+
+    def as_orm_role(self) -> Any:
+        from sentry.models import ProjectKey
+
+        if self == ProjectKeyRole.store:
+            return ProjectKey.roles.store
+        elif self == ProjectKeyRole.api:
+            return ProjectKey.roles.api
+        else:
+            raise ValueError("Unexpected project key role enum")
+
+
+@dataclass
+class ApiProjectKey:
+    dsn_public: str = ""
+
+
+class ProjectKeyService(InterfaceWithLifecycle):
+    @abstractmethod
+    def get_project_key(self, project_id: str, role: ProjectKeyRole) -> Optional[ApiProjectKey]:
+        pass
+
+
+class DatabaseBackedProjectKeyService(ProjectKeyService):
+    def close(self) -> None:
+        pass
+
+    def get_project_key(self, project_id: str, role: ProjectKeyRole) -> Optional[ApiProjectKey]:
+        from sentry.models import ProjectKey
+
+        project_keys = ProjectKey.objects.filter(
+            project=project_id, roles=F("roles").bitor(role.as_orm_role())
+        )
+
+        if project_keys:
+            return ApiProjectKey(dsn_public=project_keys[0].dsn_public)
+
+        return None
+
+
+StubProjectKeyService = CreateStubFromBase(DatabaseBackedProjectKeyService)
+
+
+project_key_service: ProjectKeyService = silo_mode_delegation(
+    {
+        SiloMode.MONOLITH: lambda: DatabaseBackedProjectKeyService(),
+        SiloMode.REGION: lambda: DatabaseBackedProjectKeyService(),
+        SiloMode.CONTROL: lambda: StubProjectKeyService(),
+    }
+)

src/sentry/utils/auth.py

@@ -16,7 +16,7 @@ from django.utils.http import is_safe_url
 from rest_framework.request import Request
 
 from sentry.models import Authenticator, Organization, User
-from sentry.services.hybrid_cloud import ApiOrganization
+from sentry.services.hybrid_cloud.organization import ApiOrganization
 from sentry.utils import metrics
 from sentry.utils.http import absolute_uri
 

src/sentry/web/client_config.py

@@ -13,7 +13,7 @@ from sentry.auth import superuser
 from sentry.auth.access import get_cached_organization_member
 from sentry.auth.superuser import is_active_superuser
 from sentry.models import Organization, OrganizationMember
-from sentry.services.hybrid_cloud import ProjectKeyRole, project_key_service
+from sentry.services.hybrid_cloud.project_key import ProjectKeyRole, project_key_service
 from sentry.utils import auth
 from sentry.utils.assets import get_frontend_app_asset_url
 from sentry.utils.email import is_smtp_enabled

src/sentry/web/frontend/auth_login.py

@@ -16,7 +16,7 @@ from sentry.auth.superuser import is_active_superuser
 from sentry.constants import WARN_SESSION_EXPIRED
 from sentry.http import get_server_hostname
 from sentry.models import AuthProvider, Organization, OrganizationMember, OrganizationStatus
-from sentry.services.hybrid_cloud import organization_service
+from sentry.services.hybrid_cloud.organization import organization_service
 from sentry.signals import join_request_link_viewed, user_signup
 from sentry.utils import auth, json, metrics
 from sentry.utils.auth import (

src/sentry/web/frontend/base.py

@@ -21,7 +21,7 @@ from sentry.api.utils import is_member_disabled_from_limit
 from sentry.auth import access
 from sentry.auth.superuser import is_active_superuser
 from sentry.models import Authenticator, Organization, Project, ProjectStatus, Team, TeamStatus
-from sentry.services.hybrid_cloud import ApiOrganization, organization_service
+from sentry.services.hybrid_cloud.organization import ApiOrganization, organization_service
 from sentry.silo import SiloMode
 from sentry.utils import auth
 from sentry.utils.audit import create_audit_entry

src/sentry/web/frontend/restore_organization.py

@@ -8,7 +8,7 @@ from django.utils.translation import ugettext_lazy as _
 from sentry import audit_log
 from sentry.api import client
 from sentry.models import Organization, OrganizationStatus
-from sentry.services.hybrid_cloud import organization_service
+from sentry.services.hybrid_cloud.organization import organization_service
 from sentry.web.frontend.base import OrganizationView
 from sentry.web.helpers import render_to_response
 

tests/conftest.py

@@ -129,13 +129,12 @@ def setup_default_hybrid_cloud_stubs():
         MockRegionToControlMessageService,
         region_to_control_message_service,
     )
-    from sentry.services.hybrid_cloud import (
+    from sentry.services.hybrid_cloud import service_stubbed
+    from sentry.services.hybrid_cloud.organization import (
         StubOrganizationService,
-        StubProjectKeyService,
         organization_service,
-        project_key_service,
-        service_stubbed,
     )
+    from sentry.services.hybrid_cloud.project_key import StubProjectKeyService, project_key_service
     from sentry.silo import SiloMode
 
     stubs = [