Skip to content

ref: improve typing of sentry.db.models #72963

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

Merged
merged 1 commit into from
Jun 20, 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
10 changes: 3 additions & 7 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -235,8 +235,6 @@ module = [
"sentry.auth.system",
"sentry.auth.view",
"sentry.db.mixin",
"sentry.db.models.paranoia",
"sentry.db.models.utils",
"sentry.db.postgres.base",
"sentry.db.router",
"sentry.digests.notifications",
Expand Down Expand Up @@ -519,11 +517,9 @@ module = [
"sentry.api.helpers.source_map_helper",
"sentry.buffer.*",
"sentry.build.*",
"sentry.db.models.manager",
"sentry.db.models.manager.base",
"sentry.db.models.manager.base_query_set",
"sentry.db.models.manager.types",
"sentry.db.models.query",
"sentry.db.models.manager.*",
"sentry.db.models.paranoia",
"sentry.db.models.utils",
"sentry.eventstore.reprocessing.redis",
"sentry.eventtypes.error",
"sentry.grouping.component",
Expand Down
4 changes: 0 additions & 4 deletions src/sentry/db/exceptions.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,2 @@
class QueryError(Exception):
pass


class CannotResolveExpression(Exception):
pass
16 changes: 13 additions & 3 deletions src/sentry/db/models/paranoia.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,14 +8,21 @@
from sentry.db.models.manager.types import M


def _bogus_delete_return_value() -> tuple[int, dict[str, int]]:
# django'd delete returns (# deleted, dict[model name, # deleted])
# but we never use this value (and aren't actually deleting!) so...
return (0, {})


class ParanoidQuerySet(BaseQuerySet[M]):
"""
Prevents objects from being hard-deleted. Instead, sets the
``date_deleted``, effectively soft-deleting the object.
"""

def delete(self) -> None:
def delete(self) -> tuple[int, dict[str, int]]:
self.update(date_deleted=timezone.now())
return _bogus_delete_return_value()


class ParanoidManager(BaseManager[M]):
Expand All @@ -35,7 +42,10 @@ class Meta:
objects: ClassVar[ParanoidManager[Self]] = ParanoidManager()
with_deleted: ClassVar[BaseManager[Self]] = BaseManager()

def delete(self) -> None:
self.update(date_deleted=timezone.now())
def delete(
self, using: str | None = None, keep_parents: bool = False
) -> tuple[int, dict[str, int]]:
self.update(using=using, date_deleted=timezone.now())
return _bogus_delete_return_value()

__repr__ = sane_repr("id")
48 changes: 44 additions & 4 deletions src/sentry/db/models/query.py
Original file line number Diff line number Diff line change
@@ -1,17 +1,16 @@
from __future__ import annotations

import itertools
import operator
from functools import reduce
from typing import TYPE_CHECKING, Any, Literal

from django.db import IntegrityError, router, transaction
from django.db.models import Model, Q
from django.db.models.expressions import CombinedExpression
from django.db.models import F, Model, Q
from django.db.models.expressions import BaseExpression, CombinedExpression, Value
from django.db.models.fields import Field
from django.db.models.signals import post_save

from .utils import resolve_combined_expression

if TYPE_CHECKING:
from sentry.db.models.base import BaseModel

Expand All @@ -21,6 +20,47 @@
"update_or_create",
)

COMBINED_EXPRESSION_CALLBACKS = {
CombinedExpression.ADD: operator.add,
CombinedExpression.SUB: operator.sub,
CombinedExpression.MUL: operator.mul,
CombinedExpression.DIV: operator.floordiv,
CombinedExpression.MOD: operator.mod,
CombinedExpression.BITAND: operator.and_,
CombinedExpression.BITOR: operator.or_,
}


class CannotResolveExpression(Exception):
pass


def resolve_combined_expression(instance: Model, node: CombinedExpression) -> BaseExpression:
def _resolve(instance: Model, node: BaseExpression | F) -> BaseExpression:
if isinstance(node, Value):
return node.value
if isinstance(node, F):
return getattr(instance, node.name)
if isinstance(node, CombinedExpression):
return resolve_combined_expression(instance, node)
return node

if isinstance(node, Value):
return node.value
if not isinstance(node, CombinedExpression):
raise CannotResolveExpression
op = COMBINED_EXPRESSION_CALLBACKS.get(node.connector, None)
if not op:
raise CannotResolveExpression
if hasattr(node, "children"):
children = node.children
else:
children = [node.lhs, node.rhs]
runner = _resolve(instance, children[0])
for n in children[1:]:
runner = op(runner, _resolve(instance, n))
return runner


def _get_field(model: type[BaseModel], key: str) -> Field[object, object]:
field = model._meta.get_field(key)
Expand Down
75 changes: 25 additions & 50 deletions src/sentry/db/models/utils.py
Original file line number Diff line number Diff line change
@@ -1,57 +1,19 @@
from __future__ import annotations

import operator
from collections.abc import Container
from typing import Any
from typing import TYPE_CHECKING, Any, Generic, Self, TypeVar, overload
from uuid import uuid4

from django.db.models import F, Field, Model
from django.db.models.expressions import BaseExpression, CombinedExpression, Value
from django.db.models import Field, Model
from django.utils.crypto import get_random_string
from django.utils.text import slugify

from sentry.db.exceptions import CannotResolveExpression

COMBINED_EXPRESSION_CALLBACKS = {
CombinedExpression.ADD: operator.add,
CombinedExpression.SUB: operator.sub,
CombinedExpression.MUL: operator.mul,
CombinedExpression.DIV: operator.floordiv,
CombinedExpression.MOD: operator.mod,
CombinedExpression.BITAND: operator.and_,
CombinedExpression.BITOR: operator.or_,
}


def resolve_combined_expression(instance: Model, node: BaseExpression) -> BaseExpression:
def _resolve(instance: Model, node: BaseExpression | F) -> BaseExpression:
if isinstance(node, Value):
return node.value
if isinstance(node, F):
return getattr(instance, node.name)
if isinstance(node, CombinedExpression):
return resolve_combined_expression(instance, node)
return node

if isinstance(node, Value):
return node.value
if not hasattr(node, "connector"):
raise CannotResolveExpression
op = COMBINED_EXPRESSION_CALLBACKS.get(node.connector, None)
if not op:
raise CannotResolveExpression
if hasattr(node, "children"):
children = node.children
else:
children = [node.lhs, node.rhs]
runner = _resolve(instance, children[0])
for n in children[1:]:
runner = op(runner, _resolve(instance, n))
return runner
if TYPE_CHECKING:
from sentry.db.models.base import Model as SentryModel


def unique_db_instance(
inst: Model,
inst: SentryModel,
base_value: str,
reserved: Container[str] = (),
max_length: int = 30,
Expand All @@ -62,7 +24,7 @@ def unique_db_instance(
if base_value is not None:
base_value = base_value.strip()
if base_value in reserved:
base_value = None
base_value = ""

if not base_value:
base_value = uuid4().hex[:12]
Expand Down Expand Up @@ -105,7 +67,7 @@ def unique_db_instance(


def slugify_instance(
inst: Model,
inst: SentryModel,
label: str,
reserved: Container[str] = (),
max_length: int = 30,
Expand All @@ -119,20 +81,33 @@ def slugify_instance(
return unique_db_instance(inst, value, reserved, max_length, field_name, *args, **kwargs)


class Creator:
# matches django-stubs for Field
_ST = TypeVar("_ST", contravariant=True)
_GT = TypeVar("_GT", covariant=True)


class Creator(Generic[_ST, _GT]):
"""
A descriptor that invokes `to_python` when attributes are set.
This provides backwards compatibility for fields that used to use
SubfieldBase which will be removed in Django1.10
"""

def __init__(self, field: Field):
def __init__(self, field: Field[_ST, _GT]) -> None:
self.field = field

def __get__(self, obj: Model, type: Any = None) -> Any:
if obj is None:
@overload
def __get__(self, inst: Model, owner: type[Any]) -> Any:
...

@overload
def __get__(self, inst: None, owner: type[Any]) -> Self:
...

def __get__(self, inst: Model | None, owner: type[Any]) -> Self | Any:
if inst is None:
return self
return obj.__dict__[self.field.name]
return inst.__dict__[self.field.name]

def __set__(self, obj: Model, value: Any) -> None:
obj.__dict__[self.field.name] = self.field.to_python(value)
Loading