Coverage for src/couchers/servicers/search.py: 83%
229 statements
« prev ^ index » next coverage.py v7.5.0, created at 2024-12-20 18:03 +0000
1"""
2See //docs/search.md for overview.
3"""
5from datetime import timedelta
7import grpc
8from sqlalchemy.sql import and_, func, or_
10from couchers import errors
11from couchers.crypto import decrypt_page_token, encrypt_page_token
12from couchers.models import (
13 Cluster,
14 ClusterSubscription,
15 Event,
16 EventOccurrence,
17 EventOccurrenceAttendee,
18 EventOrganizer,
19 EventSubscription,
20 Node,
21 Page,
22 PageType,
23 PageVersion,
24 Reference,
25 User,
26)
27from couchers.servicers.account import has_strong_verification
28from couchers.servicers.api import (
29 hostingstatus2sql,
30 meetupstatus2sql,
31 parkingdetails2sql,
32 sleepingarrangement2sql,
33 smokinglocation2sql,
34 user_model_to_pb,
35)
36from couchers.servicers.communities import community_to_pb
37from couchers.servicers.events import event_to_pb
38from couchers.servicers.groups import group_to_pb
39from couchers.servicers.pages import page_to_pb
40from couchers.sql import couchers_select as select
41from couchers.utils import (
42 create_coordinate,
43 dt_from_millis,
44 last_active_coarsen,
45 millis_from_dt,
46 now,
47 to_aware_datetime,
48)
49from proto import search_pb2, search_pb2_grpc
51# searches are a bit expensive, we'd rather send back a bunch of results at once than lots of small pages
52MAX_PAGINATION_LENGTH = 100
54REGCONFIG = "english"
55TRI_SIMILARITY_THRESHOLD = 0.6
56TRI_SIMILARITY_WEIGHT = 5
59def _join_with_space(coalesces):
60 # the objects in coalesces are not strings, so we can't do " ".join(coalesces). They're SQLAlchemy magic.
61 if not coalesces:
62 return ""
63 out = coalesces[0]
64 for coalesce in coalesces[1:]:
65 out += " " + coalesce
66 return out
69def _build_tsv(A, B=None, C=None, D=None):
70 """
71 Given lists for A, B, C, and D, builds a tsvector from them.
72 """
73 B = B or []
74 C = C or []
75 D = D or []
76 tsv = func.setweight(func.to_tsvector(REGCONFIG, _join_with_space([func.coalesce(bit, "") for bit in A])), "A")
77 if B:
78 tsv = tsv.concat(
79 func.setweight(func.to_tsvector(REGCONFIG, _join_with_space([func.coalesce(bit, "") for bit in B])), "B")
80 )
81 if C:
82 tsv = tsv.concat(
83 func.setweight(func.to_tsvector(REGCONFIG, _join_with_space([func.coalesce(bit, "") for bit in C])), "C")
84 )
85 if D:
86 tsv = tsv.concat(
87 func.setweight(func.to_tsvector(REGCONFIG, _join_with_space([func.coalesce(bit, "") for bit in D])), "D")
88 )
89 return tsv
92def _build_doc(A, B=None, C=None, D=None):
93 """
94 Builds the raw document (without to_tsvector and weighting), used for extracting snippet
95 """
96 B = B or []
97 C = C or []
98 D = D or []
99 doc = _join_with_space([func.coalesce(bit, "") for bit in A])
100 if B:
101 doc += " " + _join_with_space([func.coalesce(bit, "") for bit in B])
102 if C:
103 doc += " " + _join_with_space([func.coalesce(bit, "") for bit in C])
104 if D:
105 doc += " " + _join_with_space([func.coalesce(bit, "") for bit in D])
106 return doc
109def _similarity(statement, text):
110 return func.word_similarity(func.unaccent(statement), func.unaccent(text))
113def _gen_search_elements(statement, title_only, next_rank, page_size, A, B=None, C=None, D=None):
114 """
115 Given an sql statement and four sets of fields, (A, B, C, D), generates a bunch of postgres expressions for full text search.
117 The four sets are in decreasing order of "importance" for ranking.
119 A should be the "title", the others can be anything.
121 If title_only=True, we only perform a trigram search against A only
122 """
123 B = B or []
124 C = C or []
125 D = D or []
126 if not title_only:
127 # a postgres tsquery object that can be used to match against a tsvector
128 tsq = func.websearch_to_tsquery(REGCONFIG, statement)
130 # the tsvector object that we want to search against with our tsquery
131 tsv = _build_tsv(A, B, C, D)
133 # document to generate snippet from
134 doc = _build_doc(A, B, C, D)
136 title = _build_doc(A)
138 # trigram based text similarity between title and sql statement string
139 sim = _similarity(statement, title)
141 # ranking algo, weigh the similarity a lot, the text-based ranking less
142 rank = (TRI_SIMILARITY_WEIGHT * sim + func.ts_rank_cd(tsv, tsq)).label("rank")
144 # the snippet with results highlighted
145 snippet = func.ts_headline(REGCONFIG, doc, tsq, "StartSel=**,StopSel=**").label("snippet")
147 def execute_search_statement(session, orig_statement):
148 """
149 Does the right search filtering, limiting, and ordering for the initial statement
150 """
151 return session.execute(
152 orig_statement.where(or_(tsv.op("@@")(tsq), sim > TRI_SIMILARITY_THRESHOLD))
153 .where(rank <= next_rank if next_rank is not None else True)
154 .order_by(rank.desc())
155 .limit(page_size + 1)
156 ).all()
158 else:
159 title = _build_doc(A)
161 # trigram based text similarity between title and sql statement string
162 sim = _similarity(statement, title)
164 # ranking algo, weigh the similarity a lot, the text-based ranking less
165 rank = sim.label("rank")
167 # used only for headline
168 tsq = func.websearch_to_tsquery(REGCONFIG, statement)
169 doc = _build_doc(A, B, C, D)
171 # the snippet with results highlighted
172 snippet = func.ts_headline(REGCONFIG, doc, tsq, "StartSel=**,StopSel=**").label("snippet")
174 def execute_search_statement(session, orig_statement):
175 """
176 Does the right search filtering, limiting, and ordering for the initial statement
177 """
178 return session.execute(
179 orig_statement.where(sim > TRI_SIMILARITY_THRESHOLD)
180 .where(rank <= next_rank if next_rank is not None else True)
181 .order_by(rank.desc())
182 .limit(page_size + 1)
183 ).all()
185 return rank, snippet, execute_search_statement
188def _search_users(session, search_statement, title_only, next_rank, page_size, context, include_users):
189 if not include_users:
190 return []
191 rank, snippet, execute_search_statement = _gen_search_elements(
192 search_statement,
193 title_only,
194 next_rank,
195 page_size,
196 [User.username, User.name],
197 [User.city],
198 [User.about_me],
199 [User.things_i_like, User.about_place, User.additional_information],
200 )
202 users = execute_search_statement(session, select(User, rank, snippet).where_users_visible(context))
204 return [
205 search_pb2.Result(
206 rank=rank,
207 user=user_model_to_pb(page, session, context),
208 snippet=snippet,
209 )
210 for page, rank, snippet in users
211 ]
214def _search_pages(session, search_statement, title_only, next_rank, page_size, context, include_places, include_guides):
215 rank, snippet, execute_search_statement = _gen_search_elements(
216 search_statement,
217 title_only,
218 next_rank,
219 page_size,
220 [PageVersion.title],
221 [PageVersion.address],
222 [],
223 [PageVersion.content],
224 )
225 if not include_places and not include_guides:
226 return []
228 latest_pages = (
229 select(func.max(PageVersion.id).label("id"))
230 .join(Page, Page.id == PageVersion.page_id)
231 .where(
232 or_(
233 (Page.type == PageType.place) if include_places else False,
234 (Page.type == PageType.guide) if include_guides else False,
235 )
236 )
237 .group_by(PageVersion.page_id)
238 .subquery()
239 )
241 pages = execute_search_statement(
242 session,
243 select(Page, rank, snippet)
244 .join(PageVersion, PageVersion.page_id == Page.id)
245 .join(latest_pages, latest_pages.c.id == PageVersion.id),
246 )
248 return [
249 search_pb2.Result(
250 rank=rank,
251 place=page_to_pb(session, page, context) if page.type == PageType.place else None,
252 guide=page_to_pb(session, page, context) if page.type == PageType.guide else None,
253 snippet=snippet,
254 )
255 for page, rank, snippet in pages
256 ]
259def _search_events(session, search_statement, title_only, next_rank, page_size, context):
260 rank, snippet, execute_search_statement = _gen_search_elements(
261 search_statement,
262 title_only,
263 next_rank,
264 page_size,
265 [Event.title],
266 [EventOccurrence.address, EventOccurrence.link],
267 [],
268 [EventOccurrence.content],
269 )
271 occurrences = execute_search_statement(
272 session,
273 select(EventOccurrence, rank, snippet)
274 .join(Event, Event.id == EventOccurrence.event_id)
275 .where(EventOccurrence.end_time >= func.now()),
276 )
278 return [
279 search_pb2.Result(
280 rank=rank,
281 event=event_to_pb(session, occurrence, context),
282 snippet=snippet,
283 )
284 for occurrence, rank, snippet in occurrences
285 ]
288def _search_clusters(
289 session, search_statement, title_only, next_rank, page_size, context, include_communities, include_groups
290):
291 if not include_communities and not include_groups:
292 return []
294 rank, snippet, execute_search_statement = _gen_search_elements(
295 search_statement,
296 title_only,
297 next_rank,
298 page_size,
299 [Cluster.name],
300 [PageVersion.address, PageVersion.title],
301 [Cluster.description],
302 [PageVersion.content],
303 )
305 latest_pages = (
306 select(func.max(PageVersion.id).label("id"))
307 .join(Page, Page.id == PageVersion.page_id)
308 .where(Page.type == PageType.main_page)
309 .group_by(PageVersion.page_id)
310 .subquery()
311 )
313 clusters = execute_search_statement(
314 session,
315 select(Cluster, rank, snippet)
316 .join(Page, Page.owner_cluster_id == Cluster.id)
317 .join(PageVersion, PageVersion.page_id == Page.id)
318 .join(latest_pages, latest_pages.c.id == PageVersion.id)
319 .where(Cluster.is_official_cluster if include_communities and not include_groups else True)
320 .where(~Cluster.is_official_cluster if not include_communities and include_groups else True),
321 )
323 return [
324 search_pb2.Result(
325 rank=rank,
326 community=(
327 community_to_pb(session, cluster.official_cluster_for_node, context)
328 if cluster.is_official_cluster
329 else None
330 ),
331 group=group_to_pb(session, cluster, context) if not cluster.is_official_cluster else None,
332 snippet=snippet,
333 )
334 for cluster, rank, snippet in clusters
335 ]
338class Search(search_pb2_grpc.SearchServicer):
339 def Search(self, request, context, session):
340 page_size = min(MAX_PAGINATION_LENGTH, request.page_size or MAX_PAGINATION_LENGTH)
341 # this is not an ideal page token, some results have equal rank (unlikely)
342 next_rank = float(request.page_token) if request.page_token else None
344 all_results = (
345 _search_users(
346 session,
347 request.query,
348 request.title_only,
349 next_rank,
350 page_size,
351 context,
352 request.include_users,
353 )
354 + _search_pages(
355 session,
356 request.query,
357 request.title_only,
358 next_rank,
359 page_size,
360 context,
361 request.include_places,
362 request.include_guides,
363 )
364 + _search_events(
365 session,
366 request.query,
367 request.title_only,
368 next_rank,
369 page_size,
370 context,
371 )
372 + _search_clusters(
373 session,
374 request.query,
375 request.title_only,
376 next_rank,
377 page_size,
378 context,
379 request.include_communities,
380 request.include_groups,
381 )
382 )
383 all_results.sort(key=lambda result: result.rank, reverse=True)
384 return search_pb2.SearchRes(
385 results=all_results[:page_size],
386 next_page_token=str(all_results[page_size].rank) if len(all_results) > page_size else None,
387 )
389 def UserSearch(self, request, context, session):
390 user = session.execute(select(User).where(User.id == context.user_id)).scalar_one()
392 statement = select(User).where_users_visible(context)
393 if request.HasField("query"):
394 if request.query_name_only:
395 statement = statement.where(
396 or_(User.name.ilike(f"%{request.query.value}%"), User.username.ilike(f"%{request.query.value}%"))
397 )
398 else:
399 statement = statement.where(
400 or_(
401 User.name.ilike(f"%{request.query.value}%"),
402 User.username.ilike(f"%{request.query.value}%"),
403 User.city.ilike(f"%{request.query.value}%"),
404 User.hometown.ilike(f"%{request.query.value}%"),
405 User.about_me.ilike(f"%{request.query.value}%"),
406 User.things_i_like.ilike(f"%{request.query.value}%"),
407 User.about_place.ilike(f"%{request.query.value}%"),
408 User.additional_information.ilike(f"%{request.query.value}%"),
409 )
410 )
412 if request.HasField("last_active"):
413 raw_dt = to_aware_datetime(request.last_active)
414 statement = statement.where(User.last_active >= last_active_coarsen(raw_dt))
416 if len(request.gender) > 0:
417 if not has_strong_verification(session, user):
418 context.abort(grpc.StatusCode.FAILED_PRECONDITION, errors.NEED_STRONG_VERIFICATION)
419 elif user.gender not in request.gender:
420 context.abort(grpc.StatusCode.FAILED_PRECONDITION, errors.MUST_INCLUDE_OWN_GENDER)
421 else:
422 statement = statement.where(User.gender.in_(request.gender))
424 if len(request.hosting_status_filter) > 0:
425 statement = statement.where(
426 User.hosting_status.in_([hostingstatus2sql[status] for status in request.hosting_status_filter])
427 )
428 if len(request.meetup_status_filter) > 0:
429 statement = statement.where(
430 User.meetup_status.in_([meetupstatus2sql[status] for status in request.meetup_status_filter])
431 )
432 if len(request.smoking_location_filter) > 0:
433 statement = statement.where(
434 User.smoking_allowed.in_([smokinglocation2sql[loc] for loc in request.smoking_location_filter])
435 )
436 if len(request.sleeping_arrangement_filter) > 0:
437 statement = statement.where(
438 User.sleeping_arrangement.in_(
439 [sleepingarrangement2sql[arr] for arr in request.sleeping_arrangement_filter]
440 )
441 )
442 if len(request.parking_details_filter) > 0:
443 statement = statement.where(
444 User.parking_details.in_([parkingdetails2sql[det] for det in request.parking_details_filter])
445 )
446 if request.HasField("profile_completed"):
447 statement = statement.where(User.has_completed_profile == request.profile_completed.value)
448 if request.HasField("guests"):
449 statement = statement.where(User.max_guests >= request.guests.value)
450 if request.HasField("last_minute"):
451 statement = statement.where(User.last_minute == request.last_minute.value)
452 if request.HasField("has_pets"):
453 statement = statement.where(User.has_pets == request.has_pets.value)
454 if request.HasField("accepts_pets"):
455 statement = statement.where(User.accepts_pets == request.accepts_pets.value)
456 if request.HasField("has_kids"):
457 statement = statement.where(User.has_kids == request.has_kids.value)
458 if request.HasField("accepts_kids"):
459 statement = statement.where(User.accepts_kids == request.accepts_kids.value)
460 if request.HasField("has_housemates"):
461 statement = statement.where(User.has_housemates == request.has_housemates.value)
462 if request.HasField("wheelchair_accessible"):
463 statement = statement.where(User.wheelchair_accessible == request.wheelchair_accessible.value)
464 if request.HasField("smokes_at_home"):
465 statement = statement.where(User.smokes_at_home == request.smokes_at_home.value)
466 if request.HasField("drinking_allowed"):
467 statement = statement.where(User.drinking_allowed == request.drinking_allowed.value)
468 if request.HasField("drinks_at_home"):
469 statement = statement.where(User.drinks_at_home == request.drinks_at_home.value)
470 if request.HasField("parking"):
471 statement = statement.where(User.parking == request.parking.value)
472 if request.HasField("camping_ok"):
473 statement = statement.where(User.camping_ok == request.camping_ok.value)
475 if request.HasField("search_in_area"):
476 # EPSG4326 measures distance in decimal degress
477 # we want to check whether two circles overlap, so check if the distance between their centers is less
478 # than the sum of their radii, divided by 111111 m ~= 1 degree (at the equator)
479 search_point = create_coordinate(request.search_in_area.lat, request.search_in_area.lng)
480 statement = statement.where(
481 func.ST_DWithin(
482 # old:
483 # User.geom, search_point, (User.geom_radius + request.search_in_area.radius) / 111111
484 # this is an optimization that speeds up the db queries since it doesn't need to look up the user's geom radius
485 User.geom,
486 search_point,
487 (1000 + request.search_in_area.radius) / 111111,
488 )
489 )
490 if request.HasField("search_in_rectangle"):
491 statement = statement.where(
492 func.ST_Within(
493 User.geom,
494 func.ST_MakeEnvelope(
495 request.search_in_rectangle.lng_min,
496 request.search_in_rectangle.lat_min,
497 request.search_in_rectangle.lng_max,
498 request.search_in_rectangle.lat_max,
499 4326,
500 ),
501 )
502 )
503 if request.HasField("search_in_community_id"):
504 # could do a join here as well, but this is just simpler
505 node = session.execute(select(Node).where(Node.id == request.search_in_community_id)).scalar_one_or_none()
506 if not node:
507 context.abort(grpc.StatusCode.NOT_FOUND, errors.COMMUNITY_NOT_FOUND)
508 statement = statement.where(func.ST_Contains(node.geom, User.geom))
510 if request.only_with_references:
511 statement = statement.join(Reference, Reference.to_user_id == User.id)
513 # TODO:
514 # google.protobuf.StringValue language = 11;
515 # bool friends_only = 13;
516 # google.protobuf.UInt32Value age_min = 14;
517 # google.protobuf.UInt32Value age_max = 15;
519 page_size = min(MAX_PAGINATION_LENGTH, request.page_size or MAX_PAGINATION_LENGTH)
520 next_recommendation_score = float(decrypt_page_token(request.page_token)) if request.page_token else 1e10
522 statement = (
523 statement.where(User.recommendation_score <= next_recommendation_score)
524 .order_by(User.recommendation_score.desc())
525 .limit(page_size + 1)
526 )
527 users = session.execute(statement).scalars().all()
529 return search_pb2.UserSearchRes(
530 results=[
531 search_pb2.Result(
532 rank=1,
533 user=user_model_to_pb(user, session, context),
534 )
535 for user in users[:page_size]
536 ],
537 next_page_token=(
538 encrypt_page_token(str(users[-1].recommendation_score)) if len(users) > page_size else None
539 ),
540 )
542 def EventSearch(self, request, context, session):
543 statement = (
544 select(EventOccurrence).join(Event, Event.id == EventOccurrence.event_id).where(~EventOccurrence.is_deleted)
545 )
547 if request.HasField("query"):
548 if request.query_title_only:
549 statement = statement.where(Event.title.ilike(f"%{request.query.value}%"))
550 else:
551 statement = statement.where(
552 or_(
553 Event.title.ilike(f"%{request.query.value}%"),
554 EventOccurrence.content.ilike(f"%{request.query.value}%"),
555 EventOccurrence.address.ilike(f"%{request.query.value}%"),
556 )
557 )
559 if request.only_online:
560 statement = statement.where(EventOccurrence.geom == None)
561 elif request.only_offline:
562 statement = statement.where(EventOccurrence.geom != None)
564 if request.subscribed or request.attending or request.organizing or request.my_communities:
565 where_ = []
567 if request.subscribed:
568 statement = statement.outerjoin(
569 EventSubscription,
570 and_(EventSubscription.event_id == Event.id, EventSubscription.user_id == context.user_id),
571 )
572 where_.append(EventSubscription.user_id != None)
573 if request.organizing:
574 statement = statement.outerjoin(
575 EventOrganizer,
576 and_(EventOrganizer.event_id == Event.id, EventOrganizer.user_id == context.user_id),
577 )
578 where_.append(EventOrganizer.user_id != None)
579 if request.attending:
580 statement = statement.outerjoin(
581 EventOccurrenceAttendee,
582 and_(
583 EventOccurrenceAttendee.occurrence_id == EventOccurrence.id,
584 EventOccurrenceAttendee.user_id == context.user_id,
585 ),
586 )
587 where_.append(EventOccurrenceAttendee.user_id != None)
588 if request.my_communities:
589 my_communities = (
590 session.execute(
591 select(Node.id)
592 .join(Cluster, Cluster.parent_node_id == Node.id)
593 .join(ClusterSubscription, ClusterSubscription.cluster_id == Cluster.id)
594 .where(ClusterSubscription.user_id == context.user_id)
595 .where(Cluster.is_official_cluster)
596 .order_by(Node.id)
597 .limit(100000)
598 )
599 .scalars()
600 .all()
601 )
602 where_.append(Event.parent_node_id.in_(my_communities))
604 statement = statement.where(or_(*where_))
606 if not request.include_cancelled:
607 statement = statement.where(~EventOccurrence.is_cancelled)
609 if request.HasField("search_in_area"):
610 # EPSG4326 measures distance in decimal degress
611 # we want to check whether two circles overlap, so check if the distance between their centers is less
612 # than the sum of their radii, divided by 111111 m ~= 1 degree (at the equator)
613 search_point = create_coordinate(request.search_in_area.lat, request.search_in_area.lng)
614 statement = statement.where(
615 func.ST_DWithin(
616 # old:
617 # User.geom, search_point, (User.geom_radius + request.search_in_area.radius) / 111111
618 # this is an optimization that speeds up the db queries since it doesn't need to look up the user's geom radius
619 EventOccurrence.geom,
620 search_point,
621 (1000 + request.search_in_area.radius) / 111111,
622 )
623 )
624 if request.HasField("search_in_rectangle"):
625 statement = statement.where(
626 func.ST_Within(
627 EventOccurrence.geom,
628 func.ST_MakeEnvelope(
629 request.search_in_rectangle.lng_min,
630 request.search_in_rectangle.lat_min,
631 request.search_in_rectangle.lng_max,
632 request.search_in_rectangle.lat_max,
633 4326,
634 ),
635 )
636 )
637 if request.HasField("search_in_community_id"):
638 # could do a join here as well, but this is just simpler
639 node = session.execute(select(Node).where(Node.id == request.search_in_community_id)).scalar_one_or_none()
640 if not node:
641 context.abort(grpc.StatusCode.NOT_FOUND, errors.COMMUNITY_NOT_FOUND)
642 statement = statement.where(func.ST_Contains(node.geom, EventOccurrence.geom))
644 if request.HasField("after"):
645 statement = statement.where(EventOccurrence.start_time > to_aware_datetime(request.after))
646 if request.HasField("before"):
647 statement = statement.where(EventOccurrence.end_time < to_aware_datetime(request.before))
649 page_size = min(MAX_PAGINATION_LENGTH, request.page_size or MAX_PAGINATION_LENGTH)
650 # the page token is a unix timestamp of where we left off
651 page_token = (
652 dt_from_millis(int(request.page_token)) if request.page_token and not request.page_number else now()
653 )
654 page_number = request.page_number or 1
655 # Calculate the offset for pagination
656 offset = (page_number - 1) * page_size
658 if not request.past:
659 statement = statement.where(EventOccurrence.end_time > page_token - timedelta(seconds=1)).order_by(
660 EventOccurrence.start_time.asc()
661 )
662 else:
663 statement = statement.where(EventOccurrence.end_time < page_token + timedelta(seconds=1)).order_by(
664 EventOccurrence.start_time.desc()
665 )
667 total_items = session.execute(select(func.count()).select_from(statement.subquery())).scalar()
668 # Apply pagination by page number
669 statement = statement.offset(offset).limit(page_size) if request.page_number else statement.limit(page_size + 1)
670 occurrences = session.execute(statement).scalars().all()
672 return search_pb2.EventSearchRes(
673 events=[event_to_pb(session, occurrence, context) for occurrence in occurrences[:page_size]],
674 next_page_token=(str(millis_from_dt(occurrences[-1].end_time)) if len(occurrences) > page_size else None),
675 total_items=total_items,
676 )