Базы из Python: psycopg, параметры, батчи и транзакции

«Какие библиотеки используете для обработки данных?» звучит как вопрос про список, но интервьюер проверяет другое: понимаете ли вы, когда нужен драйвер базы, когда pandas, а когда Spark. Этот урок про первый слой, прямую работу с базой из Python, потому что именно там происходят самые дорогие ошибки: инъекция в запрос, загрузка по одной строке на миллион строк и половина батча в витрине после падения.

Какие библиотеки и для чего

Честный ответ на собеседовании называет две-три библиотеки с задачей для каждой, а не весь список. «psycopg для загрузок в Postgres, pandas для разведки и отчётов, PySpark когда данные на кластере» звучит убедительнее, чем перечисление.

Соединение, курсор, параметры

import psycopg
 
with psycopg.connect("postgresql://user:pass@host:5432/dwh") as conn:
    with conn.cursor() as cur:
        cur.execute(
            "SELECT order_id, amount FROM orders WHERE city = %s AND created_at >= %s",
            (city, since),
        )
        for order_id, amount in cur:
            ...

Соединение это ресурс: with закроет его при любом исходе, а в psycopg ещё и закоммитит транзакцию при успешном выходе и откатит при исключении. Курсор выполняет запросы и отдаёт результат итератором: for row in cur читает строки по мере надобности, cur.fetchall() читает всё в список. Для больших выборок второе это та же ловушка с памятью, что readlines().

Параметры передаются отдельно от текста запроса через %s и кортеж. Это единственный правильный способ подставить значение в SQL. Строка, собранная через f-строку, это инъекция:

cur.execute(f"SELECT ... WHERE city = '{city}'")      # нельзя

Если city пришёл из API или файла и содержит '; DROP TABLE orders; --, вы выполните чужой код. Даже без злого умысла значение с апострофом ("Марьина роща" без проблем, "O'Brien" сломает запрос) даст синтаксическую ошибку в три часа ночи. Параметризация экранирует всё за вас, и заодно база может переиспользовать план запроса.

Параметрами можно передавать только значения. Имена таблиц и колонок параметрами не передаются: WHERE %s = 1 не станет WHERE city = 1. Для динамических имён есть psycopg.sql.Identifier, который экранирует идентификатор, или белый список допустимых имён в коде.

Загрузка батчами

Загрузка по одной строке это сетевой запрос на строку. Миллион строк это миллион обменов с базой, часы работы. Три способа быстрее, в порядке возрастания скорости:

rows = [(1, "paid", 100), (2, "new", 50)]
 
cur.executemany(
    "INSERT INTO fact_orders (order_id, status, amount) VALUES (%s, %s, %s)",
    rows,
)

executemany отправляет строки пачкой, в psycopg 3 это действительно один обмен на батч. Для десятков тысяч строк достаточно.

Для миллионов строк в Postgres есть COPY, самый быстрый путь: база читает поток данных напрямую, без разбора отдельных команд:

with cur.copy("COPY fact_orders (order_id, status, amount) FROM STDIN") as copy:
    for row in rows:
        copy.write_row(row)

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

Размер батча это компромисс: слишком маленький делает много обменов, слишком большой держит память и длинную транзакцию. Обычно 1 000–10 000 строк, и в коде это параметр, а не константа внутри цикла.

Транзакция и повторный запуск

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

В psycopg транзакция открывается при первом запросе и завершается conn.commit() или conn.rollback(); with conn: делает это автоматически. Правильная структура шага загрузки:

with psycopg.connect(dsn) as conn:
    with conn.cursor() as cur:
        cur.execute("DELETE FROM fact_orders WHERE load_date = %s", (day,))
        cur.executemany(insert_sql, rows_for(day))
    # выход из with: COMMIT, при исключении ROLLBACK

Удаление дня и вставка в одной транзакции делают шаг идемпотентным: повторный запуск за тот же день даёт тот же результат. Это та же идея, что в уроке SQL про повторную загрузку, только теперь границы транзакции рисует Python.

Включённый autocommit=True превращает каждую команду в свою транзакцию. Он нужен для VACUUM и CREATE DATABASE, которые не работают внутри транзакции, и вреден для загрузок.

SQLAlchemy и ORM

SQLAlchemy даёт два слоя. Core: единый способ собирать запросы и подключаться к разным базам, нужен, когда проект ходит и в Postgres, и в MySQL, или когда pandas читает через read_sql. ORM: классы вместо таблиц, объекты вместо строк, нужен в приложениях с бизнес-логикой.

Для загрузок ORM обычно лишний: вы думаете таблицами и батчами, а не объектами, и накладные расходы на создание объекта на каждую строку заметны на миллионах. Инженеры данных используют SQLAlchemy чаще как универсальный коннектор (create_engine плюс сырой SQL), чем как ORM.

Spark и когда его брать

PySpark это другой режим работы: код описывает трансформацию, а выполняется она на кластере, партициями, не в памяти вашего процесса. Брать его стоит, когда данные не помещаются на одну машину или когда они уже лежат в озере в Parquet и нужны соединения между большими таблицами. Для файла на 5 ГБ и Postgres на одном сервере Spark это overhead на запуск и инфраструктуру без выигрыша. Отдельный модуль про Spark дальше в roadmap.

Как отвечать на собеседовании

Какие библиотеки используете для данных? Драйвер базы для загрузок (psycopg для Postgres, свой для ClickHouse), pandas для разведки и небольших трансформаций, pyarrow для Parquet, PySpark для данных на кластере, SQLAlchemy как универсальный коннектор.

Как защититься от SQL-инъекции? Параметризованные запросы: значения передаются отдельно от текста через %s и кортеж, драйвер их экранирует. Имена таблиц параметрами не передать, для них sql.Identifier или белый список.

Как быстро загрузить миллион строк? Не по одной: executemany батчами по несколько тысяч, для Postgres COPY из генератора. Всё в одной транзакции вместе с удалением старых данных за период, чтобы повторный запуск был идемпотентным.

Зачем транзакция в загрузке? Чтобы падение на середине не оставило половину данных. Удаление и вставка в одной транзакции либо проходят целиком, либо откатываются.

Практика

Сети и базы в песочнице нет, поэтому вместо настоящего соединения в задачах используется класс FakeCursor: он записывает, какие запросы и с какими параметрами вы выполнили. Интерфейс как у psycopg: execute(sql, params), executemany(sql, rows).

class FakeCursor:
    """Запоминает запросы вместо выполнения."""

    def __init__(self, rows=()):
        self.calls = []
        self.rows = list(rows)

    def execute(self, sql, params=None):
        self.calls.append((" ".join(sql.split()), tuple(params) if params is not None else None))

    def executemany(self, sql, rows):
        rows = list(rows)
        self.calls.append((" ".join(sql.split()), len(rows)))

    def fetchall(self):
        return list(self.rows)