Усиление PostgreSQL с помощью PL/Python

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

Сегодня мы прокачаем 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. Запись