|
18 | 18 | from django.core.exceptions import ValidationError |
19 | 19 | from django.core.files.storage import default_storage |
20 | 20 | from django.core.validators import URLValidator |
21 | | -from django.db import DatabaseError, transaction |
| 21 | +from django.db import DatabaseError, connection, transaction |
22 | 22 | from django.db import models as db |
23 | 23 | from django.db.models.expressions import RawSQL |
24 | 24 | from django.db.models.functions import Greatest, Left, Length |
|
104 | 104 |
|
105 | 105 | # pylint: disable=too-many-ancestors |
106 | 106 |
|
| 107 | +# Minimum trigram similarity between a user search and an email or a full name |
| 108 | +USERS_SEARCH_SIMILARITY_THRESHOLD = 0.2 |
| 109 | + |
| 110 | + |
| 111 | +def set_trigram_similarity_threshold(threshold): |
| 112 | + """ |
| 113 | + Set the threshold of the trigram `%` operator for the current transaction only, |
| 114 | + so that it never leaks to the next users of a pooled connection. |
| 115 | + """ |
| 116 | + with connection.cursor() as cursor: |
| 117 | + cursor.execute( |
| 118 | + "SELECT set_config('pg_trgm.similarity_threshold', %s, true)", |
| 119 | + [str(threshold)], |
| 120 | + ) |
| 121 | + |
107 | 122 |
|
108 | 123 | class NestedGenericViewSet(viewsets.GenericViewSet): |
109 | 124 | """ |
@@ -233,94 +248,89 @@ def get_queryset(self): |
233 | 248 |
|
234 | 249 | filter_data = filterset.form.cleaned_data |
235 | 250 | query = remove_accents(filter_data["q"]) |
| 251 | + if not query: |
| 252 | + return [] |
236 | 253 |
|
237 | | - # For emails, match emails by Levenstein distance to prevent typing errors |
| 254 | + limit = settings.API_USERS_LIST_LIMIT |
| 255 | + |
| 256 | + # The trigram indexes of migration 0027 are built on these two expressions: |
| 257 | + # only a `%` filter on them uses the indexes, any other condition on the |
| 258 | + # user table (a similarity score, the raw columns) scans all the users. |
| 259 | + queryset = queryset.alias( |
| 260 | + unaccented_email=db.Func( |
| 261 | + "email", function="f_unaccent", output_field=db.TextField() |
| 262 | + ), |
| 263 | + unaccented_full_name=db.Func( |
| 264 | + "full_name", function="f_unaccent", output_field=db.TextField() |
| 265 | + ), |
| 266 | + ) |
| 267 | + |
| 268 | + # For emails, match emails by Levenstein distance to prevent typing errors, |
| 269 | + # among the users that the trigram index finds close enough to the query |
238 | 270 | if "@" in query: |
239 | | - return ( |
240 | | - queryset.annotate( |
241 | | - distance=RawSQL( |
242 | | - "levenshtein(unaccent(email::text), %s::text)", (query,) |
| 271 | + with transaction.atomic(): |
| 272 | + set_trigram_similarity_threshold(USERS_SEARCH_SIMILARITY_THRESHOLD) |
| 273 | + return list( |
| 274 | + queryset.filter(unaccented_email__trigram_similar=query) |
| 275 | + .annotate( |
| 276 | + distance=RawSQL( |
| 277 | + "levenshtein(unaccent(email::text), %s::text)", (query,) |
| 278 | + ) |
243 | 279 | ) |
| 280 | + .filter(distance__lte=3) |
| 281 | + .order_by("distance", "email")[:limit] |
244 | 282 | ) |
245 | | - .filter(distance__lte=3) |
246 | | - .order_by("distance", "email")[: settings.API_USERS_LIST_LIMIT] |
247 | | - ) |
248 | 283 |
|
249 | | - # Use trigram similarity for non-email-like queries |
250 | | - # For performance reasons we filter first by similarity, which relies on an |
251 | | - # index, then only calculate precise similarity scores for sorting purposes. |
| 284 | + # Use trigram similarity for non-email-like queries. |
252 | 285 | # |
253 | | - # Additionally results are reordered to prefer users "closer" to the current |
254 | | - # user: users they recently shared documents with, then same email domain. |
255 | | - # To achieve that without complex SQL, we build a proximity score in Python |
256 | | - # and return the top N results. |
257 | | - # For security results, users that match neither of these proximity criteria |
| 286 | + # Results are ordered to prefer users "closer" to the current user: users |
| 287 | + # they recently shared documents with, then users of the same email domain. |
| 288 | + # For security reasons, users that match neither of these proximity criteria |
258 | 289 | # are not returned at all, to prevent email enumeration. |
259 | 290 | current_user = self.request.user |
260 | 291 | shared_map = users_sharing_documents_with(current_user.id) |
| 292 | + shared_ids = list(shared_map) |
| 293 | + user_email_domain = get_domain_from_email(current_user.email) |
261 | 294 |
|
262 | | - user_email_domain = get_domain_from_email(current_user.email) or "" |
263 | | - |
264 | | - candidates = list( |
265 | | - queryset.annotate( |
266 | | - sim_email=TrigramSimilarity("email", query), |
267 | | - sim_name=TrigramSimilarity("full_name", query), |
| 295 | + similar_users = queryset.filter( |
| 296 | + db.Q(unaccented_email__trigram_similar=query) |
| 297 | + | db.Q(unaccented_full_name__trigram_similar=query) |
| 298 | + ).annotate( |
| 299 | + similarity=Greatest( |
| 300 | + TrigramSimilarity("unaccented_email", query), |
| 301 | + TrigramSimilarity("unaccented_full_name", query), |
268 | 302 | ) |
269 | | - .annotate(similarity=Greatest("sim_email", "sim_name")) |
270 | | - .filter(similarity__gt=0.2) |
271 | | - .order_by("-similarity") |
272 | 303 | ) |
273 | 304 |
|
274 | | - # Keep only users that either share documents with the current user |
275 | | - # or have an email with the same domain as the current user. |
276 | | - filtered_candidates = [] |
277 | | - for u in candidates: |
278 | | - candidate_domain = get_domain_from_email(u.email) or "" |
279 | | - if shared_map.get(u.id) or ( |
280 | | - user_email_domain and candidate_domain == user_email_domain |
281 | | - ): |
282 | | - filtered_candidates.append(u) |
283 | | - |
284 | | - candidates = filtered_candidates |
285 | | - |
286 | | - # Build ordering key for each candidate |
287 | | - def _sort_key(u): |
288 | | - # shared priority: most recent first |
289 | | - # Use shared_last_at timestamp numeric for secondary ordering when shared. |
290 | | - shared_last_at = shared_map.get(u.id) |
291 | | - if shared_last_at: |
292 | | - is_shared = 1 |
293 | | - shared_score = int(shared_last_at.timestamp()) |
294 | | - else: |
295 | | - is_shared = 0 |
296 | | - shared_score = 0 |
297 | | - |
298 | | - # domain proximity |
299 | | - candidate_email_domain = get_domain_from_email(u.email) or "" |
| 305 | + with transaction.atomic(): |
| 306 | + set_trigram_similarity_threshold(USERS_SEARCH_SIMILARITY_THRESHOLD) |
300 | 307 |
|
301 | | - same_full_domain = ( |
302 | | - 1 |
303 | | - if candidate_email_domain |
304 | | - and candidate_email_domain == user_email_domain |
305 | | - else 0 |
| 308 | + # Bounded by the number of users the current user shares documents with |
| 309 | + shared_users = ( |
| 310 | + list(similar_users.filter(id__in=shared_ids)) if shared_ids else [] |
306 | 311 | ) |
307 | 312 |
|
308 | | - # similarity fallback |
309 | | - sim = getattr(u, "similarity", 0) or 0 |
310 | | - |
311 | | - return ( |
312 | | - is_shared, |
313 | | - shared_score, |
314 | | - same_full_domain, |
315 | | - sim, |
316 | | - ) |
| 313 | + # A domain can count many users: only fetch the ones that can still fit |
| 314 | + domain_users = [] |
| 315 | + if user_email_domain and len(shared_users) < limit: |
| 316 | + domain_users = list( |
| 317 | + similar_users.filter(email__endswith=f"@{user_email_domain}") |
| 318 | + .exclude(id__in=shared_ids) |
| 319 | + .order_by("-similarity")[: limit - len(shared_users)] |
| 320 | + ) |
317 | 321 |
|
318 | | - # Sort candidates by the key descending and return top N as a queryset-like |
319 | | - # list. Keep return type consistent with previous behavior (QuerySet slice |
320 | | - # was returned) by returning a list of model instances. |
321 | | - candidates.sort(key=_sort_key, reverse=True) |
| 322 | + # Most recently shared first, then same domain, then the closest match |
| 323 | + shared_users.sort( |
| 324 | + key=lambda user: ( |
| 325 | + shared_map[user.id], |
| 326 | + bool(user_email_domain) |
| 327 | + and get_domain_from_email(user.email) == user_email_domain, |
| 328 | + user.similarity, |
| 329 | + ), |
| 330 | + reverse=True, |
| 331 | + ) |
322 | 332 |
|
323 | | - return candidates[: settings.API_USERS_LIST_LIMIT] |
| 333 | + return (shared_users + domain_users)[:limit] |
324 | 334 |
|
325 | 335 | @drf.decorators.action( |
326 | 336 | detail=False, |
|
0 commit comments