Как наполнить базу сгенерированными джейсонами

Предположим, у вас Postgres с миллионами JSONb-документов. Вы хотите проверить нагрузку на стейджинге, но данные с прода брать нельзя – их нужно сгенерить. Тут начинаются проблемы.

Размножить один и тот же джейсон будет неправильным, потому что все они дадут одинаковое значение индекса. В итоге индекс будет “перекошен”: миллион записей попадут в один блок, а остальные сто тысяч будут равномерны. В реальности так не бывает.

Второй вопрос – как генерить. Можно взять язык, на котором вы пишете, и наколбасить случайные словари в CSV. Потом сжать в gzip, перетащить на сервер с базой, распаковать и вставить через COPY IN. Но придется писать код и воевать с передачей файлов по SCP/SSH.

Я склоняюсь к вот такому костылику.

Берем с прода любой JSON, записываем в локальный файл с красивым форматированием. Например:

{
    "meta": {
        "eyes": "brown"
    },
    "attrs": {
        "email": "petr@test.com",
        "name": "Petr Ivanov"
    },
    "roles": [
        "user",
        "admin"
    ]
}

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

cat sample.json | sed \
     -e "s/\"/'/g" -e 's/: /, /g' \
     -e 's/{/jsonb_build_object(/g' \
     -e 's/}/\)/g' \
     -e 's/\[/jsonb_build_array(/g' \
     -e 's/\]/)/g'

Получается SQL-код, который строит тот же самый JSON:

jsonb_build_object(
    'meta', jsonb_build_object(
        'eyes', 'brown'
    ),
    'attrs', jsonb_build_object(
        'email', 'petr@test.com',
        'name', 'Petr Ivanov'
    ),
    'roles', jsonb_build_array(
        'user',
        'admin'
    )
)

Наберите в консоли SELECT и вставьте эту колбасу. База выплюнет в точности тот JSON, который был в файле.

Если форматировать JSON лень, добавьте в пайп утилиту jq — по умолчанию она просто форматирует документ.

Это был только один джейсон. Теперь размножим его с помощью select ... from generate_series...

select jsonb_build_object(
    'meta', jsonb_build_object(
        'eyes', 'brown'
    ),
    'attrs', jsonb_build_object(
        'email', 'petr@test.com',
        'name', 'Petr Ivanov'
    ),
    'roles', jsonb_build_array(
        'user',
        'admin'
    )
) as document
from
    generate_series(1, 5) as x;

Выборка вернет столько документов, сколько чисел в диапазоне generate_series. Но это один и тот же документ, что не подходит. Пройдитесь по значимым полям документа и замените статичные строки на форматирование, например так (переменная x ссылается на текущее значение generate_series):

select jsonb_build_object(
    'meta', jsonb_build_object(
        'eyes', format('color-%s', x)
    ),
    'attrs', jsonb_build_object(
        'email', format('user-%s@test.com', x),
        'name', format('Test Name %s', x)
    ),
    'roles', jsonb_build_array(
        'user',
        'admin'
    )
) as document
from
    generate_series(1, 5) as x;

Новая выборка станет такой:

                                                     document
-------------------------------------------------------------------------------------------------------------------------
 {"meta": {"eyes": "color-1"}, "attrs": {"name": "Test Name 1", "email": "user-1@test.com"}, "roles": ["user", "admin"]}
 {"meta": {"eyes": "color-2"}, "attrs": {"name": "Test Name 2", "email": "user-2@test.com"}, "roles": ["user", "admin"]}
 {"meta": {"eyes": "color-3"}, "attrs": {"name": "Test Name 3", "email": "user-3@test.com"}, "roles": ["user", "admin"]}
 {"meta": {"eyes": "color-4"}, "attrs": {"name": "Test Name 4", "email": "user-4@test.com"}, "roles": ["user", "admin"]}
 {"meta": {"eyes": "color-5"}, "attrs": {"name": "Test Name 5", "email": "user-5@test.com"}, "roles": ["user", "admin"]}
(5 rows)

Вставим выборку в таблицу документов:

create table documents(id uuid, document jsonb)

Для этого допишем в запрос шапку INSERT:

insert into documents (id, document)
select
    gen_random_uuid() as id,
    jsonb_build_object(
        '__generated__', true,
        'meta', jsonb_build_object(
            'eyes', format('color-%s', x)
        ),
        'attrs', jsonb_build_object(
            'email', format('user-%s@test.com', x),
            'name', format('Test Name %s', x)
        ),
        'roles', jsonb_build_array(
            'user',
            'admin'
        )
    ) as document
from
    generate_series(1, 500) as x
returning
    id;

Обратите внимание на поле __generated__: я добавил его, чтобы отличить обычный документ от сгенерированного.

Запустите запрос, и в таблице окажется миллион случайных документов. Вставка может занять до пары минут, потому что JSONb — дорогое удовольствие.

Если что-то пошло не так, удалите сгенерированные документы запросом:

delete from documents
where document -> '__generated__' = 'true'::jsonb;

Удобство в том, что не нужны питоны-джавы, все делается силами SQL. Не придется генерировать CSV и перетаскивать на сервер. Просто выполнили запрос — и пошли дальше.

Обязательно сохраните скрипт в недрах проекта. Он понадобится если не завтра, то через месяц, а если не вам, то коллеге. И вы такой раз — как Джонни Старк на картинке.