Coverage for extras/lookups.py: 67%

1from django.contrib.postgres.fields import ArrayField 

2from django.contrib.postgres.fields.ranges import RangeField 

3from django.db.models import CharField, JSONField, Lookup 

4from django.db.models.expressions import Col 

5from django.db.models.fields.json import KeyTextTransform 

6from django.db.models.lookups import IContains, IEndsWith, IExact, IStartsWith 

7 

8from .fields import CachedValueField, ChoiceSetField 

9 

10__all__ = ( 

11 'ChoiceValueLookup', 

12 'CollatedIContains', 

13 'CollatedIEndsWith', 

14 'CollatedIExact', 

15 'CollatedIStartsWith', 

16 'Empty', 

17 'JSONEmpty', 

18 'NetContainsOrEquals', 

19 'NetHost', 

20 'RangeContains', 

21) 

22 

23# The ICU collation created by dcim.migrations.0197_natural_sort_collation and applied to 

24# the name field of most models. 

25NATURAL_SORT_COLLATION = 'natural_sort' 

26 

27 

28class RangeContains(Lookup): 

29 """ 

30 Filter ArrayField(RangeField) columns where ANY element-range contains the scalar RHS. 

31 

32 Usage (ORM): 

33 Model.objects.filter(<range_array_field>__range_contains=<scalar>) 

34 

35 Works with int4range[], int8range[], daterange[], tstzrange[], etc. 

36 """ 

37 

38 lookup_name = 'range_contains' 

39 

40 def as_sql(self, compiler, connection): 

41 # Compile LHS (the array-of-ranges column/expression) and RHS (scalar) 

42 lhs, lhs_params = self.process_lhs(compiler, connection) 

43 rhs, rhs_params = self.process_rhs(compiler, connection) 

44 

45 # Guard: only allow ArrayField whose base_field is a PostgreSQL RangeField 

46 field = getattr(self.lhs, 'output_field', None) 

47 if not (isinstance(field, ArrayField) and isinstance(field.base_field, RangeField)): 

48 raise TypeError('range_contains is only valid for ArrayField(RangeField) columns') 

49 

50 # Range-contains-element using EXISTS + UNNEST keeps the range on the LHS: r @> value 

51 sql = f"EXISTS (SELECT 1 FROM unnest({lhs}) AS r WHERE r @> {rhs})" 

52 params = lhs_params + rhs_params 

53 return sql, params 

54 

55 

56class ChoiceValueLookup(Lookup): 

57 """ 

58 Match rows where any [value, label] pair in a ChoiceSetField has the given value. 

59 

60 Compares the RHS against the first element (the value) of each pair. 

61 """ 

62 lookup_name = 'choice_value' 

63 prepare_rhs = False 

64 

65 def as_sql(self, compiler, connection): 

66 lhs, lhs_params = self.process_lhs(compiler, connection) 

67 rhs, rhs_params = self.process_rhs(compiler, connection) 

68 # Slice the value column of the two-dimensional array and match any element 

69 return f'{rhs} = ANY({lhs}[:][1:1])', [*rhs_params, *lhs_params] 

70 

71 

72class Empty(Lookup): 

73 """ 

74 Filter on whether a string is empty. 

75 """ 

76 lookup_name = 'empty' 

77 prepare_rhs = False 

78 

79 def as_sql(self, compiler, connection): 

80 sql, params = compiler.compile(self.lhs) 

81 if self.rhs: 

82 return f"CAST(LENGTH({sql}) AS BOOLEAN) IS NOT TRUE", params 

83 return f"CAST(LENGTH({sql}) AS BOOLEAN) IS TRUE", params 

84 

85 

86class JSONEmpty(Lookup): 

87 """ 

88 Support "empty" lookups for JSONField keys. 

89 

90 A key is considered empty if it is "", null, or does not exist. 

91 """ 

92 lookup_name = 'empty' 

93 

94 def as_sql(self, compiler, connection): 

95 # self.lhs.lhs is the parent expression (could be a JSONField or another KeyTransform) 

96 # Rebuild the expression using KeyTextTransform to guarantee ->> (text) 

97 text_expr = KeyTextTransform(self.lhs.key_name, self.lhs.lhs) 

98 lhs_sql, lhs_params = compiler.compile(text_expr) 

99 

100 value = self.rhs 

101 if value not in (True, False): 

102 raise ValueError("The 'empty' lookup only accepts True or False.") 

103 

104 condition = '' if value else 'NOT ' 

105 sql = f"(NULLIF({lhs_sql}, '') IS {condition}NULL)" 

106 

107 return sql, lhs_params 

108 

109 

110class NetHost(Lookup): 

111 """ 

112 Similar to ipam.lookups.NetHost, but casts the field to INET. 

113 """ 

114 lookup_name = 'net_host' 

115 

116 def as_sql(self, qn, connection): 

117 lhs, lhs_params = self.process_lhs(qn, connection) 

118 rhs, rhs_params = self.process_rhs(qn, connection) 

119 params = lhs_params + rhs_params 

120 return f'HOST(CAST({lhs} AS INET)) = HOST({rhs})', params 

121 

122 

123class NetContainsOrEquals(Lookup): 

124 """ 

125 Similar to ipam.lookups.NetContainsOrEquals, but casts the field to INET. 

126 """ 

127 lookup_name = 'net_contains_or_equals' 

128 

129 def as_sql(self, qn, connection): 

130 lhs, lhs_params = self.process_lhs(qn, connection) 

131 rhs, rhs_params = self.process_rhs(qn, connection) 

132 params = lhs_params + rhs_params 

133 return f'CAST({lhs} AS INET) >>= {rhs}', params 

134 

135 

136class CollatedCaseInsensitiveMixin: 

137 """ 

138 Apply the column's collation to the right-hand side of a case-insensitive comparison. 

139 

140 UPPER() folds according to the collation of its argument. Django uppercases the column 

141 under the column's own collation but the parameter under the database default, so for a 

142 column using natural_sort the two sides disagree: UPPER('ß') is 'SS' on the left and 

143 'ß' on the right, and the comparison silently matches nothing (#23012). 

144 

145 The COLLATE clause must sit inside UPPER(), not after the comparison, or it applies to 

146 the comparison's result rather than to its operand and has no effect. 

147 

148 Tested in dcim.tests.test_filtersets.DeviceCollatedFilterTestCase, which is where the 

149 collated fields these lookups act upon are defined. 

150 """ 

151 def process_rhs(self, compiler, connection): 

152 rhs, params = super().process_rhs(compiler, connection) 

153 collation = getattr(self.lhs.output_field, 'db_collation', None) 

154 

155 # Restricted to a bare column compared against a single placeholder. An expression 

156 # wrapping the column (Collate() and CollateAsChar() in particular) may already 

157 # carry an explicit collation, and PostgreSQL rejects two explicit collations in 

158 # one comparison. Requiring a Col also avoids reading a collation from an 

159 # annotation's output_field which the annotation itself does not carry, as Concat() 

160 # and Coalesce() both do. 

161 # 

162 # The placeholder is compared literally rather than inspected structurally: a field 

163 # declaring its own get_placeholder() compiles to something other than '%s', and 

164 # splicing a COLLATE clause into that is not safe. Any other rhs is a deliberate 

165 # opt-out which leaves the lookup at its previous behaviour. 

166 if collation == NATURAL_SORT_COLLATION and rhs == '%s' and isinstance(self.lhs, Col): 

167 # The collation name cannot be passed as a query parameter, but it originates 

168 # from the field definition rather than from user input. 

169 rhs = f'%s COLLATE "{collation}"' 

170 

171 return rhs, params 

172 

173 

174class CollatedIContains(CollatedCaseInsensitiveMixin, IContains): 

175 pass 

176 

177 

178class CollatedIExact(CollatedCaseInsensitiveMixin, IExact): 

179 pass 

180 

181 

182class CollatedIStartsWith(CollatedCaseInsensitiveMixin, IStartsWith): 

183 pass 

184 

185 

186class CollatedIEndsWith(CollatedCaseInsensitiveMixin, IEndsWith): 

187 pass 

188 

189 

190ArrayField.register_lookup(RangeContains) 

191ChoiceSetField.register_lookup(ChoiceValueLookup) 

192CharField.register_lookup(Empty) 

193JSONField.register_lookup(JSONEmpty) 

194CachedValueField.register_lookup(NetHost) 

195CachedValueField.register_lookup(NetContainsOrEquals) 

196 

197# Override the built-in case-insensitive lookups so that they respect the collation of the 

198# column being searched. 

199CharField.register_lookup(CollatedIContains) 

200CharField.register_lookup(CollatedIExact) 

201CharField.register_lookup(CollatedIStartsWith) 

202CharField.register_lookup(CollatedIEndsWith)