Учимся писать расширения на языке C для PostgreSQL | Записки программиста

Александр Алексеев
для блога «Записки программиста»
3 октября 2016

          Недавно я выложил на GitHub ZSON, расширение к PostgreSQL для сжатия JSONB. Сжатие происходит путем анализа существующих в базе документов и создания словаря с наиболее часто встречающимися в документах строками. Притом строки могут быть не только именами ключей, но и значениями в массивах, и так далее. В этой статье на примере ZSON мы разберемся, как вообще пишутся расширения к PostgreSQL, как они покрываются тестами, как происходит их установка и удаление, и так далее.

          Взглянем на код.

          Файл zson.control содержит основные сведения о расширении — главным образом его название и версию:


comment = 'ZSON'
default_version = '1.0'
module_pathname = '$libdir/zson'
relocatable = true

          Файл zson--1.0.sql содержит запросы, которые будут выполнены, когда пользователь скажет CREATE EXTENSION:


-- complain if script is sourced in psql
-- rather than via CREATE EXTENSION
\echo USE "CREATE EXTENSION zson" TO LOAD this file. \quit

CREATE TYPE zson;

CREATE TABLE zson_dict (
    dict_id SERIAL NOT NULL,
    word_id INTEGER NOT NULL,
    word text NOT NULL,
    PRIMARY KEY(dict_id, word_id)
);

-- ... и еще много-много кода...

          PostgreSQL — очень гибкая СУБД. Она позволяет объявлять собственные процедуры и типы, используя только SQL и PL/pgSQL. Многие расширения могут содержать весь код в одном только .sql файле, без единой сточки на языке C.

          Makefile содержит информацию о том, как собрать расширение, с какими флагами его компилировать, какие библиотеки следует прилинковать, и прочее:


EXTENSION = zson
MODULES = zson
DATA = zson--1.0.sql
OBJS = zson.o
REGRESS = zson

MODULE_big = zson
PG_CPPFLAGS = -g -O2
SHLIB_LINK = # -lz -llz4

PGXS := $(shell pg_config --pgxs)
include $(PGXS)

          В файле sql/zson.sql содержится код тестов:


CREATE EXTENSION zson;
SELECT zson_extract_strings('true');
SELECT zson_extract_strings('"aaa"');

-- ... и так далее ...

… а в expected/zson.out — вывод, который ожидается в результате выполнения этих тестов:


CREATE EXTENSION zson;
SELECT zson_extract_strings('true');
 zson_extract_strings 
----------------------
 {}
(1 row)

SELECT zson_extract_strings('"aaa"');
 zson_extract_strings 
----------------------
 {aaa}
(1 row)

... и так далее ...

          Получить такой .out файл проще всего командой:


psql -eq test_database < t.sql > t.out

          Пока не слишком сложно, не так ли? Теперь перейдем, пожалуй, к самому интересному — коду на C из файла zson.c.

          Начинается он с довольно типичных для расширений к PostgreSQL заголовочных файлов:


#include <postgres.h>
#include <port.h>
#include <catalog/pg_type.h>
#include <executor/spi.h>
#include <utils/builtins.h>
#include <utils/jsonb.h>
#include <limits.h>
#include <string.h>

          Дальше идет «волшебный» макрос, который не нужно понимать, нужно просто не забыть написать :)


PG_MODULE_MAGIC;

          Объявление процедур на C, экспортируемых расширением:


PG_FUNCTION_INFO_V1(zson_in);
PG_FUNCTION_INFO_V1(zson_out);
// ... и так далее ...

          С передачей аргументов процедурам на C все несколько хитро. Так, к примеру, выглядит тело jsonb_to_zson:


Datum
jsonb_to_zson(PG_FUNCTION_ARGS)
{
  Jsonb *jsonb = PG_GETARG_JSONB(0);
  uint8* jsonb_data = (uint8*)VARDATA(jsonb);
  Size jsonb_data_size = VARSIZE(jsonb) - VARHDRSZ;

  // ... пропущено ...

  SET_VARSIZE(encoded_buff, encoded_size);
  PG_RETURN_BYTEA_P(encoded_buff);
}

          Для приема и возвращения различных типов предусмотрены макросы PG_GETARG_* и PG_RETURN_* соответственно. Datum — это местный аналог uintptr_t, типа, который может содержать любое число, и при этом достаточно велик для хранения указателей на данной платформе. Передача аргументов и возвращение значений на границах работы с расширением осуществляется с его помощью.

          Также мы видим здесь пример работы с типами переменной длины, к которым относится JSONB, TEXT и другие. Никакой магии — просто есть заголовок, хранящий размер данных, за которыми следуют сами данные. Также предусмотрены макросы для получения указателя на сами данные, их размера и указания нового размера.

          Иногда такие «граничные» процедуры требуется вызывать прямо из кода на C. Для этого предусмотрены макросы DirectFunctionCall*:


Datum
zson_in(PG_FUNCTION_ARGS)
{
  Datum string_datum = PG_GETARG_DATUM(0);
  Datum jsonb_datum = DirectFunctionCall1(jsonb_in, string_datum);
  Datum zson_datum = DirectFunctionCall1(jsonb_to_zson, jsonb_datum);
  bytea* zson_bytea = DatumGetByteaP(zson_datum);
  PG_RETURN_BYTEA_P(zson_bytea);
}

          Для доступа к базе данных из расширения предусмотрен так называемый Server Programming Interface или SPI. Пользоваться им не сложно. Так, например, выглядит часть кода процедуры get_current_dict_id, возвращающей максимальный id словаря из базы:


SPI_connect();

if(savedPlanGetDictId == NULL)
{
  savedPlanGetDictId = SPI_prepare(
    "select max(dict_id) from zson_dict;", 0, NULL);
  if (savedPlanGetDictId == NULL)
    elog(ERROR, "Error preparing query");
  if (SPI_keepplan(savedPlanGetDictId))
    elog(ERROR, "Error keeping plan");
}

if (SPI_execute_plan(savedPlanGetDictId, NULL, NULL, true, 1) < 0 ||
    SPI_processed != 1)
  elog(ERROR, "Failed to get current dict_id");

id = DatumGetInt32(
    SPI_getbinval(SPI_tuptable->vals[0], SPI_tuptable->tupdesc,
                  1, &isnull)
  );

SPI_finish();

          Все подробности есть в официальной документации PostgreSQL.

          Макрос elog нам ранее не встречался. Вызванный с уровнем INFO, NOTICE или WARNING он просто выводит предупреждение в сессию пользователя. С уровнем ERROR он аварийно завершает выполнение кода и откатывает транзакцию. Чтобы это работало в C, PostgreSQL имеет свою небольшую реализацию исключений на setjmp и longjmp. Макрос может быть вызван с форматной строкой и переменным числом аргументов, в точности, как printf.

          Кстати, elog очень удобно использовать для отладки расширения. Конечно, отладчики GDB и LLDB тоже никто не отменял.

          Установка расширения производится так:


cd /path/to/source/code
make
sudo make install

          Прогон тестов:


make installcheck

          Включение и выключение расширения для базы данных:


CREATE EXTENSION zson;
DROP EXTENSION zson;

          Само собой разумеется, при удалении расширения все созданные им типы, таблицы и так далее автоматически вычищаются.

          Наконец, полное удаление расширения производится командой:


sudo make uninstall

          Как видите, в первом приближении все довольно просто. К сожалению, в рамках одной статьи невозможно осветить абсолютно все аспекты написания расширений к PostgreSQL. Все сильно зависит от того, что именно вы хотите сделать. Например, начиная с PostgreSQL 9.6 стало возможным создавать в расширении собственные типы индексов. Но это совершенно отдельный механизм, заслуживающий целой статьи, а то и нескольких. Если вас интересует эта тема, посмотрите, как работает расширение contrib/bloom, добавляющее в PostgreSQL индекс, работающий по принципу фильтра Блума. Если же вы, например, хотите завести новый тип данных, то это совершенно другой вопрос.

          Из дополнительных источников информации можно рекомендовать:

  • Помимо прочего, книга PostgreSQL Server Programming рассматривает и написание расширений к PostgreSQL;
  • В исходниках PostgreSQL есть каталог contrib, содержащий множество примеров совершенно настоящих, используемых в боевых системах, расширений;
  • У компании adjust есть замечательная серия статей из пяти частей, посвященная написанию расширений к PostgreSQL: часть один, два, три, четыре, пять;

          У меня все. Как всегда, буду несказанно рад любым вашим вопросам и дополнениям!