253 lines
9.9 KiB
Python
253 lines
9.9 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
|
||
|
|
|
||
|
|
|
||
|
|
class PostgresBackend(SearchBackend):
|
||
|
|
name = "postgres"
|
||
|
|
|
||
|
|
def search(self, request: SearchRequest) -> SearchResult:
|
||
|
|
started = time.perf_counter()
|
||
|
|
queryset = self._apply_filters(Profile.objects.all(), request)
|
||
|
|
queryset = self._apply_keyword(queryset, request.q)
|
||
|
|
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(queryset) 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) -> QuerySet[Profile]:
|
||
|
|
if request.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:
|
||
|
|
queryset = queryset.filter(job_title_role__in=request.roles)
|
||
|
|
if request.seniority:
|
||
|
|
queryset = queryset.filter(seniority_levels__overlap=list(request.seniority))
|
||
|
|
if request.countries:
|
||
|
|
queryset = queryset.filter(location_country__in=request.countries)
|
||
|
|
if request.industries:
|
||
|
|
queryset = queryset.filter(industry__in=request.industries)
|
||
|
|
if request.companies:
|
||
|
|
queryset = queryset.filter(current_company__name__in=request.companies)
|
||
|
|
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, queryset: QuerySet[Profile]) -> dict[str, list[FacetBucket]]:
|
||
|
|
ids = list(queryset.values_list("id", flat=True))
|
||
|
|
if not ids:
|
||
|
|
return {
|
||
|
|
name: []
|
||
|
|
for name in (
|
||
|
|
"skills",
|
||
|
|
"job_title_role",
|
||
|
|
"seniority",
|
||
|
|
"country",
|
||
|
|
"industry",
|
||
|
|
"company",
|
||
|
|
)
|
||
|
|
}
|
||
|
|
|
||
|
|
skills = Skill.objects.filter(profile_skills__profile_id__in=ids).values_list(
|
||
|
|
"name", flat=True
|
||
|
|
)
|
||
|
|
rows = Profile.objects.filter(id__in=ids).values(
|
||
|
|
"job_title_role", "location_country", "industry", "current_company__name",
|
||
|
|
"seniority_levels",
|
||
|
|
) # fmt: skip
|
||
|
|
|
||
|
|
counters: dict[str, Counter] = {
|
||
|
|
"skills": Counter(skills),
|
||
|
|
"job_title_role": Counter(),
|
||
|
|
"seniority": Counter(),
|
||
|
|
"country": Counter(),
|
||
|
|
"industry": Counter(),
|
||
|
|
"company": Counter(),
|
||
|
|
}
|
||
|
|
for row in rows:
|
||
|
|
counters["job_title_role"][row["job_title_role"]] += 1
|
||
|
|
counters["country"][row["location_country"]] += 1
|
||
|
|
counters["industry"][row["industry"]] += 1
|
||
|
|
counters["company"][row["current_company__name"] or ""] += 1
|
||
|
|
for level in row["seniority_levels"] or []:
|
||
|
|
counters["seniority"][level] += 1
|
||
|
|
|
||
|
|
return {
|
||
|
|
name: [
|
||
|
|
FacetBucket(key=key, count=count)
|
||
|
|
for key, count in counter.most_common(_FACET_LIMIT)
|
||
|
|
if key
|
||
|
|
]
|
||
|
|
for name, counter in counters.items()
|
||
|
|
}
|
||
|
|
|
||
|
|
# -- 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
|