Как повысить безопасность приложений с помощью фабрик строк в psycopg

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

Эти специализированные модели обычно определяются разработчиками приложений в виде классов данных Python, сопоставляемых с одной или более связей (обычно таблиц) в PostgreSQL. Именно здесь зачастую в игру вступает ORM (объектно-реляционное отображение), предоставляя абстрактный слой для написания запросов из классов моделей и отображения результатов в экземпляры моделей.

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

▍ Что такое фабрика строк?

Представим, что работаем с данными о погоде, собирая отчёты по датам для разных городов в таблицу weather:

CREATE TABLE weather (
    city            varchar(80),
    temp_lo         int,
    temp_hi         int,
    prcp            real,
    date            date
);

При запросе этой таблицы через psycopg мы получим строки в виде значений кортежа, где каждый столбец будет иметь подходящий тип Python, соответствующий типу PostgreSQL:

>>> with psycopg.connect() as conn:
...     print(conn.execute("SELECT * FROM weather").fetchone())
('San Francisco', 46, 50, 0.25, datetime.date(1994, 11, 27))

Самый простой способ получить что-то другое — это использовать фабрики строк, например, dict_row, которая создаёт словарь, где имена столбцов результатов сопоставляются со значениями:

>>> with psycopg.connect() as conn, conn.cursor(row_factory=psycopg.rows.dict_row) as cur:
...     print(cur.execute("SELECT * FROM weather").fetchone())
{'city': 'San Francisco', 'temp_lo': 46, 'temp_hi': 50, 'prcp': 0.25, 'date': datetime.date(1994, 11, 27)}

Фабрика строк передаётся при инициализации курсора, и строка, возвращаемая fetchone(), на деле оказывается словарём.

А теперь кое-что поинтереснее. Предположим, что у нас в приложении Python есть следующая специализированная модель:

from dataclasses import dataclass
from datetime import date

@dataclass
class Weather:
    city: str
    temperature: tuple[int, int]
    precipitation: float
    date: date

Целью будет получить экземпляры класса Weather при выполнении запросов вроде SELECT * FROM weather. Здесь нам помогут уже пользовательские фабрики строк.

Фабрика строк обычно представляет собой функцию, преобразующую сырые данные Python (которые psycopg адаптировала из SQL значений) в некий конечный тип — в данном случае это экземпляр класса Weather. Написать подобную функцию можно так:

def weather_from_row(city, temp_lo, temp_hi, prcp, date):
    return Weather(
        city=city, temperature=(temp_lo, temp_hi), precipitation=prcp, date=date
    )

Однако этого будет недостаточно, поскольку psycopg недостаёт информации о столбцах в наборе результатов SQL для использования с этой функцией. Именно поэтому фабрика строк дополнительно задействует курсор, используемый для текущего запроса. Соответственно, эту функцию необходимо обернуть в фабрику:

def weather_row_factory(cursor):
    # Извлекает имена столбцов результатов.
    columns = [column.name for column in cursor.description]

    def make_row(values):
        # Сопоставляет имена столбцов со значениями.
        row = dict(zip(columns, values))
        return weather_from_row(**row)

    return make_row

Кроме того, мы получим более удачную структуру, если эти функции будут являться методами класса Weather, приведённого выше:

@dataclass
class Weather:
    city: str
    ...

    @classmethod
    def from_row(cls, *, city, temp_lo, temp_hi, prcp, date):
        return cls(
            city=city, temperature=(temp_lo, temp_hi), precipitation=prcp, date=date
        )

    @classmethod
    def row_factory(cls, cursor):
        columns = [column.name for column in cursor.description]

        def make_row(values):
            row = dict(zip(columns, values))
            return cls.from_row(**row)

        return make_row

Теперь, когда всё на своих местах, посмотрим, как оно работает:

>>> with psycopg.connect() as conn:
...     with conn.cursor(row_factory=Weather.row_factory) as cur:
...         cur.execute("SELECT * FROM weather")
...         row = cur.fetchone()
>>> print(row)
Weather(city='San Francisco', temperature=(46, 50), precipitation=0.25, date=datetime.date(1994, 11, 27))

Как видите, переменная row — это экземпляр Weather, значит, можно соответственным образом ей оперировать, например, обращаться к её атрибутам:

>>> import statictics
>>> statictics.mean(row.temperature)
48

▍ Почему это безопаснее?

В контексте баз данных, или ввода-вывода в общем, задействованный на границе код — в месте, где происходит преобразование данных между разными системами типов — обычно является «небезопасным» и требует особого внимания, например, валидации при кодировании/декодировании.

Реализуя отправку запросов из psycopg посредством фабрик строк, как это показано выше, мы выполняем преобразование «сырых» типов Python в специализированные модели на ранней стадии. Таким образом, время, проведённое на этой границе ввода-вывода, сокращается. Как правило, необходимо с особой внимательностью создавать корректные SQL-запросы и соответствующие фабрики строк, а также старательно их тестировать. Тогда остальная часть кода приложения будет прекрасно работать, не озадачиваясь типом данных, поступающих из БД. В этом также зачастую помогает ORM.

Ещё одно существенное преимущество проявляется при совмещении фабрик строк со статической типизацией Python. Статическая типизация с модулями проверки типов вроде mypy обеспечивает дополнительные гарантии безопасности, не влияя на производительность в среде выполнения. Однако использовать этот подход на границе ввода-вывода будет сложнее, поскольку он обычно подразумевает «слабую» типизацию (то есть типы str или dict в противоположность, например, Weather, считающемуся «строгим» типом). В этом контексте значительно помогут фабрики строк, так как привнесут строгую типизацию на границу ввода-вывода. В результате можно будет спокойно использовать типы моделей со строгими возможностями валидации (например, библиотеки вроде pydantic).

Итак, вернёмся к предыдущему примеру и добавим в него аннотации типов:

from typing import Any, Sequence

@dataclass
class Weather:
    city: str
    ...

    @classmethod
    def from_row(cls, *, city: str, temp_lo: int, temp_hi: int, prcp: float, date: date) -> Weather:
        return cls(
            city=city, temperature=(temp_lo, temp_hi), precipitation=prcp, date=date
        )

    @classmethod
    def row_factory(cls, cursor: Cursor[Any]) -> Callable[[Sequence[Any]], Weather]:
        columns = [column.name for column in cursor.description]

        def make_row(values: Sequence[Any]) -> Weather:
            row = dict(zip(columns, values))
            return cls.from_row(**row)

        return make_row

Следующий пример спокойно проходит проверку типов:

def get_weather_reports(conn: Connection[Any]) -> list[Weather]:
    with conn.cursor(row_factory=Weather.row_factory) as cur:
        cur.execute("SELECT * FROM weather")
        return cur.fetchall()

При передаче типизированной фабрики строк в курсор, используемый для получения запросов, значение, возвращаемое вызовом fetchall(), будет выведено с соответствующим типом строки. Конкретно здесь курсор содержит фабрику, возвращающую строки с типом Weather, в результате чего fetchall() возвращает значение list[Weather], соответствующее тому, что объявлено в get_weather_reports() (и будет проверено модулем проверки типов вроде mypy). Такой подход безопаснее, поскольку данные поступают из БД в виде строгих типов (например, Weather), а не tuple (или dict) анонимных значений (int, date и т.д.).

Безопасность также увеличивает тот факт, что мы передаём эти строго типизированные значения, тем самым избегая ошибок в среде выполнения (например, IndexError для tuple или KeyError для dict), которые модуль проверки типов выявить неспособен. При этом типичную AttributeError, возникающую в случае строгих типов, обнаружить проблем не составит. К примеру, нижеприведённую IndexError невозможно перехватить до среды выполнения:

>>> with conn.cursor() as cur:
...     rows = cur.execute("SELECT * FROM weather").fetchall()  # тип list[tuple[Any, ...]]
>>> rows[0][5]
Traceback (most recent call last):
  ...
IndexError: tuple index out of range

А вот следующую AttributeError модуль проверки уже отловит, и в среде выполнения она не возникнет:

>>> with conn.cursor(row_factory=Weather.row_factory) as cur:
...     rows = cur.execute("SELECT * FROM weather").fetchall() # тип list[Weather]
>>> rows[0].prcp
Traceback (most recent call last):
  ...
AttributeError: 'Weather' object has no attribute 'prcp'

Наконец, с помощью фабрик строк мы повышаем не только безопасность, но и выразительность, так как при обращении к базе данных избегаем неуклюжих dict и прочего.

▍ Как это работает изнутри?

Теперь последняя часть предыдущего раздела, демонстрирующая, что типы, определённые в фабрике строк, распространяются через методы Cursor, может показаться несколько магической. Давайте тогда ещё раз возьмём последний пример и добавим несколько вызовов reveal_type(), отразив результаты в комментариях:

def get_weather_reports(conn: Connection[Any]) -> list[Weather]:
    with conn.cursor(row_factory=Weather.row_factory) as cur:
        reveal_type(cur)
        # note: раскрыт тип "psycopg.cursor.Cursor[Weather]"
        cur.execute("SELECT * FROM weather")
        rset = cur.fetchall()
        reveal_type(rset)
        # note: раскрыт тип "builtins.list[Weather]"
        return rset

Видно, что значение cur: Cursor[Weather] параметризуется в «строковом» типе Weather, как объявлено в Weather.row_factory(cursor: Cursor[Any]) -> Callable[[Sequence[Any]], Weather] (важная часть — это тип Weather, объявленный как возвращаемое значение объекта вызова, возвращаемого фабрикой строк). Аналогичным образом rset: list[Weather] на деле выводится как список объектов Weather.

Чтобы понять, как это работает, мы сначала разберём определение Connection в psycopg (в упрощённом виде):

Row = TypeVar("Row", covariant=True)
RowFactory = Callable[[Sequence[Any]], Row]  # упрощённая форма

class Connection:

    @overload
    def cursor(self) -> Cursor[Tuple[Any, ...]]:
        ...

    @overload
    def cursor(self, *, row_factory: RowFactory[Row]) -> Cursor[Row]:
        ...

    def cursor(self, *, row_factory: Optional[RowFactory[Any]] = None) -> Cursor[Any]:
        # здесь идёт реализация

Метод Connection.cursor() определяется как перегрузка, завися от значения параметра row_factory и возвращая объект Cursor с переменной особого типа Row. Иными словами, эта переменная типа Row привязывается от параметра фабрики строк к возвращаемому значению Cursor.

Затем та же переменная с типом Row используется для определения обобщённого класса Cursor, который, в свою очередь, позволяет методам fetch*() возвращать значения Row (также в упрощённой форме):

class Cursor(Generic[Row]):

    def fetchone(self) -> Optional[Row]:
        ...

    def fetchall(self) -> List[Row]:
        ...

Получается, если вернуться к последнему примеру:

>>> cur = conn.cursor(row_factory=Weather.row_factory)
>>> cur.execute("SELECT * FROM weather")
>>> rows = cur.fetchall()

• cur — это Cursor[Weather], поскольку Weather.row_factory имеет тип RowFactory[Weather], а
• rows — это list объектов Weather, поскольку cur является Cursor[Weather].

Вот и весь обзор фабрик строк в psycopg. Далее рекомендую обратиться к документации этого инструмента, в частности, к странице о статической типизации (где объясняется, каким образом Cursor является обобщением для Row), а также к основной странице фабрик строк и странице, посвящённой модулю rows, где приводятся полезные вспомогательные генераторы фабрик строк.

Telegram-канал и уютный чат для клиентов