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

Часто при работе с 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:10WORKDIR /home/RUN apt-get -y update && apt-get -y install curlADD 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.debRUN rm zombodb_jessie_pg10-10-1.0.3_amd64.debRUN apt-get -y clean

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

FROM python:stretchWORKDIR /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, modelsfrom django.contrib.postgres.operations import CreateExtensionclass 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, modelsfrom django.contrib.postgres.operations import CreateExtensionimport main.modelsclass 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_lookupclass 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, Funcclass 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 RawSQLhighlighted = (Article.objects .filter(text__zombo_search='delete') .values(highlight_text=RawSQL("zdb.highlight(ctid, %s)", ('text',))))

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

UPD: Появилась библиотека django-zombodb