Усиление PostgreSQL с помощью PL/Python09.10.2024 21:46

1369942a388e823c142ebdfe11d393d4.png

Привет, Хабр!

Сегодня мы прокачаем PostgreSQL, добавив в него Python. А именно — PL/Python. Это расширение позволяет писать функции на Python прямо внутри базы данных. Лично для меня это как объединение двух лучших миров: любимого PostgreSQL и могучего Python.

PL/Python — это про то, когда стандартного SQL мало. Когда надо сделать что-то действительно интересное: сложные расчеты, массивы данных, или интеграция с аналитикой прямо в базе. А самое крутое — можно тянуть любые Python-библиотеки.

Для начала нужно просто установить расширение в PostgreSQL: 

CREATE EXTENSION plpython3u;

Основы создания функций на PL/Python

PL/Python — это процедурный язык внутри PostgreSQL, который позволяет писать функции на Python. По сути, это как обычная SQL-функция, только вместо SQL-запросов ты пишешь Python-код.

Вот базовая конструкция: 

CREATE FUNCTION function_name(argument_list)
RETURNS return_type
AS $$
    # Здесь идет Python-код
$$ LANGUAGE plpythonu;

Всё, что между $ и $, — это чистый Python. PostgreSQL автоматически передаёт параметры как глобальные переменные в функцию, что позволяет работать с ними, как с обычными переменными Python.

Создадим функцию, которая возвращает большее из двух чисел: 

CREATE FUNCTION pymax(a integer, b integer)
RETURNS integer
AS $$
    if a > b:
        return a
    return b
$$ LANGUAGE plpythonu;

Эта функция принимает два целых числа и возвращает большее из них. Как видишь, очень похоже на обычный Python-код, только он работает внутри PostgreSQL.

Если не вернуть значение, PostgreSQL вернёт NULL.

Теперь копнем глубже. В PL/Python аргументы передаются как глобальные переменные, что открывает массу возможностей. Но есть нюанс: если попытаться переназначить аргумент внутри функции, можно получить ошибку.

CREATE FUNCTION pystrip(x text)
RETURNS text
AS $$
    x = x.strip()  # ошибка
    return x
$$ LANGUAGE plpythonu;

Ошибка возникает, потому что при присвоении переменной нового значения она становится локальной. В этом случае Python считает, что переменная ещё не инициализирована, и выдаёт ошибку.

Нужно явно указать, что используем глобальные переменные: 

CREATE FUNCTION pystrip(x text)
RETURNS text
AS $$
    global x
    x = x.strip()  # теперь всё работает
    return x
$$ LANGUAGE plpythonu;

Теперь всё работает правильно, и Python больше не ругается.

Это мелочь, но такая ошибка может испортить весь день.

PL/Python также есть возможность взаимодействовать с таблицами напрямую через Python. Например, можно написать функцию, которая достает данные из таблицы, обрабатывает их и возвращает результат.

Создадим функцию, которая извлекает email пользователя по его ID: 

CREATE FUNCTION get_user_email(user_id integer)
RETURNS text
AS $$
    query = plpy.execute(f"SELECT email FROM users WHERE id = {user_id}")
    if query:
        return query[0]['email']
    return None
$$ LANGUAGE plpythonu;

Здесь используем встроенную функцию plpy.execute() для выполнения SQL-запроса прямо из Python. Она возвращает результат в виде списка словарей, что удобно для дальнейшей обработки.

Интеграция с внешними библиотеками

PL/Python поддерживает все библиотеки Python, установленные на сервере, где крутится PostgreSQL.

Пример с Pandas. Допустим, есть данные о продажах, и нужно их быстро проанализировать: сумму, среднее значение и медиану по месяцам.

CREATE FUNCTION analyze_sales()
RETURNS table(month text, total_sales numeric, average_sales numeric, median_sales numeric)
AS $$
    import pandas as pd

    result = plpy.execute("SELECT month, sales FROM sales_data")
    df = pd.DataFrame(result)

    df_summary = df.groupby('month')['sales'].agg(['sum', 'mean', 'median']).reset_index()

    return df_summary.to_dict(orient='records')
$$ LANGUAGE plpythonu;

Здесь мы:

  1. Получаем данные из таблицы sales_data с помощью plpy.execute().

  2. Преобразуем результат в DataFrame через pandas для анализа.

  3. Группируем данные по месяцам и считаем сумму, среднее и медиану.

  4. Возвращаем результат обратно в базу данных.

Переходим к массивам данных. Если нужно работать с большими массивами или матрицами, numpy приходит на помощь.

Допустим, нужно вычислить среднее значение и стандартное отклонение по массиву данных: 

CREATE FUNCTION calculate_statistics(arr double precision[])
RETURNS table(mean double precision, stddev double precision)
AS $$
    import numpy as np

    np_arr = np.array(arr)

    mean = np.mean(np_arr)
    stddev = np.std(np_arr)

    return [{'mean': mean, 'stddev': stddev}]
$$ LANGUAGE plpythonu;

Пенедаем массив данных в PostgreSQL, преобразуем его в numpy-массив, а затем выполняем нужные расчёты.

Исключения и обработка ошибок

Обрабатывать исключения в PL/Python так же легко, как в обычном Python. Вот пример функции, которая делит два числа, обрабатывая возможные ошибки: 

CREATE FUNCTION safe_divide(a float, b float)
RETURNS float
AS $$
    try:
        return a / b
    except ZeroDivisionError:
        plpy.error("Деление на ноль невозможно!")
    except Exception as e:
        plpy.error(f"Произошла ошибка: {e}")
$$ LANGUAGE plpythonu;

Здесь обрабатываем ZeroDivisionError и любые другие исключения, отправляя ошибки в лог PostgreSQL через plpy.error().

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

Триггеры и транзакции

Триггеры в PL/Python работают так же, как и в обычном SQL. Они срабатывают при вставке, обновлении или удалении данных, и могут выполнять определённые действия.

Создадим триггерную функцию, которая проверяет количество товаров в заказе перед вставкой: 

CREATE FUNCTION validate_order_quantity()
RETURNS trigger
AS $$
    if NEW.quantity <= 0:
        raise plpy.Error('Количество товаров должно быть больше нуля!')
    return NEW
$$ LANGUAGE plpythonu;

Этот триггер проверяет каждую вставку в таблицу заказов и гарантирует, что количество товаров больше нуля.

Теперь создадим сам триггер: 

CREATE TRIGGER check_quantity
BEFORE INSERT OR UPDATE ON orders
FOR EACH ROW
EXEUTE FUNCTION validate_order_quantity();

Триггер срабатывает до вставки или обновления и вызывает нашу функцию validate_order_quantity().

Управление транзакциями

PL/Python можно управлять транзакциями вручную. Пример, когда мы явно управляем транзакцией: 

CREATE FUNCTION transaction_test()
RETURNS void
AS $$
    try:
        plpy.execute("BEGIN;")
        plpy.execute("INSERT INTO test_table VALUES (1);")
        plpy.execute("INSERT INTO test_table VALUES (2);")
        plpy.execute("COMMIT;")
    except:
        plpy.execute("ROLLBACK;")
        raise
$$ LANGUAGE plpythonu;

Здесь начинаем транзакцию с BEGIN, выполняем несколько операций и завершаем её командой COMMIT. Если что-то пойдёт не так, откатываем транзакцию с помощью ROLLBACK.

В завершение напомню про ближайшие открытые уроки, которые пройдут в рамках курса «PostgreSQL Advanced»:

  • 10 октября: Правила работы с аналитическими запросами. Оптимизация в OLTP и OLAP. Запись

  • 23 октября: PostgreSQL и DevOps — управляем базой данных через CI/CD и Kubernetes. Запись

© Habrahabr.ru