cyberyan_linkedin_search/backend/apps/search/backends/postgres.py

"""
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