Расширяем PostgreSQL с помощью Rust

Свободная система управления базами данных PostgreSQL не только предоставляет высокопроизводительный движок для выполнения запросов, но и может быть расширена с помощью расширений, которые могут добавлять новые типы данных (например, для ГИС-расширений или астрономических координат), дополнительные типы индекса и возможности поиска (например, полнотекстовый поиск), сбор статистики, поддержку новых языков для встроенных функций и многое другое. Большой список существующих расширений может быть найден по этой ссылке. В этой статье мы рассмотрим один из возможных вариантов по созданию собственного расширения для PostgreSQL с использованием библиотеки pgx.

Для разработки расширения будет необходимо установить актуальную версию Rust не менее 1.6.4 и cargo для управления пакетами. Также должен быть установлен rustfmt, clang и libclang-dev, а также libreadline-dev (на Debian/Ubuntu):

curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
source "$HOME/.cargo/env"
sudo apt install libclang libclang-dev libreadline-dev

Установка библиотеки выполняется через cargo:

cargo install --locked cargo-pgx

Далее становится возможным настройка окружение и создание каркаса нового расширения через подкоманды cargo pgx:

• cargo pgx init - подготовка среды выполнения (загружает также версии PostgreSQL с 11 до 15-й);

• cargo pgx new <name> - создает новый проект с заготовкой плагина;

• cargo pgx run - запускает экземпляр PostgreSQL, компилирует плагин и устанавливает его в PostgreSQL и предоставляет доступ к консоли для выполнения команд (можно дополнительно указать версию pg11, pg12, pg13, pg14 или pg15);

• cargo pgx connect - подключается к активному экземпляру PostgreSQL через psql;

• cargo pgx install - установка расширения в активный экземпляр PostgreSQL (из pg_config);

• cargo pgx schema - создание схемы данных для расширения;

• cargo pgx start - запуск управляемого сервера PostgreSQL;

• cargo pgx stop - остановка управляемого сервера PostgreSQL;

• cargo pgx status - получить статус управляемого сервера;

• cargo pgx package - упаковка расширения для дальнейшей установки в любой поддерживаемый сервер PostgreSQL (через CREATE EXTENSION);

• cargo pgx test - запуск автоматических тестов для расширения.

Создаваемое по умолчанию расширение регистрирует функцию hello_<название_проекта>. Проверим корректность работы расширения, для этого сначала подключим расширение к PostgreSQL, а затем вызовем функцию. Для определенности будем считать, что расширение называется testext (было создано через cargo pgx new testext)

CREATE EXTENSION testext;
SELECT hello_testext();

Результатом выполнения будет отображена строка Hello, testext! Для сборки расширения как самостоятельного пакета добавим путь к pg_config:

export PATH=~/.pgx/13.10/pgx-install/bin:$PATH
cargo pgx package

В результате сборки проекта мы получим файл с библиотекой (.so для linux) в подкаталоге target.

Посмотрим на организацию проекта и предоставляемые возможности. Функции регистрируется в файле lib.rs. Здесь внешняя функция регистрируется как pg_extern.

use pgx::prelude::*;

pgx::pg_module_magic!();

#[pg_extern]
fn hello_testext() -> &'static str {
    "Hello, testext"
}

Функции могут принимать аргументы, например:

#[pg_extern]
fn my_to_lowercase(input: &'static str) -> String {
    input.to_lowercase()
}

При необходимости могут быть созданы новые типы данных. Для перечислений можно создать тип от PostgresEnum, для структур данных нужно обеспечить сериализацию/десериализацию через serde.

#[derive(PostgresEnum, Serialize)]
pub enum GenderValue {
    Male,
    Female,
}

use serde::{Serialize, Deserialize};

#[derive(Serialize, Deserialize, PostgresType)]
struct MyType {
    values: Vec<String>,
    thing: Option<Box<MyType>>;
}

Для выполнения запросов к базе данных можно использовать возможности SPI (подключается из pgx::spi), например:

• Spi::run("запрос") выполняет запрос без результата;

• Spi::get_one("запрос") возвращает Result<Option<i64>, pgx::spi::Error>

• Spi::get_one_with_args("запрос", vec!([аргументы]) используется для подстановки значений вместо $1, $2 и т.д. в запросе;

• Spi::connect() возвращает клиента, через который могут выполняться запросы (select, insert и др.).

Также можно использовать [pg_trigger] для создания функции для использования в триггере. (сама регистрация триггера выполняется через extension_sql!, например так:

extension_sql!("CREATE TRIGGER test_trigger BEFORE INSERT ON test FOR EACH ROW EXECUTE PROCEDURE trigger_example();)

Также может быть создан фоновый процесс, в этом случае он должен быть добавлен как предзагруженная библиотека в postgresql.conf (shared_preload_library='bgw.so'), здесь в _PG_init() регистрируется и запускается фоновый процесс (также здесь может быть установлено разрешение на доступ к серверному API PostgreSQL через вызов enable_spi_access).

use pgx::bgworkers::*;
use pgx::datum::{FromDatum, IntoDatum};
use pgx::log;
use pgx::prelude::*;
use std::time::Duration;

pgx::pg_module_magic!();

#[allow(non_snake_case)]
#[pg_guard]
pub extern "C" fn _PG_init() {
    BackgroundWorkerBuilder::new("bgw")
        .set_function("background_worker_main")
        .set_library("bgw")
        .set_argument(42i32.into_datum())
        .enable_spi_access()
        .load();
}

#[pg_guard]
#[no_mangle]
pub extern "C" fn background_worker_main(arg: pg_sys::Datum) {
    let arg = unsafe { i32::from_datum(arg, false) };

    BackgroundWorker::attach_signal_handlers(SignalWakeFlags::SIGHUP | SignalWakeFlags::SIGTERM);
    BackgroundWorker::connect_worker_to_spi(Some("postgres"), None);
    log!(
        "Background Worker '{}' is starting.  Argument={}",
        BackgroundWorker::get_name(),
        arg.unwrap()
    );
    while BackgroundWorker::wait_latch(Some(Duration::from_secs(10))) {
        if BackgroundWorker::sighup_received() {
            // здесь можно перезагрузить конфигурацию 
        }

        BackgroundWorker::transaction(|| {
            Spi::execute(|client| {
                let tuple_table = client.select(
                    "SELECT 'Hi', id, ''||a FROM (SELECT id, 42 from generate_series(1,10) id) a ",
                    None,
                    None,
                );
                tuple_table.for_each(|tuple| {
                    let a = tuple.by_ordinal(1).unwrap().value::<String>().unwrap();
                    let b = tuple.by_ordinal(2).unwrap().value::<i32>().unwrap();
                    let c = tuple.by_ordinal(3).unwrap().value::<String>().unwrap();
                    log!("from bgworker: ({}, {}, {})", a, b, c);
                });
            });
        });
    }
}

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

Для более подробного погружения в возможности библиотеки pgx можно использовать официальную документацию. Также хорошим местом для старта могут статьи примеры использования pgx в этом репозитории.

---

В заключение статьи рекомендую всем желающим посетить открытый урок «Работа с геоданными в Postgres», который состоится сегодня вечером в OTUS. На нем участники рассмотрят основные операции с геоданными. Записаться можно по ссылке.