289 lines
12 KiB
Python
289 lines
12 KiB
Python
"""
|
|
PostgreSQL full-text search backend.
|
|
|
|
Why it exists
|
|
-------------
|
|
Elasticsearch is the primary engine, but a search API that hard-fails without
|
|
it is fragile: the test suite would need a cluster, a reviewer would need one to
|
|
run the project, and an Elasticsearch outage would take the whole product down.
|
|
Setting ``SEARCH_BACKEND=postgres`` switches to this implementation, which is
|
|
built on the ``tsvector`` maintained by the importer plus a trigram similarity
|
|
pass for typo tolerance.
|
|
|
|
It is intentionally *not* a re-implementation of Elasticsearch: ranking is
|
|
coarser and there is no stemming beyond what the ``english`` text search
|
|
configuration provides. For a 300 profile dataset the difference is barely
|
|
visible, which is exactly why it is a good fallback.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import re
|
|
import time
|
|
from collections import Counter
|
|
|
|
from django.contrib.postgres.search import SearchQuery, SearchRank, TrigramSimilarity
|
|
from django.db.models import Case, F, FloatField, Q, QuerySet, Value, When
|
|
from django.db.models.functions import Greatest
|
|
|
|
from apps.profiles.models import Profile, Skill
|
|
from apps.search.backends.base import SearchBackend
|
|
from apps.search.types import (
|
|
SORT_EXPERIENCE_ASC,
|
|
SORT_EXPERIENCE_DESC,
|
|
SORT_NAME,
|
|
FacetBucket,
|
|
SearchHit,
|
|
SearchRequest,
|
|
SearchResult,
|
|
)
|
|
|
|
_TOKEN_RE = re.compile(r"[\w']+", re.UNICODE)
|
|
_FACET_LIMIT = 25
|
|
_MIN_TRIGRAM_SIMILARITY = 0.3
|
|
|
|
#: Maps a facet (as returned to the frontend) to the request attribute that
|
|
#: selects values *within that same facet*. Used to exclude a facet's own
|
|
#: filter from its own count - see ``_facets``.
|
|
_FACET_TO_FILTER_ATTR = {
|
|
"skills": "skills",
|
|
"job_title_role": "roles",
|
|
"seniority": "seniority",
|
|
"country": "countries",
|
|
"industry": "industries",
|
|
"company": "companies",
|
|
}
|
|
|
|
|
|
class PostgresBackend(SearchBackend):
|
|
name = "postgres"
|
|
|
|
def search(self, request: SearchRequest) -> SearchResult:
|
|
started = time.perf_counter()
|
|
keyword_queryset = self._apply_keyword(Profile.objects.all(), request.q)
|
|
queryset = self._apply_filters(keyword_queryset, request)
|
|
total = queryset.count()
|
|
|
|
page = (
|
|
self._apply_sort(queryset, request)
|
|
.select_related("current_company")
|
|
.prefetch_related("profile_skills__skill")[
|
|
request.offset : request.offset + request.page_size
|
|
]
|
|
)
|
|
hits = [self._to_hit(profile, request.q) for profile in page]
|
|
facets = self._facets(keyword_queryset, request) if request.include_facets else {}
|
|
|
|
return SearchResult(
|
|
total=total,
|
|
hits=hits,
|
|
facets=facets,
|
|
took_ms=int((time.perf_counter() - started) * 1000),
|
|
backend=self.name,
|
|
)
|
|
|
|
def suggest(self, prefix: str, limit: int = 8) -> list[str]:
|
|
if not prefix:
|
|
return []
|
|
return list(
|
|
Profile.objects.filter(
|
|
Q(full_name__istartswith=prefix) | Q(full_name__icontains=f" {prefix}")
|
|
)
|
|
.order_by("full_name")
|
|
.values_list("full_name", flat=True)[:limit]
|
|
)
|
|
|
|
def ping(self) -> bool:
|
|
Profile.objects.exists()
|
|
return True
|
|
|
|
# -- Query construction -------------------------------------------------
|
|
@staticmethod
|
|
def _apply_filters(
|
|
queryset: QuerySet[Profile], request: SearchRequest, *, exclude: str | None = None
|
|
) -> QuerySet[Profile]:
|
|
"""
|
|
Apply every active filter except ``exclude``.
|
|
|
|
``exclude`` is how a facet's own count query avoids narrowing itself:
|
|
pass e.g. ``exclude="roles"`` when computing the job-family facet so
|
|
ticking "Engineering" does not make "Human resources" disappear from
|
|
that same list. The main hits queryset calls this with no exclusion.
|
|
"""
|
|
if request.skills and exclude != "skills":
|
|
# Every selected skill widens the result set (OR), mirroring the
|
|
# ``terms`` filter used by the Elasticsearch backend.
|
|
queryset = queryset.filter(
|
|
skills__slug__in=[Skill.make_slug(s) for s in request.skills]
|
|
)
|
|
if request.roles and exclude != "roles":
|
|
queryset = queryset.filter(job_title_role__in=request.roles)
|
|
if request.seniority and exclude != "seniority":
|
|
queryset = queryset.filter(seniority_levels__overlap=list(request.seniority))
|
|
if request.countries and exclude != "countries":
|
|
queryset = queryset.filter(location_country__in=request.countries)
|
|
if request.industries and exclude != "industries":
|
|
queryset = queryset.filter(industry__in=request.industries)
|
|
if request.companies and exclude != "companies":
|
|
queryset = queryset.filter(current_company__name__in=request.companies)
|
|
# The experience range is not a facet with its own checkboxes, so it
|
|
# always applies - it is never the thing being excluded.
|
|
if request.min_experience is not None:
|
|
queryset = queryset.filter(years_experience__gte=request.min_experience)
|
|
if request.max_experience is not None:
|
|
queryset = queryset.filter(years_experience__lte=request.max_experience)
|
|
return queryset.distinct()
|
|
|
|
@staticmethod
|
|
def _apply_keyword(queryset: QuerySet[Profile], term: str) -> QuerySet[Profile]:
|
|
if not term:
|
|
return queryset.annotate(rank=Value(0.0, output_field=FloatField()))
|
|
|
|
# ``websearch`` accepts the syntax users already know from search
|
|
# engines: quoted phrases, ``or``, and ``-excluded``.
|
|
query = SearchQuery(term, search_type="websearch", config="english")
|
|
return (
|
|
queryset.annotate(
|
|
text_rank=SearchRank(F("search_vector"), query),
|
|
name_similarity=TrigramSimilarity("full_name", term),
|
|
)
|
|
.annotate(
|
|
rank=Greatest(
|
|
F("text_rank"),
|
|
# A close name match is worth as much as a strong text hit,
|
|
# which is what makes "jospeh" still find "joseph".
|
|
Case(
|
|
When(
|
|
name_similarity__gte=_MIN_TRIGRAM_SIMILARITY, then=F("name_similarity")
|
|
),
|
|
default=Value(0.0),
|
|
output_field=FloatField(),
|
|
),
|
|
output_field=FloatField(),
|
|
)
|
|
)
|
|
.filter(Q(search_vector=query) | Q(name_similarity__gte=_MIN_TRIGRAM_SIMILARITY))
|
|
)
|
|
|
|
@staticmethod
|
|
def _apply_sort(queryset: QuerySet[Profile], request: SearchRequest) -> QuerySet[Profile]:
|
|
if request.sort == SORT_NAME:
|
|
return queryset.order_by("full_name", "id")
|
|
if request.sort == SORT_EXPERIENCE_DESC:
|
|
return queryset.order_by(F("years_experience").desc(nulls_last=True), "id")
|
|
if request.sort == SORT_EXPERIENCE_ASC:
|
|
return queryset.order_by(F("years_experience").asc(nulls_last=True), "id")
|
|
return queryset.order_by("-rank", "full_name", "id")
|
|
|
|
# -- Facets -------------------------------------------------------------
|
|
def _facets(
|
|
self, keyword_queryset: QuerySet[Profile], request: SearchRequest
|
|
) -> dict[str, list[FacetBucket]]:
|
|
"""
|
|
One count per facet, each computed on a queryset that excludes that
|
|
facet's *own* filter.
|
|
|
|
This is what keeps a multi-select facet fully populated once one of
|
|
its values is ticked: the job-family facet must still list "Human
|
|
resources" after "Engineering" is selected, while every *other* facet
|
|
(country, seniority, ...) does narrow to the engineering-only result
|
|
set. Six independent queries instead of one is a deliberate
|
|
trade-off - see the module docstring on why this backend does not aim
|
|
to match Elasticsearch's scale, only its contract.
|
|
"""
|
|
return {
|
|
name: self._facet_counts(
|
|
name, self._apply_filters(keyword_queryset, request, exclude=attribute)
|
|
)
|
|
for name, attribute in _FACET_TO_FILTER_ATTR.items()
|
|
}
|
|
|
|
@staticmethod
|
|
def _facet_counts(facet_name: str, queryset: QuerySet[Profile]) -> list[FacetBucket]:
|
|
ids = list(queryset.values_list("id", flat=True))
|
|
if not ids:
|
|
return []
|
|
|
|
if facet_name == "skills":
|
|
counter = Counter(
|
|
Skill.objects.filter(profile_skills__profile_id__in=ids).values_list(
|
|
"name", flat=True
|
|
)
|
|
)
|
|
elif facet_name == "seniority":
|
|
counter = Counter()
|
|
for levels in Profile.objects.filter(id__in=ids).values_list(
|
|
"seniority_levels", flat=True
|
|
):
|
|
for level in levels or []:
|
|
counter[level] += 1
|
|
else:
|
|
column = {
|
|
"job_title_role": "job_title_role",
|
|
"country": "location_country",
|
|
"industry": "industry",
|
|
"company": "current_company__name",
|
|
}[facet_name]
|
|
counter = Counter(
|
|
value or "" for value in Profile.objects.filter(id__in=ids).values_list(
|
|
column, flat=True
|
|
)
|
|
)
|
|
|
|
return [
|
|
FacetBucket(key=key, count=count)
|
|
for key, count in counter.most_common(_FACET_LIMIT)
|
|
if key
|
|
]
|
|
|
|
# -- Result mapping -----------------------------------------------------
|
|
@classmethod
|
|
def _to_hit(cls, profile: Profile, term: str) -> SearchHit:
|
|
skills = [link.skill.name for link in profile.profile_skills.all()[:12]]
|
|
return SearchHit(
|
|
id=profile.id,
|
|
score=round(float(getattr(profile, "rank", 0.0) or 0.0), 4),
|
|
full_name=profile.full_name,
|
|
headline=profile.headline,
|
|
job_title=profile.job_title,
|
|
job_title_role=profile.job_title_role,
|
|
seniority_levels=profile.seniority_levels,
|
|
company=profile.current_company.name if profile.current_company_id else "",
|
|
industry=profile.industry,
|
|
location_name=profile.location_name,
|
|
location_country=profile.location_country,
|
|
years_experience=profile.years_experience,
|
|
linkedin_url=profile.linkedin_url,
|
|
skills=skills,
|
|
highlight=cls._highlight(term, profile, skills),
|
|
)
|
|
|
|
@staticmethod
|
|
def _highlight(term: str, profile: Profile, skills: list[str]) -> dict[str, list[str]]:
|
|
"""
|
|
Approximate Elasticsearch's highlighter.
|
|
|
|
PostgreSQL offers ``ts_headline`` but it needs an extra round trip per
|
|
row; for the fields shown on a result card a simple token wrap gives the
|
|
same visual result at no query cost.
|
|
"""
|
|
tokens = [token for token in _TOKEN_RE.findall(term.lower()) if len(token) > 1]
|
|
if not tokens:
|
|
return {}
|
|
pattern = re.compile("|".join(re.escape(token) for token in tokens), re.IGNORECASE)
|
|
|
|
highlight: dict[str, list[str]] = {}
|
|
for field_name, value in (
|
|
("full_name", profile.full_name),
|
|
("job_title", profile.job_title),
|
|
):
|
|
if value and pattern.search(value):
|
|
highlight[field_name] = [pattern.sub(lambda m: f"<mark>{m.group(0)}</mark>", value)]
|
|
matched_skills = [skill for skill in skills if pattern.search(skill)]
|
|
if matched_skills:
|
|
highlight["skills"] = [
|
|
pattern.sub(lambda m: f"<mark>{m.group(0)}</mark>", skill)
|
|
for skill in matched_skills[:3]
|
|
]
|
|
return highlight
|