Как подключить Elasticsearch для Django REST Framework

Этот блог начнётся с непрофильной для него темы. Но задача интересная, и возможно кому-то мой способ поможет.

Фото Jametlene Reskp / Unsplash

Ко мне обратился заказчик с просьбой усовершенствовать поиск интернет-магазина. Проект основан на Django и имеет REST API (с помощью Django REST framework) для работы с клиентом, для поиска использовался полнотекстовый поиск PostgreSQL. Т.к. база товаров была большая, алгоритмы поиска планировались сложными, а работа API быстрой, то выбор пал на Elasticsearch. И как оказалось, добавить его в проект можно довольно просто.

Установка Elasticsearch

Проект использует Docker, поэтому вначале добавляем новый контейнер для Elasticsearch в Docker Compose. Вы разумеется, можете запустить его отдельно.

elasticsearch:  
    image: elasticsearch:7.17.8
    environment:
        - discovery.type=single-node
        - cluster.name=es-docker
        - node.name=node1
        - bootstrap.memory_lock=true
        - "ES_JAVA_OPTS=-Xms4g -Xmx4g"
    deploy:
        resources:
            limits:
                memory: 8G # Ограничили количество оперативной памяти
    restart: always
    ports:
        - 9200:9200
        - 9300:9300
    volumes:
        - esdata:/usr/share/elasticsearch/data # Volume для хранения данных
    depends_on:
        - server # Запускаем контейнер после запуска нашего API (Django)
    networks: 
        - prnetwork # Указываем сеть для работы контейнера

Т.к. сервер имеет ограничение по количеству свободной оперативной памяти, то пришлось ограничить его контейнер 8 ГБ. И чтобы не терять наши данные при пересоздании контейнера, сохраняем их в volume.

Настройка сервера

Для настройки взаимодействия с поисковиком из Django REST Framework будем использовать django-elasticsearch-dsl-drf. Этот пакет позволяет производить детальную настройку, и настраивать API под наши нужды. Поэтому добавляем в проект следующие пакеты:

elasticsearch==7.17.8
elasticsearch-dsl==7.4.0
django-elasticsearch-dsl==7.2.2
django-elasticsearch-dsl-drf==0.22.5

В файле settings.py записываем пакеты в INSTALLED_APPS, и указываем хост для подключения:

INSTALLED_APPS = [
    # .........
    'rest_framework', # REST framework
    'django_elasticsearch_dsl', # Интеграция с Elasticsearch
    'django_elasticsearch_dsl_drf', # Пакет для работы с API
    # .........
]

ELASTICSEARCH_DSL = {
    'default': {
        'hosts': os.environ.get("ELASTICSEARCH_HOST")
    },
}

В нашем случае работа осуществлялась в пределах одной Docker сети, поэтому в .env файл добавили:

ELASTICSEARCH_HOST=http://elasticsearch:9200

Работа с моделями

Мы исходим из ситуации, что у вас уже имеется ваша модель, поэтому процесс её создания я упускаю. Для примера, приложу свой вариант модели, который я упростил для визуальной понятности:

class Product(models.Model):
    """Модель продукта"""
    name = models.CharField(max_length=255, db_index=True, verbose_name="Название")
    slug = models.SlugField(max_length=255, db_index=True, unique=True, verbose_name="Url адрес")
    barcode = models.CharField(max_length=300, db_index=True, blank=True, null=True, verbose_name="Штрихкод")
    brand = models.ForeignKey(Brand, on_delete=models.CASCADE, related_name='products', verbose_name="Производитель")
    category = models.ForeignKey(Category, on_delete=models.CASCADE, related_name='products', verbose_name="Категория")
    price = models.IntegerField(default=0, blank=True, null=True, verbose_name="Цена")
    visible = models.BooleanField(default=True, verbose_name="Видимость")

    class Meta:
        ordering = ['name']
        verbose_name = 'product'
        verbose_name_plural = 'Products'

    def __str__(self):
        return self.name
    
    @property
    def get_brand(self):
        return self.brand.name
    
    @property
    def get_category(self):
        return self.category.name

Методы get_brand() и get_category() нужны для того, чтобы мы могли обработать название брендов и категорий, как простое текстовое поле.

Теперь необходимо создать документы (аналог моделей для поиска) для работы с моделями товаров. Для этого в папке с нашим API создаём файл documents.py, подключаем необходимые библиотеки и описываем наши сущности.

from django_elasticsearch_dsl import Document, fields
from elasticsearch_dsl import analyzer, Index
from django_elasticsearch_dsl.registries import registry
from shop.models import Product

product = Index('products') # Имя индекса в Elasticsearch

product.settings(
    number_of_shards=1,
    number_of_replicas=1
) # Расшифровку этой настройки лучше почитать в документации Elasticsearch
   
html_strip = analyzer(
    'html_strip',
    tokenizer="standard",
    filter=["lowercase", "stop", "snowball"],
    char_filter=["html_strip"]
) # Анализатор текстовых полей
        
@registry.register_document
@product.document
class ProductDocument(Document):
    """Документ продукта"""
    
    id = fields.IntegerField(attr='id')
    
    name = fields.TextField(
        analyzer=html_strip,
        fields={
            'raw': fields.TextField(analyzer='keyword'),
        }
    )
    
    barcode = fields.TextField(
        analyzer=html_strip,
        fields={
            'raw': fields.TextField(analyzer='keyword'),
        }
    )
    
    brand = fields.TextField(
        attr='get_brand', # Результат метода get_catget_brandegory()
        analyzer=html_strip,
        fields={
            'raw': fields.TextField(analyzer='keyword'),
        }
    )
    
    category = fields.TextField(
        attr='get_category', # Результат метода get_category()
        analyzer=html_strip,
        fields={
            'raw': fields.TextField(analyzer='keyword'),
        }
    )

    slug = fields.FileField(attr='slug')
    price = fields.IntegerField(attr='price')
    visible = fields.BooleanField(attr='visible')
    
    class Django:
        model = Product # Наша Django модель

Я не буду описывать процесс создания документа т.к. документация пакета очень хорошая, поэтому лучше сразу прочитать там. Ещё вам может понадобиться документация по Elasticsearch, чтобы правильно настроить алгоритмы.

Теперь, когда документы готовы, можно индексировать данные:

python manage.py search_index --create -f
python manage.py search_index --populate -f

Если вам необходимо удалить все данные, то используйте команду python manage.py search_index --delete

Получаем данные по REST API

Нам осталось создать сериалайзер, описать алгоритм поиска и всё будет готово. Для простоты будем выводить только 4 поля:

class ProductDocumentSerializer(DocumentSerializer):
    """Сериалайзер товара для поиска"""

    class Meta:
        document = ProductDocument

        fields = (
            'id',
            'name',
            'slug',
            'price',
        )

Для это проекта нужно было изменить пагинацию, чтобы вместе с результатами в определённом формате выводилась максимальная и минимальная цена. Поэтому дополнительно создаём собственный класс пагинации:

from django_elasticsearch_dsl_drf.pagination import PageNumberPagination

class SearchPagination(PageNumberPagination):
    """Настройка пагинации для поиска товаров"""

    def get_paginated_response_context(self, data):
        min_price = 0
        max_price = 0
        
        __facets = self.get_facets()
        if __facets is not None:
            if __facets['min']['value'] != None:
                min_price = int(__facets['min']['value'])
            if __facets['max']['value'] != None:
                max_price = int(__facets['max']['value'])

        # Переопределяем ответ под наши нужды
        return [
            ('page', self.page.number),
            ('count', self.page.count),
            ('min_price', min_price),
            ('max_price', max_price),
            ('results', data)
        ]

А теперь дописываем views.py, и сразу настроим, чтобы показывались только товары с включенным отображением:

from elasticsearch_dsl import Q
from django_elasticsearch_dsl_drf.viewsets import BaseDocumentViewSet
from django_elasticsearch_dsl_drf.constants import ( LOOKUP_FILTER_RANGE )
from django_elasticsearch_dsl_drf.filter_backends import (
    FilteringFilterBackend,
    OrderingFilterBackend,
    DefaultOrderingFilterBackend,
    SearchFilterBackend,
)
from .documents import ProductDocument

class SearchProductsView(BaseDocumentViewSet):
    """Поиск товаров"""
    document = ProductDocument # Указываем документ
    serializer_class = ProductDocumentSerializer
    pagination_class = SearchPagination
    lookup_field = 'id'
    filter_backends = [
        FilteringFilterBackend,
        OrderingFilterBackend,
        DefaultOrderingFilterBackend,
        SearchFilterBackend,
    ]

    # Определяем поля по которым будет осуществляться поиск
    # И устанавливаем правила приоритета и погрешности
    search_fields = {
        'name': { 'fuzziness': 'AUTO', 'boost': 2 },
        'barcode': { 'boost': 3 },
        'brand': { 'boost': 1 },
        'category': { 'boost': 1 },
    }
    
    # Определяем поля для фильтрации
    filter_fields = {
        'price': {
            'field': 'price',
            'lookups': [
                LOOKUP_FILTER_RANGE,
            ],
        },
    }
    
    # Определяем поля для сортировки
    ordering_fields = {
        'price': 'price',
    }
    
    # Сортировка по умолчанию
    ordering = ['_score']
    
    def get_queryset(self):
        """Оставляем только видимые товары"""
        queryset = self.search.query(Q('term', visible='true'))
        queryset.model = self.document.Django.model
        return queryset
        
    def paginate_queryset(self, queryset): 
        """Высчитываем максимальную и минимальную цену товаров"""       
        queryset.aggs.metric('max', 'max', field='price')
        queryset.aggs.metric('min', 'min', field='price')
        if self.paginator is None: return None
        return self.paginator.paginate_queryset(queryset, self.request, view=self)

Сортировка _score означает, что результаты будут в порядки релевантности запросу.

Добавляем ссылку в urls.py нашего API:

router.register(r'search', SearchProductsView, basename='search_products')

И всё, поиск работает. Спасибо за прочтение!