schon/engine/core/elasticsearch/documents.py

from typing import Any

from django.db.models import Model, QuerySet
from django_elasticsearch_dsl import Document, fields
from django_elasticsearch_dsl.registries import registry
from health_check.db.models import TestModel

from engine.core.elasticsearch import (
    COMMON_ANALYSIS,
    ActiveOnlyMixin,
    add_multilang_fields,
)
from engine.core.models import Brand, Category, Product


class BaseDocument(Document):
    name = fields.TextField(
        attr="name",
        analyzer="standard",
        fields={
            "raw": fields.KeywordField(ignore_above=256),
            "ngram": fields.TextField(
                analyzer="name_ngram", search_analyzer="icu_query"
            ),
            "phonetic": fields.TextField(analyzer="name_phonetic"),
            "auto": fields.TextField(
                analyzer="autocomplete", search_analyzer="autocomplete_search"
            ),
            "translit": fields.TextField(
                analyzer="translit_index", search_analyzer="translit_query"
            ),
            "ci": fields.TextField(analyzer="name_exact", search_analyzer="name_exact"),
        },
    )
    description = fields.TextField(
        attr="description",
        analyzer="standard",
        fields={
            "raw": fields.KeywordField(ignore_above=256),
            "ngram": fields.TextField(
                analyzer="name_ngram", search_analyzer="icu_query"
            ),
            "phonetic": fields.TextField(analyzer="name_phonetic"),
            "auto": fields.TextField(
                analyzer="autocomplete", search_analyzer="autocomplete_search"
            ),
            "translit": fields.TextField(
                analyzer="translit_index", search_analyzer="translit_query"
            ),
        },
    )
    slug = fields.KeywordField(attr="slug")

    class Index:
        settings = {
            "number_of_shards": 1,
            "number_of_replicas": 0,
            "analysis": COMMON_ANALYSIS,
            "index": {"max_ngram_diff": 20},
        }

    def prepare_name(self, instance: Model) -> str:
        return getattr(instance, "name", "") or ""

    def prepare_description(self, instance: Model) -> str:
        return getattr(instance, "description", "") or ""


class ProductDocument(ActiveOnlyMixin, BaseDocument):
    rating = fields.FloatField(attr="rating")
    total_orders = fields.IntegerField(attr="total_orders")
    personal_orders_only = fields.BooleanField(attr="personal_orders_only")
    brand_priority = fields.IntegerField(
        attr="brand.priority",
        index=True,
        fields={"raw": fields.KeywordField()},
    )
    category_priority = fields.IntegerField(
        attr="category.priority",
        index=True,
        fields={"raw": fields.KeywordField()},
    )

    brand_name = fields.TextField(
        attr="brand.name",
        analyzer="standard",
        fields={
            "raw": fields.KeywordField(ignore_above=256),
            "ngram": fields.TextField(
                analyzer="name_ngram", search_analyzer="icu_query"
            ),
            "phonetic": fields.TextField(analyzer="name_phonetic"),
            "auto": fields.TextField(
                analyzer="autocomplete", search_analyzer="autocomplete_search"
            ),
            "translit": fields.TextField(
                analyzer="translit_index", search_analyzer="translit_query"
            ),
        },
    )
    category_name = fields.TextField(
        attr="category.name",
        analyzer="standard",
        fields={
            "raw": fields.KeywordField(ignore_above=256),
            "ngram": fields.TextField(
                analyzer="name_ngram", search_analyzer="icu_query"
            ),
            "phonetic": fields.TextField(analyzer="name_phonetic"),
            "auto": fields.TextField(
                analyzer="autocomplete", search_analyzer="autocomplete_search"
            ),
            "translit": fields.TextField(
                analyzer="translit_index", search_analyzer="translit_query"
            ),
        },
    )

    sku = fields.KeywordField(
        attr="sku",
        normalizer="lc_norm",
        fields={
            "raw": fields.KeywordField(normalizer="lc_norm"),
            "ngram": fields.TextField(
                analyzer="name_ngram", search_analyzer="icu_query"
            ),
            "auto": fields.TextField(
                analyzer="autocomplete", search_analyzer="autocomplete_search"
            ),
        },
    )

    partnumber = fields.KeywordField(
        attr="partnumber",
        normalizer="lc_norm",
        fields={
            "raw": fields.KeywordField(normalizer="lc_norm"),
            "ngram": fields.TextField(
                analyzer="name_ngram", search_analyzer="icu_query"
            ),
            "auto": fields.TextField(
                analyzer="autocomplete", search_analyzer="autocomplete_search"
            ),
        },
    )

    def get_queryset(self) -> QuerySet[Product]:
        return (
            super()
            .get_queryset()
            .filter(
                brand__is_active=True,
                category__is_active=True,
                stocks__vendor__is_active=True,
            )
        )

    class Index(BaseDocument.Index):
        name = "products"

    class Django:
        model = Product
        fields = ["uuid"]


add_multilang_fields(ProductDocument)
registry.register_document(ProductDocument)


class CategoryDocument(ActiveOnlyMixin, BaseDocument):
    priority = fields.IntegerField(attr="priority")

    class Index(BaseDocument.Index):
        name = "categories"

    class Django:
        model = Category
        fields = ["uuid"]


add_multilang_fields(CategoryDocument)
registry.register_document(CategoryDocument)


class BrandDocument(ActiveOnlyMixin, BaseDocument):
    priority = fields.IntegerField(attr="priority")

    class Index(BaseDocument.Index):
        name = "brands"

    class Django:
        model = Brand
        fields = ["uuid"]


add_multilang_fields(BrandDocument)
registry.register_document(BrandDocument)


class TestModelDocument(Document):
    class Index:
        name = "testmodels"

    class Django:
        model = TestModel
        fields = ["title"]
        ignore_signals = True
        related_models: list[Any] = []
        auto_refresh = False


registry.register_document(TestModelDocument)