Подключение к реляционным базам данных в скриптах на Python

Расширенные примеры для разных сценариев

В этом блоке собраны примеры, которые выходят за рамки простого выполнения SELECT и INSERT. Они показывают возможности модуля sqlite3, SQLAlchemy и драйвера psycopg2.

Словарь строк через row_factory

Объект sqlite3.Row позволяет получать значения по имени колонки. Это удобно при большом количестве полей и уменьшает вероятность ошибки с индексами.

import sqlite3

conn = sqlite3.connect(':memory:')
conn.row_factory = sqlite3.Row

cursor = conn.cursor()
cursor.execute('CREATE TABLE employees (id INTEGER PRIMARY KEY, name TEXT, salary INTEGER)')
cursor.executemany(
    'INSERT INTO employees (name, salary) VALUES (?, ?)',
    [('Анна', 60000), ('Борис', 85000), ('Виктор', 95000)]
)
conn.commit()

cursor.execute('SELECT name, salary FROM employees ORDER BY salary DESC')
rows = cursor.fetchmany(2)
for row in rows:
    print(row['name'], row['salary'])
conn.close()

Виктор 95000
Борис 85000

Метод fetchmany возвращает ограниченное число строк, что помогает при постраничном выводе.

Пользовательская функция SQLite

SQLite позволяет создавать функции на Python и использовать их прямо в SQL. Следующий пример считает цену со скидкой.

import sqlite3

def discount(price, percent):
    return round(price * (100 - percent) / 100, 2)

conn = sqlite3.connect(':memory:')
conn.create_function('discount', 2, discount)
cursor = conn.cursor()
cursor.execute('CREATE TABLE products (name TEXT, price REAL)')
cursor.executemany(
    'INSERT INTO products VALUES (?, ?)',
    [('Монитор', 25000.00), ('Клавиатура', 5000.00)]
)
conn.commit()

cursor.execute('SELECT name, discount(price, 10) FROM products')
for name, new_price in cursor.fetchall():
    print(name, new_price)
conn.close()

Монитор 22500.0
Клавиатура 4500.0

Функция create_function принимает имя функции, число аргументов и ссылку на функцию Python.

Явное управление транзакцией в SQLite

При isolation_level=None SQLite не открывает транзакцию автоматически. Это позволяет явно управлять началом и завершением блока обновлений.

import sqlite3

conn = sqlite3.connect(':memory:', isolation_level=None)
cursor = conn.cursor()
cursor.execute('CREATE TABLE accounts (id INTEGER PRIMARY KEY, balance INTEGER)')
cursor.execute('INSERT INTO accounts VALUES (1, 1000)')
cursor.execute('INSERT INTO accounts VALUES (2, 500)')

try:
    cursor.execute('BEGIN')
    cursor.execute('UPDATE accounts SET balance = balance - 200 WHERE id = 1')
    cursor.execute('UPDATE accounts SET balance = balance + 200 WHERE id = 2')
    cursor.execute('COMMIT')
except Exception:
    cursor.execute('ROLLBACK')
    raise

cursor.execute('SELECT * FROM accounts')
print(cursor.fetchall())
conn.close()

[(1, 800), (2, 700)]

Если вторая UPDATE приведёт к ошибке, ROLLBACK отменит первую. Такая схема защищает от частичного изменения данных.

Отражение таблицы в SQLAlchemy Core

Отражение позволяет загрузить структуру существующей таблицы в объект Table без описания колонок вручную.

from sqlalchemy import create_engine, MetaData, Table, select

engine = create_engine('sqlite:///reflect.db')
with engine.begin() as conn:
    conn.exec_driver_sql('DROP TABLE IF EXISTS authors')
    conn.exec_driver_sql('CREATE TABLE authors (id INTEGER PRIMARY KEY, name TEXT)')
    conn.exec_driver_sql("INSERT INTO authors (id, name) VALUES (1, 'Лев Толстой')")

metadata = MetaData()
authors = Table('authors', metadata, autoload_with=engine)

with engine.connect() as conn:
    result = conn.execute(select(authors.c.name).where(authors.c.id == 1))
    print(result.scalar())

Лев Толстой

Благодаря autoload_with загружаются названия колонок и их типы. Указание where строится через сравнение объектов колонок.

Связь многие-к-одному в SQLAlchemy ORM

ORM позволяет работать с графом объектов. При обращении к author.books загружаются связанные строки из таблицы books.

from sqlalchemy import create_engine, Column, Integer, String, ForeignKey
from sqlalchemy.orm import declarative_base, Session, relationship

Base = declarative_base()

class Author(Base):
    __tablename__ = 'authors'
    id = Column(Integer, primary_key=True)
    name = Column(String(50))
    books = relationship('Book', back_populates='author')

class Book(Base):
    __tablename__ = 'books'
    id = Column(Integer, primary_key=True)
    title = Column(String(100))
    author_id = Column(Integer, ForeignKey('authors.id'))
    author = relationship('Author', back_populates='books')

engine = create_engine('sqlite:///demo_orm.db')
Base.metadata.drop_all(engine)
Base.metadata.create_all(engine)

with Session(engine) as session:
    author = Author(name='Анна Ахматова')
    author.books.append(Book(title='Вечер'))
    author.books.append(Book(title='Чётки'))
    session.add(author)
    session.commit()

with Session(engine) as session:
    author = session.query(Author).filter_by(name='Анна Ахматова').first()
    print(author.name)
    for book in author.books:
        print(f' - {book.title}')

Анна Ахматова
 - Вечер
 - Чётки

Метод drop_all перед create_all очищает таблицы, поэтому повторный запуск не создаёт дубликатов.

Массовая загрузка данных в PostgreSQL через COPY

Оператор COPY увеличивает скорость вставки большого числа строк по сравнению с отдельными INSERT. Данные передаются через специальный буфер.

from psycopg2 import connect
from io import StringIO

conn = connect(
    host='127.0.0.1',
    dbname='shop',
    user='shop_user',
    password='secret'
)
buffer = StringIO()
buffer.write('1\tКружка\t250\n')
buffer.write('2\tТарелка\t400\n')
buffer.seek(0)

with conn.cursor() as cursor:
    cursor.execute('CREATE TEMP TABLE products (id INTEGER PRIMARY KEY, name TEXT, price NUMERIC)')
    cursor.copy_expert(
        'COPY products (id, name, price) FROM STDIN WITH (FORMAT text)',
        buffer
    )
    cursor.execute('SELECT count(*) FROM products')
    print(cursor.fetchone()[0])
conn.commit()
conn.close()

2

Временная таблица исчезнет после закрытия соединения, поэтому пример безопасен для реальной базы данных.