Кастомизация 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. Дополнения и исправления также приветствуются.
Автор: menstenebris