Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

ref: get proper base manager typing #72221

Merged
merged 1 commit into from
Jul 29, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

talked a bit with @markstory about these particular models -- we believe there's a way to convert these to free functions to capture the dependent-types inheritance hierarchy with the type system and make these sound -- but as a followup

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
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

the issue here is HC converts the bitfield to an integer over the wire and the assignment then confuses our workaround for bitfields

I plan to follow up by making bitfield an actual mypy plugin which understands the assignments and values properly here

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
Loading