1
- import abc
2
- from typing import Any
1
+ from __future__ import annotations
2
+
3
+ from typing import Any , Self
3
4
4
5
from django .core import exceptions
5
6
from django .core .exceptions import EmptyResultSet
6
7
from django .db import connections , router , transaction
7
8
from django .db .models import QuerySet , sql
8
9
10
+ from sentry .db .models .manager import M
9
11
from sentry .signals import post_update
10
12
11
13
12
- class BaseQuerySet (QuerySet , abc . ABC ):
13
- def __init__ (self , * args , ** kwargs ) :
14
+ class BaseQuerySet (QuerySet [ M ] ):
15
+ def __init__ (self , * args : Any , ** kwargs : Any ) -> None :
14
16
super ().__init__ (* args , ** kwargs )
15
17
self ._with_post_update_signal = False
16
18
17
- def with_post_update_signal (self , enable : bool ) -> "BaseQuerySet" :
19
+ def with_post_update_signal (self , enable : bool ) -> Self :
18
20
"""
19
21
Enables sending a `post_update` signal after this queryset runs an update command. Note that this is less
20
22
efficient than just running the update. To get the list of group ids affected, we first run the query to
@@ -24,12 +26,12 @@ def with_post_update_signal(self, enable: bool) -> "BaseQuerySet":
24
26
qs ._with_post_update_signal = enable
25
27
return qs
26
28
27
- def _clone (self ) -> "BaseQuerySet" :
29
+ def _clone (self ) -> Self :
28
30
qs = super ()._clone () # type: ignore[misc]
29
31
qs ._with_post_update_signal = self ._with_post_update_signal
30
32
return qs
31
33
32
- def update_with_returning (self , returned_fields : list [str ], ** kwargs ) :
34
+ def update_with_returning (self , returned_fields : list [str ], ** kwargs : Any ) -> list [ tuple [ int ]] :
33
35
"""
34
36
Copied and modified from `Queryset.update()` to support `RETURNING <returned_fields>`
35
37
"""
@@ -77,7 +79,9 @@ def update_with_returning(self, returned_fields: list[str], **kwargs):
77
79
78
80
def update (self , ** kwargs : Any ) -> int :
79
81
if self ._with_post_update_signal :
80
- pk = self .model ._meta .pk .name
82
+ pk_field = self .model ._meta .pk
83
+ assert pk_field is not None
84
+ pk = pk_field .name
81
85
ids = [result [0 ] for result in self .update_with_returning ([pk ], ** kwargs )]
82
86
if ids :
83
87
updated_fields = list (kwargs .keys ())
@@ -86,17 +90,17 @@ def update(self, **kwargs: Any) -> int:
86
90
else :
87
91
return super ().update (** kwargs )
88
92
89
- def using_replica (self ) -> "BaseQuerySet" :
93
+ def using_replica (self ) -> Self :
90
94
"""
91
95
Use read replica for this query. Database router is expected to use the
92
96
`replica=True` hint to make routing decision.
93
97
"""
94
98
return self .using (router .db_for_read (self .model , replica = True ))
95
99
96
- def defer (self , * args : Any , ** kwargs : Any ) -> "BaseQuerySet" :
100
+ def defer (self , * args : Any , ** kwargs : Any ) -> Self :
97
101
raise NotImplementedError ("Use ``values_list`` instead [performance]." )
98
102
99
- def only (self , * args : Any , ** kwargs : Any ) -> "BaseQuerySet" :
103
+ def only (self , * args : Any , ** kwargs : Any ) -> Self :
100
104
# In rare cases Django can use this if a field is unexpectedly deferred. This
101
105
# mostly can happen if a field is added to a model, and then an old pickle is
102
106
# passed to a process running the new code. So if you see this error after a
0 commit comments