Skip to content

Commit

Permalink
ref: get proper base manager typing
Browse files Browse the repository at this point in the history
  • Loading branch information
asottile-sentry committed Jul 29, 2024
1 parent 3dc95d7 commit 0d84170
Show file tree
Hide file tree
Showing 9 changed files with 23 additions and 19 deletions.
2 changes: 1 addition & 1 deletion requirements-dev-frozen.txt
Original file line number Diff line number Diff line change
Expand Up @@ -180,7 +180,7 @@ selenium==4.16.0
sentry-arroyo==2.16.5
sentry-cli==2.16.0
sentry-devenv==1.7.0
sentry-forked-django-stubs==5.0.2.post7
sentry-forked-django-stubs==5.0.2.post8
sentry-forked-djangorestframework-stubs==3.15.0.post1
sentry-kafka-schemas==0.1.101
sentry-ophio==0.2.7
Expand Down
2 changes: 1 addition & 1 deletion requirements-dev.txt
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@ pip-tools>=7.1.0
packaging>=21.3

# for type checking
sentry-forked-django-stubs>=5.0.2.post7
sentry-forked-django-stubs>=5.0.2.post8
sentry-forked-djangorestframework-stubs>=3.15.0.post1
lxml-stubs
msgpack-types>=0.2.0
Expand Down
2 changes: 1 addition & 1 deletion src/sentry/backup/imports.py
Original file line number Diff line number Diff line change
Expand Up @@ -72,7 +72,7 @@ def _clear_model_tables_before_import():
for model in reversed:
using = router.db_for_write(model)
manager = model.with_deleted if issubclass(model, ParanoidModel) else model.objects
manager.all().delete() # type: ignore[attr-defined]
manager.all().delete()

# TODO(getsentry/team-ospo#190): Remove the "Node" kludge below in favor of a more permanent
# solution.
Expand Down
9 changes: 6 additions & 3 deletions src/sentry/db/models/manager/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,13 +6,13 @@
from collections.abc import Callable, Collection, Generator, Mapping, MutableMapping, Sequence
from contextlib import contextmanager
from enum import IntEnum, auto
from typing import Any, Generic
from typing import Any

from django.conf import settings
from django.db import models, router
from django.db.models import Model
from django.db.models.fields import Field
from django.db.models.manager import BaseManager as DjangoBaseManager
from django.db.models.manager import Manager as DjangoBaseManager
from django.db.models.signals import class_prepared, post_delete, post_init, post_save
from django.utils.encoding import smart_str

Expand Down Expand Up @@ -69,7 +69,10 @@ def make_key(model: Any, prefix: str, kwargs: Mapping[str, Model | int | str]) -
return f"{prefix}:{model.__name__}:{md5_text(kwargs_bits_str).hexdigest()}"


class BaseManager(DjangoBaseManager.from_queryset(BaseQuerySet), Generic[M]): # type: ignore[misc]
_base_manager_base = DjangoBaseManager.from_queryset(BaseQuerySet, "_base_manager_base")


class BaseManager(_base_manager_base[M]):
lookup_handlers = {"iexact": lambda x: x.upper()}
use_for_related_fields = True

Expand Down
9 changes: 6 additions & 3 deletions src/sentry/models/files/abstractfile.py
Original file line number Diff line number Diff line change
Expand Up @@ -219,7 +219,8 @@ class Meta:

def _get_chunked_blob(self, mode=None, prefetch=False, prefetch_to=None, delete=True):
return ChunkedFileBlobIndexWrapper(
self.FILE_BLOB_INDEX_MODEL.objects.filter(file=self)
# TODO: file blob inheritance hierarchy is unsound
self.FILE_BLOB_INDEX_MODEL.objects.filter(file=self) # type: ignore[misc]
.select_related("blob")
.order_by("offset"),
mode=mode,
Expand Down Expand Up @@ -295,7 +296,8 @@ def putfile(self, fileobj, blob_size=DEFAULT_BLOB_SIZE, commit=True, logger=noop
blob_fileobj = ContentFile(contents)
blob = self.FILE_BLOB_MODEL.from_file(blob_fileobj, logger=logger)
results.append(
self.FILE_BLOB_INDEX_MODEL.objects.create(file=self, blob=blob, offset=offset)
# TODO: file blob inheritance hierarchy is unsound
self.FILE_BLOB_INDEX_MODEL.objects.create(file=self, blob=blob, offset=offset) # type: ignore[misc]
)
offset += blob.size
self.size = offset
Expand Down Expand Up @@ -334,7 +336,8 @@ def assemble_from_file_blob_ids(self, file_blob_ids, checksum):
offset = 0
for blob in file_blobs:
try:
self.FILE_BLOB_INDEX_MODEL.objects.create(file=self, blob=blob, offset=offset)
# TODO: file blob inheritance hierarchy is unsound
self.FILE_BLOB_INDEX_MODEL.objects.create(file=self, blob=blob, offset=offset) # type: ignore[misc]
except IntegrityError:
# Most likely a `ForeignKeyViolation` like `SENTRY-11P5`, because
# the blob we want to link does not exist anymore
Expand Down
3 changes: 2 additions & 1 deletion src/sentry/models/files/abstractfileblob.py
Original file line number Diff line number Diff line change
Expand Up @@ -92,7 +92,8 @@ def _ensure_blob_owned(blob):
return
try:
with atomic_transaction(using=router.db_for_write(cls.FILE_BLOB_OWNER_MODEL)):
cls.FILE_BLOB_OWNER_MODEL.objects.create(
# TODO: file blob inheritance hierarchy is unsound
cls.FILE_BLOB_OWNER_MODEL.objects.create( # type: ignore[misc]
organization_id=organization.id, blob=blob
)
except IntegrityError:
Expand Down
4 changes: 2 additions & 2 deletions src/sentry/organizations/services/organization/impl.py
Original file line number Diff line number Diff line change
Expand Up @@ -475,7 +475,7 @@ def get_or_create_team_member(

def update_membership_flags(self, *, organization_member: RpcOrganizationMember) -> None:
model = OrganizationMember.objects.get(id=organization_member.id)
model.flags = self._deserialize_member_flags(organization_member.flags)
model.flags = self._deserialize_member_flags(organization_member.flags) # type: ignore[assignment] # TODO: make BitField a mypy plugin
model.save()

@classmethod
Expand Down Expand Up @@ -518,7 +518,7 @@ def merge_users(self, *, organization_id: int, from_user_id: int, to_user_id: in
return

if to_member is None:
to_member = OrganizationMember.objects.create(
to_member = OrganizationMember.objects.create( # type: ignore[misc] # TODO: make BitField a mypy plugin
organization_id=organization_id,
user_id=to_user_id,
role=from_member.role,
Expand Down
2 changes: 1 addition & 1 deletion src/sentry/reprocessing2.py
Original file line number Diff line number Diff line change
Expand Up @@ -547,7 +547,7 @@ def start_group_reprocessing(

# Create a duplicate row that has the same attributes by nulling out
# the primary key and saving
group.pk = group.id = None
group.pk = group.id = None # type: ignore[assignment] # XXX: intentional resetting pk
new_group = group # rename variable just to avoid confusion
del group
new_group.status = original_status
Expand Down
9 changes: 3 additions & 6 deletions tests/sentry/db/test_transactions.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
from unittest.mock import patch

import pytest
from django.db import IntegrityError, router, transaction
from django.db import router, transaction
from django.test import override_settings

from sentry.db.postgres.transactions import (
Expand Down Expand Up @@ -30,11 +30,8 @@ def test_collect_transaction_queries(self):
User.objects.filter(username="user1").first()

with transaction.atomic(using=router.db_for_write(Organization)):
try:
with transaction.atomic(using=router.db_for_write(Organization)):
Organization.objects.create(name=None)
except (IntegrityError, MaxSnowflakeRetryError):
pass
with pytest.raises(MaxSnowflakeRetryError):
Organization.objects.create(name=None) # type: ignore[misc] # intentional to trigger error

with transaction.atomic(using=router.db_for_write(Organization)):
Organization.objects.create(name="org3")
Expand Down

0 comments on commit 0d84170

Please sign in to comment.