Skip to content

Commit

Permalink
ref: fix some typing in hybrid cloud modules
Browse files Browse the repository at this point in the history
fixes most errors in these modules when BaseManager becomes typed
  • Loading branch information
asottile-sentry committed Jun 12, 2024
1 parent 7c76075 commit 6f0084c
Show file tree
Hide file tree
Showing 13 changed files with 33 additions and 27 deletions.
14 changes: 9 additions & 5 deletions src/sentry/db/models/outboxes.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@

import contextlib
import logging
from collections.abc import Collection, Iterable, Mapping
from collections.abc import Collection, Iterable, Mapping, Sequence
from typing import TYPE_CHECKING, Any, Protocol, TypeVar

from django.db import connections, router, transaction
Expand Down Expand Up @@ -80,7 +80,7 @@ class RegionOutboxProducingManager(BaseManager[_RM]):
Provides bulk update and delete methods that respect outbox creation.
"""

def bulk_create(self, objs: Iterable[_RM], *args: Any, **kwds: Any) -> Collection[_RM]:
def bulk_create(self, objs: Iterable[_RM], *args: Any, **kwds: Any) -> list[_RM]:
from sentry.models.outbox import outbox_context

tuple_of_objs: tuple[_RM, ...] = tuple(objs)
Expand All @@ -107,7 +107,9 @@ def bulk_create(self, objs: Iterable[_RM], *args: Any, **kwds: Any) -> Collectio
type(outboxes[0]).objects.bulk_create(outboxes)
return super().bulk_create(tuple_of_objs, *args, **kwds)

def bulk_update(self, objs: Iterable[_RM], fields: list[str], *args: Any, **kwds: Any) -> Any:
def bulk_update(
self, objs: Iterable[_RM], fields: Sequence[str], *args: Any, **kwds: Any
) -> Any:
from sentry.models.outbox import outbox_context

tuple_of_objs: tuple[_RM, ...] = tuple(objs)
Expand Down Expand Up @@ -262,7 +264,7 @@ class ControlOutboxProducingManager(BaseManager[_CM]):
Provides bulk update and delete methods that respect outbox creation.
"""

def bulk_create(self, objs: Iterable[_CM], *args: Any, **kwds: Any) -> Collection[_CM]:
def bulk_create(self, objs: Iterable[_CM], *args: Any, **kwds: Any) -> list[_CM]:
from sentry.models.outbox import outbox_context

tuple_of_objs: tuple[_CM, ...] = tuple(objs)
Expand All @@ -289,7 +291,9 @@ def bulk_create(self, objs: Iterable[_CM], *args: Any, **kwds: Any) -> Collectio
type(outboxes[0]).objects.bulk_create(outboxes)
return super().bulk_create(tuple_of_objs, *args, **kwds)

def bulk_update(self, objs: Iterable[_CM], fields: list[str], *args: Any, **kwds: Any) -> Any:
def bulk_update(
self, objs: Iterable[_CM], fields: Sequence[str], *args: Any, **kwds: Any
) -> Any:
from sentry.models.outbox import outbox_context

tuple_of_objs: tuple[_CM, ...] = tuple(objs)
Expand Down
4 changes: 2 additions & 2 deletions src/sentry/hybridcloud/models/webhookpayload.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
from __future__ import annotations

import datetime
from typing import Any, Self
from typing import Any

from django.db import models
from django.http import HttpRequest
Expand Down Expand Up @@ -79,7 +79,7 @@ def create_from_request(
identifier: int | str,
request: HttpRequest,
integration_id: int | None = None,
) -> Self:
) -> WebhookPayload:
metrics.incr("hybridcloud.deliver_webhooks.saved")
return cls.objects.create(
mailbox_name=f"{provider}:{identifier}",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -90,7 +90,7 @@ def _get_previously_provisioned_org_and_validate(
if matching_organizations_qs.count() > 1:
raise PreProvisionCheckException("Multiple conflicting organizations found")

matching_org: Organization = matching_organizations_qs.first()
matching_org = matching_organizations_qs.get()

try:
provisioning_user_is_org_owner = (
Expand Down
8 changes: 4 additions & 4 deletions src/sentry/receivers/outbox/control.py
Original file line number Diff line number Diff line change
Expand Up @@ -57,8 +57,8 @@ def process_sentry_app_updates(object_identifier: int, region_name: str, **kwds:
)
# There isn't a constraint on org : sentryapp so we have to handle lists
install_map: dict[int, list[int]] = defaultdict(list)
for row in install_query:
install_map[row["organization_id"]].append(row["id"])
for install_row in install_query:
install_map[install_row["organization_id"]].append(install_row["id"])

# Clear application_id cache
region_caching_service.clear_key(
Expand All @@ -70,8 +70,8 @@ def process_sentry_app_updates(object_identifier: int, region_name: str, **kwds:
region_query = OrganizationMapping.objects.filter(
organization_id__in=list(install_map.keys()), region_name=region_name
).values("organization_id")
for row in region_query:
installs = install_map[row["organization_id"]]
for region_row in region_query:
installs = install_map[region_row["organization_id"]]
for install_id in installs:
region_caching_service.clear_key(
key=get_installation.key_from(install_id), region_name=region_name
Expand Down
2 changes: 1 addition & 1 deletion src/sentry/services/hybrid_cloud/app/impl.py
Original file line number Diff line number Diff line change
Expand Up @@ -166,7 +166,7 @@ class _AppServiceFilterQuery(
):
def base_query(self, select_related: bool = True) -> QuerySet[SentryAppInstallation]:
if not select_related:
return SentryAppInstallation.objects
return SentryAppInstallation.objects.all()
return SentryAppInstallation.objects.select_related("sentry_app")

def filter_arg_validator(
Expand Down
2 changes: 1 addition & 1 deletion src/sentry/services/hybrid_cloud/identity/impl.py
Original file line number Diff line number Diff line change
Expand Up @@ -109,7 +109,7 @@ def apply_filters(
return query

def base_query(self, select_related: bool = True) -> QuerySet[Identity]:
return Identity.objects
return Identity.objects.all()

def filter_arg_validator(self) -> Callable[[IdentityFilterArgs], str | None]:
return self._filter_has_any_key_validator(*IdentityFilterArgs.__annotations__.keys())
Expand Down
2 changes: 1 addition & 1 deletion src/sentry/services/hybrid_cloud/integration/impl.py
Original file line number Diff line number Diff line change
Expand Up @@ -471,7 +471,7 @@ def get_integration_external_projects(
except OrganizationIntegration.DoesNotExist:
return []

iep_kwargs = {"organization_integration_id": oi.id}
iep_kwargs: dict[str, Any] = {"organization_integration_id": oi.id}
if external_id is not None:
iep_kwargs["external_id"] = external_id
external_projects = IntegrationExternalProject.objects.filter(**iep_kwargs)
Expand Down
4 changes: 2 additions & 2 deletions src/sentry/services/hybrid_cloud/organization/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
# from __future__ import annotations
# in modules such as this one where hybrid cloud data models or service classes are
# defined, because we want to reflect on type annotations and avoid forward references.
from collections.abc import Mapping, Sequence
from collections.abc import Iterable, Mapping, Sequence
from datetime import datetime
from enum import IntEnum
from typing import Any, TypedDict
Expand Down Expand Up @@ -255,7 +255,7 @@ def get_owners(self) -> Sequence[RpcUser]:
from sentry.services.hybrid_cloud.user.service import user_service

if SiloMode.get_current_mode() == SiloMode.CONTROL:
owners = OrganizationMemberMapping.objects.filter(
owners: Iterable[int | None] = OrganizationMemberMapping.objects.filter(
organization_id=self.id, role__in=[roles.get_top_dog().id]
).values_list("user_id", flat=True)
else:
Expand Down
4 changes: 2 additions & 2 deletions src/sentry/services/hybrid_cloud/organization/serial.py
Original file line number Diff line number Diff line change
Expand Up @@ -142,10 +142,10 @@ def serialize_rpc_organization(
)

if include_projects:
projects: list[Project] = Project.objects.filter(organization=org)
projects = Project.objects.filter(organization=org)
rpc_org.projects.extend(serialize_project(project) for project in projects)
if include_teams:
teams: list[Team] = Team.objects.filter(organization=org)
teams = Team.objects.filter(organization=org)
rpc_org.teams.extend(serialize_rpc_team(team) for team in teams)

return rpc_org
Expand Down
10 changes: 5 additions & 5 deletions src/sentry/services/hybrid_cloud/organization_mapping/impl.py
Original file line number Diff line number Diff line change
Expand Up @@ -88,18 +88,18 @@ def _check_organization_mapping_integrity(
def _upsert_organization_slug_reservation_for_monolith(
self, organization_id: int, mapping_update: RpcOrganizationMappingUpdate
) -> None:
org_slug_reservation_qs = OrganizationSlugReservation.objects.filter(
org_slug_reservation = OrganizationSlugReservation.objects.filter(
organization_id=organization_id
)
if not org_slug_reservation_qs.exists():
).first()
if org_slug_reservation is None:
OrganizationSlugReservation(
region_name=mapping_update.region_name,
slug=mapping_update.slug,
organization_id=organization_id,
user_id=-1,
).save(unsafe_write=True)
elif org_slug_reservation_qs.first().slug != mapping_update.slug:
org_slug_reservation_qs.first().update(slug=mapping_update.slug, unsafe_write=True)
elif org_slug_reservation.slug != mapping_update.slug:
org_slug_reservation.update(slug=mapping_update.slug, unsafe_write=True)

def upsert(self, organization_id: int, update: RpcOrganizationMappingUpdate) -> None:
update_dict: dict[str, Any] = dict(
Expand Down
4 changes: 3 additions & 1 deletion src/sentry/services/hybrid_cloud/project/impl.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,9 @@
class DatabaseBackedProjectService(ProjectService):
def get_by_id(self, *, organization_id: int, id: int) -> RpcProject | None:
try:
project = Project.objects.get_from_cache(id=id, organization=organization_id)
project: Project | None = Project.objects.get_from_cache(
id=id, organization=organization_id
)
except ValueError:
project = Project.objects.filter(id=id, organization=organization_id).first()
except Project.DoesNotExist:
Expand Down
2 changes: 1 addition & 1 deletion src/sentry/services/hybrid_cloud/user/impl.py
Original file line number Diff line number Diff line change
Expand Up @@ -333,7 +333,7 @@ def apply_filters(self, query: QuerySet[User], filters: UserFilterArgs) -> Query

def base_query(self, select_related: bool = True) -> QuerySet[User]:
if not select_related:
return User.objects
return User.objects.all()

return User.objects.extra(
select={
Expand Down
2 changes: 1 addition & 1 deletion src/sentry/services/hybrid_cloud/user_option/impl.py
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,7 @@ class _UserOptionFilterQuery(
FilterQueryDatabaseImpl[UserOption, UserOptionFilterArgs, RpcUserOption, None]
):
def base_query(self, select_related: bool = True) -> QuerySet[UserOption]:
return UserOption.objects
return UserOption.objects.all()

def filter_arg_validator(self) -> Callable[[UserOptionFilterArgs], str | None]:
return self._filter_has_any_key_validator("user_ids")
Expand Down

0 comments on commit 6f0084c

Please sign in to comment.