Кастомизация Django ORM на примере ZomboDB

Часто при работе с Django и PostgreSQL возникает необходимость в дополнительных расширениях для базы данных. И если например с hstore или PostGIS (благодаря GeoDjango) всё достаточно удобно, то c более редкими расширениями — вроде pgRouting, ZomboDB и пр. — приходится либо писать на RawSQL, либо кастомизировать Django ORM. Чем я предлагаю, в данной статье, и заняться, используя в качестве примера ZomboDB и его getting started tutorial. И заодно рассмотрим как можно подключить ZomboDB к проекту на Django.

У PostgreSQL есть свой полнотекстовый поиск и работает он, судя по последним бенчмаркам, довольно быстро. Но его возможности именно в поиске всё ещё оставляют желать лучшего. Вследствие чего без решений на базе Lucene — ElasticSearch, например, — приходится туго. ElasticSearch внутри имеет свою БД, по которой проводит поиск. Основное решение на текущий момент — это ручное управление консистентностью данных между PostgreSQL и ElasticSearch с помощью сигналов или ручных функций обратного вызова.

ZomboDB — это расширение, которое реализует собственный тип индекса, превращая значение таблицы в указатель на ElasticSearch, что позволяет проводить полнотекстовый поиск по таблице, используя ElasticSearch DSL как часть синтаксиса SQL.

На момент написания статьи поиск по сети к результатам не привел. Из статей на Хабре про ZomboDB только одна. Статьи по интеграции ZomboDB и Django отсутствуют.

В описании ZomboDB сказано, что обращения в Elasticsearch идут через RESTful API, поэтому производительность вызывает сомнения, но сейчас мы ее касаться не будем. Также вопросов корректного удаления ZomboDB без потери данных.

Далее все тесты будем проводить в Docker, поэтому соберем небольшой docker-compose файл

version: '3'
services:
  postgres:
    build: docker/postgres
    environment:
      - POSTGRES_USER=django
      - POSTGRES_PASSWORD=123456
      - POSTGRES_DB=zombodb
      - PGDATA=/home/postgresql/data
    ports:
      - 5432:5432

  # sudo sysctl -w vm.max_map_count=262144
  elasticsearch:
    image: elasticsearch:6.5.4
    environment:
      - cluster.name=zombodb
      - bootstrap.memory_lock=true
      - ES_JAVA_OPTS=-Xms512m -Xmx512m
    ulimits:
      memlock:
        soft: -1
        hard: -1
    ports:
      - 9200:9200

  django:
    build: docker/python
    command: python3 manage.py runserver 0.0.0.0:8000
    volumes:
      - ./:/home/
    ports:
      - 8000:8000
    depends_on:
      - postgres
      - elasticsearch

Последняя версия ZomboDB работает максимум с 10-ой версией Postgres и из зависимостей требует curl (полагаю, чтобы делать запросы в ElasticSearch).

FROM postgres:10

WORKDIR /home/
RUN apt-get -y update && apt-get -y install curl
ADD https://www.zombodb.com/releases/v10-1.0.3/zombodb_jessie_pg10-10-1.0.3_amd64.deb ./
RUN dpkg -i zombodb_jessie_pg10-10-1.0.3_amd64.deb
RUN rm zombodb_jessie_pg10-10-1.0.3_amd64.deb
RUN apt-get -y clean

Контейнер для Django типичный. В него мы поставим только последние версии Django и psycopg2.

FROM python:stretch

WORKDIR /home/
RUN pip3 install --no-cache-dir django psycopg2-binary

ElasticSearch в Linux не стартует с базовыми настройками vm.max_map_count, поэтому нам придется их немного увеличить (кто знает как это автоматизировать через docker — отпишитесь в комментариях).

sudo sysctl -w vm.max_map_count=262144

Итак, тестовое окружение готово. Можно переходить к проекту на Django. Целиком я его приводить не буду, желающие могут посмотреть его в репозитории на GitLab. Остановлюсь только на критичных моментах.

Первое, что нам нужно сделать, это подключить ZomboDB как extension в PostgreSQL. Можно, конечно, подключиться к базе и включить расширение через SQL CREATE EXTENSION zombodb;. Можно даже для этого использовать docker-entrypoint-initdb.d hook в официальном контейнере для Postgres. Но раз у нас Django, то и пойдем его путем.
После создания проекта и создания первой миграции добавим в нее подключение расширения.

from django.db import migrations, models
from django.contrib.postgres.operations import CreateExtension

class Migration(migrations.Migration):
    initial = True
    dependencies = [
    ]
    operations = [
        CreateExtension('zombodb'),
    ]

Во-вторых, нам нужна модель, которая будет описывать тестовую таблицу. Для этого нам необходимо поле, которое бы работало с типом данных zdb.fulltext. Ну что же, напишем свое. Так как этот тип данных для django ведет себя так же, как и нативный postgresql text, то при создании своего поля мы унаследуем наш класс от models.TextField. Вдобавок нужно сделать две важных вещи: выключить возможность использовать Btree-индекс на этом поле и ограничить backend для базы данных. В конечном результате это выглядит следующим образом:

class ZomboField(models.TextField):
    description = "Alias for Zombodb field"

    def __init__(self, *args, **kwargs):
        kwargs['db_index'] = False
        super().__init__(*args, **kwargs)

    def db_type(self, connection):
        databases = [
            'django.db.backends.postgresql_psycopg2',
            'django.db.backends.postgis'
        ]
        if connection.settings_dict['ENGINE'] in databases:
            return 'zdb.fulltext'
        else:
            raise TypeError('This database not support')

В-третьих, объясним ZomboDB где искать наш ElasticSearch. В самой базе с этой целью используется кастомный индекс от ZomboDB. Поэтому если адрес поменяется, то и индекс нужно изменить.
Django именует таблицы по шаблону app_model: в нашем случае приложение называется main, а модель — article. elasticsearch — это dns-имя, которое докер присваивает по названию контейнера.
В SQL это выглядит так:

CREATE INDEX idx_main_article
          ON main_article
       USING zombodb ((main_article.*))
        WITH (url='elasticsearch:9200/');

В Django нам тоже нужно создать кастомный индекс. Индексы там пока еще не очень гибкие: в частности, zombodb индекс указывает не на конкретную колонку, а на всю таблицу целиком. В Django же индекс требует обязательное указание на поле. Поэтому я подменил statement.parts['columns'] на ((main_article.*)), но методы construct и deconstruct по-прежнему требуют указывать атрибут fields при создании поля. Так же нам нужно передать дополнительный параметр в params. Для чего переопределим метод __init__, deconstruct и get_with_params.
В целом, конструкция получилась рабочая. Миграции применяются и отменяются без проблем.

class ZomboIndex(models.Index):
    def __init__(self, *, url=None, **kwargs):
        self.url = url
        super().__init__(**kwargs)

    def create_sql(self, model, schema_editor, using=''):
        statement = super().create_sql(model, schema_editor, using=' USING zombodb')
        statement.parts['columns'] = '(%s.*)' % model._meta.db_table
        with_params = self.get_with_params()
        if with_params:
            statement.parts['extra'] = " WITH (%s) %s" % (
                ', '.join(with_params),
                statement.parts['extra'],
            )
        print(statement)
        return statement

    def deconstruct(self):
        path, args, kwargs = super().deconstruct()
        if self.url is not None:
            kwargs['url'] = self.url
        return path, args, kwargs

    def get_with_params(self):
        with_params = []
        if self.url:
            with_params.append("url='%s'" % self.url)
        return with_params

Кому такой подход не по душе могут использовать миграции с RunSQL, напрямую добавив индекс. Только придется следить за названием таблицы и индекса самостоятельно.

migrations.RunSQL(
    sql = (
        "CREATE INDEX idx_main_article "
        "ON main_article "
        "USING zombodb ((main_article.*)) "
        "WITH (url='elasticsearch:9200/');"
    ),
    reverse_sql='DROP INDEX idx_main_article'
)

В итоге получилась вот такая модель. ZomboField принимает те же самые аргументы, что и TextField, с одним исключением — index_db ни на что не влияет, так же как и атрибут fields в ZomboIndex.

class Article(models.Model):
    text = ZomboField()

    class Meta:
        indexes = [
            ZomboIndex(url='elasticsearch:9200/', name='zombo_idx', fields=['text'])
        ]

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

from django.db import migrations, models
from django.contrib.postgres.operations import CreateExtension
import main.models

class Migration(migrations.Migration):

    initial = True

    dependencies = [
    ]

    operations = [
        CreateExtension('zombodb'),
        migrations.CreateModel(
            name='Article',
            fields=[
                ('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
                ('text', main.models.ZomboField()),
            ],
        ),
        migrations.AddIndex(
            model_name='article',
            index=main.models.ZomboIndex(fields=['text'], name='zombo_idx', url='elasticsearch:9200/'),
        )
    ]

Для интересующихся прилагаю SQL, который выдает Django ORM (можно посмотреть через sqlmigrate, ну, или с учетом докера: sudo docker-compose exec django python3 manage.py sqlmigrate main 0001)

BEGIN;
--
-- Creates extension zombodb
--
CREATE EXTENSION IF NOT EXISTS "zombodb";
--
-- Create model Article
--
CREATE TABLE "main_article" ("id" serial NOT NULL PRIMARY KEY, "text" zdb.fulltext NOT NULL);
--
-- Create index zombo_idx on field(s) text of model article
--
CREATE INDEX "zombo_idx" ON "main_article" USING zombodb ((main_article.*)) WITH (url='elasticsearch:9200/') ;
COMMIT;

Итак, модель у нас есть. Осталось теперь сделать поиск через filter. Для этого опишем свой lookup и зарегистрируем его.

@ZomboField.register_lookup
class ZomboSearch(models.Lookup):
    lookup_name = 'zombo_search'

    def as_sql(self, compiler, connection):
        lhs, lhs_params = self.process_lhs(compiler, connection)
        rhs, rhs_params = self.process_rhs(compiler, connection)
        params = lhs_params + rhs_params
        return "%s ==> %s" % (lhs.split('.')[0], rhs), params

Поиск в таком случае будет выглядеть следующим образом:

Article.objects.filter(text__zombo_search='(call OR box)')

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

from django.db.models import FloatField, Func

class ZomboScore(Func):
    lookup_name = 'score'
    function = 'zdb.score'
    template = "%(function)s(ctid)"
    arity = 0

    @property
    def output_field(self):
        return FloatField()

Теперь можно строить довольно сложные запросы без особых проблем.

scores = (Article.objects
          .filter(text__zombo_search='delete')
          .annotate(score=ZomboScore())
          .values_list(F('score'))
          .order_by('-score'))

Подсветка результата (highlight) оказалась несколько сложнее, красиво не получилось. Django psycopg2 backend в любых ситуациях преобразует имя_колонки в таблица.имя_колонки. Если было text, то будет "main_article"."text", чего ZomboDB категорически не приемлет. Указание колонки должно быть исключительно текстовым именем колонки. Но и здесь нам на помощь приходит RawSQL.

from django.db.models.expressions import RawSQL

highlighted = (Article.objects
               .filter(text__zombo_search='delete')
               .values(highlight_text=RawSQL("zdb.highlight(ctid, %s)", ('text',))))

Полную версию проекта с тестами можно посмотреть в репозитории. Все примеры из статьи оформлены там в виде тестов. Надеюсь для кого-нибудь эта статья окажется полезной и подтолкнет не писать велосипед на сигналах, с возможностью отстрелить себе всю консистентность, а использовать готовое решение не теряя все положительные стороны ORM. Дополнения и исправления также приветствуются.