Beyond SQLAlchemy's abstraction — Postgres offers speed improvements in a sea of complexity. This is a swim in that complexity — a swim in search of the elementary, dear Watson.
Setting this up will take some time. You can jump to the conclusion or sit back and enjoy observing someone else set up a project and finally get to the point where the work is done…
To access Postgres from python you can use the psycopg package. Version 3 affords an async interface that works effortlessly with SQLAlchemy 2.0 async extension. So lets install both, but first we need an environment:
uv venv
source .venv/bin/activate
uv pip install "psycopg[binary]" sqlalchemy
uv — install this, if you can, in your base python and see the most wonderful speed improvements when creating virtual environments and installing packages from pypi.
Now we can connect to Postgres using:
from asyncio import current_task
from sqlalchemy.ext.asyncio import (
async_scoped_session,
async_sessionmaker,
create_async_engine,
)
async_session_factory = async_sessionmaker(expire_on_commit=False)
Session = async_scoped_session(async_session_factory, scopefunc=current_task)
def init_db(echo: bool = False):
"""initialise our db"""
db_url = os.getenv(
'DATABASE_URL',
'postgresql://postgres:secret@localhost/folly_db'
).replace('postgresql://', 'postgresql+psycopg://', 1)
engine = create_async_engine(db_url, echo=echo)
async_session_factory.configure(bind=engine)
return db_url
init_db is called at runtime, enabling us to import this context and use Session before it is configured. Why mangle the DATABASE_URL, because we can use it with the psycopg package directly:
db_url = os.getenv(
'DATABASE_URL',
'postgresql://postgres:secret@localhost/folly_db'
)
with psycopg.connect(db_url) as conn:
...
Now that we can connect, we need something to connect to. I'm using Docker, so our Dockerfile:
FROM postgres:16
RUN apt-get update
RUN apt-get -y install python3 postgresql-plpython3-16
And our docker-compose.yml:
version: '3.9'
services:
db:
build: .
restart: unless-stopped
ports:
- '5432:5432'
environment:
- POSTGRES_PASSWORD=secret
- POSTGRES_USER=postgres
- POSTGRES_DB=folly_db
volumes:
- ./postgres:/var/lib/postgresql/data
- ./start-up:/docker-entrypoint-initdb.d
The start-up file installs plpython3u:
CREATE EXTENSION IF NOT EXISTS plpython3u;
Now we need a schema:
from sqlalchemy import (
Column,
DateTime,
ForeignKey,
Index,
Integer,
MetaData,
String,
Table,
)
metadata = MetaData()
user = Table(
'user',
metadata,
Column('id', Integer, primary_key=True),
Column('email', String(128)),
)
log = Table(
'log',
metadata,
Column('id', Integer, primary_key=True),
Column('created', DateTime, nullable=False),
Column(
'activity_id',
Integer,
ForeignKey(
'activity.id',
name='fk_log_activity_id',
ondelete='CASCADE'
),
nullable=False,
),
)
activity = Table(
'activity',
metadata,
Column('id', Integer, primary_key=True),
Column('name', String(255), nullable=False),
Column(
'user_id',
Integer,
ForeignKey('user.id', name='fk_activity_user_id', ondelete='CASCADE'),
nullable=False,
),
Index('activity_user_idx', 'name', 'user_id', unique=True),
)
Now we can use alembic to create a migration script:
uv pip install alembic
alembic init -template async ./src/scripts
Change the alembic.ini to reflect our DATABASE_URL with SQLAlchemy protocol, and edit the src/scripts/env.py to import our tables and point to our metadata. Then we can auto generate our first script and upgrade the database:
alembic revision --autogenerate -m 'first pass'
alembic upgrade head
Now we get to do work:
@dataclass
class LogEntry:
id: int
activity: str
created: datetime
user_id: int
async def log(activity: str) -> LogEntry:
"""log an activity for current user starting now"""
user = context.get_auth_user()
async with context.Session() as session:
log_id = await session.scalar(
insert(tables.log)
.values(
activity_id=await id_for_activity(
session, user.id, activity
),
created=(created := utils.now()),
)
.returning(tables.log.c.id)
)
result = LogEntry(log_id, activity, created, user.id)
await session.execute(sa.text(f'NOTIFY logs, {utils.dumps(result)!r}'))
await session.commit()
return result
Our context also holds:
import contextvars
class AccessControlException(Exception):
"""You could not do that"""
USER = contextvars.ContextVar('_USER')
def get_auth_user():
"""return the context user or raise an exception"""
user = USER.get(None)
if user is None:
raise AccessControlException('you are not logged in')
return user
And id_for_activity is a find or create:
async def id_for_activity(session, user_id: int, name: str) -> int:
"""find or create an activity for user with name and return id"""
result = await session.execute(
insert(tables.activity)
.values(user_id=user_id, name=name)
.on_conflict_do_nothing()
.returning(tables.activity.c.id)
)
if row := result.first():
await session.commit()
LOGGER.debug('created activity: %s', row.id)
return row.id
result = await session.scalar(
sa.select(tables.activity.c.id).where(
tables.activity.c.name == name,
tables.activity.c.user_id == user_id,
)
)
LOGGER.debug('using activity: %s', result)
return result
We're using Postgres NOTIFY to let everyone know that a log entry has been made. But we are not using your standard INSERT, we're using sqlalchemy.dialects.postgresql.insert which allows us to ignore conflicts. You'll also notice the handy use of the walrus operator.
Let's test it:
uv pip install pytest pytest-asyncio pytest-benchmark
In our src/tests we'll create a conftest.py which setup our database using the migration scripts as a pytest fixture:
@pytest.fixture(scope='module')
def init_db():
db_url = context.init_db()
alembic_config = config.Config(
config_args={
'script_location': 'src/scripts',
'prepend_sys_path': 'src',
'version_path_separator': 'os',
'sqlalchemy.url': db_url,
}
)
command.downgrade(alembic_config, 'base')
command.upgrade(alembic_config, 'head')
Now we can write a test:
import logging
import pytest
from folly import api, context
class MockUser:
id = 1
@pytest.mark.asyncio
async def test_logs(init_db):
context.USER.set(MockUser)
result = await api.log('foo')
assert result.id == 1
assert result.activity == 'foo'
That should work. So, we're not using SQLAlchemy ORM, we're using the expression language. This should be as fast as using the driver itself.
How can we improve performance? Stored procedures: plpython3u.
CREATE OR REPLACE FUNCTION enter_log (activity_name VARCHAR, user_id INT)
RETURNS TABLE(id INT, activity VARCHAR, created TIMESTAMP WITH TIME ZONE, user_id INT)
AS $$
import datetime
import json
plan = plpy.prepare(
"INSERT INTO activity (name, user_id) VALUES ($1, $2) ON CONFLICT DO NOTHING RETURNING id",
["VARCHAR", "INT"]
)
rv = plpy.execute(plan, [activity_name, user_id])
if len(rv) == 0:
plan = plpy.prepare(
"SELECT id FROM activity WHERE name=$1 AND user_id = $2",
["VARCHAR", "INT"]
)
rv = plpy.execute(plan, [activity_name, user_id])
activity_id = int(rv[0]['id'])
created = datetime.datetime.now()
trunc = int(created.microsecond / 1000) * 1000
created = created.replace(microsecond=trunc, tzinfo=created.tzinfo)
plan = plpy.prepare(
"INSERT INTO log (activity_id, created) VALUES ($1, $2) RETURNING id",
["INT", "TIMESTAMP"]
)
rv = plpy.execute(plan, [activity_id, created])
return_id = rv[0]['id']
result = {
'id': return_id,
'activity': activity_name,
'created': created.isoformat(),
'user_id': user_id
}
notify = json.dumps(result)
plpy.execute(f"NOTIFY logs, {notify!r}")
return [result]
$$ LANGUAGE plpython3u;
Our function inserts or selects to get the activity_id, it then inserts our log entry, creates the result, notifies, and returns the result. And how do we call this?
async def enter_log(activity: str) -> LogEntry:
"""log an activity for current user starting now"""
user = context.get_auth_user()
async with context.Session() as session:
id, _, created, _ = await session.scalar(
sa.select(sa.func.enter_log(activity, user.id))
)
await session.commit()
return LogEntry(
id=int(id),
activity=activity,
created=datetime.fromisoformat(created),
user_id=user.id,
)
The return type of my procedure is a Table — and I have specified that the column id is an INT, and yet it returns a string. There are plenty of other return types, but they all become strings when they get back to SQLAlchemy. I tried declaring a return type in SQL — still a big string — and then used JSON to build a big string, but decided to use table as it produced several little strings… The documentation is minimal at best and the world does not seem to be using this idea — but at least the error messages are meaningful and help you in the right direction.
If you learn SQLAlchemy, and then learn psycopg and now learn plpython3u — you'll have learned three different ways to span the distance between Postgres and Python! And Postgres is responsible for two of them!
How do we integrate this amazing function — you could add to the docker-entrypoint-initdb.d — but to make it easier to debug I added this to the alembic scripts. The migration looks as follows:
def upgrade():
with open('src/sql/enter_log.sql') as file:
op.execute(sa.text(file.read()))
def downgrade():
op.execute(sa.text('DROP FUNCTION IF EXISTS enter_log'))
I put the function in a separate sql file ease of editing. And we can test it:
@pytest.mark.asyncio
async def test_enter_log(init_db):
context.USER.set(MockUser)
result = await api.enter_log('bar')
assert result.id == 2
assert result.activity == 'bar'
And if we benchmark it?
import asyncio
from folly import api, context
class MockUser:
id = 1
context.USER.set(MockUser)
def run_log(what):
return asyncio.run(api.log(what))
def run_enter_log(what):
return asyncio.run(api.enter_log(what))
def test_log(init_db, benchmark):
benchmark(run_log, 'foobert')
def test_enter_log(init_db, benchmark):
benchmark(run_enter_log, 'dilbert')
We get the following results:
-------- benchmark: 2 tests --------
Name (time in us) Min Max Mean StdDev Median IQR Outliers OPS Rounds Iterations
test_enter_log 940.0 1,929.5 1,040.5 142.1 1,012.1 71.2 3;5 961.0 52 1
test_log 1,695.7 2,977.1 2,031.7 328.0 1,875.4 357.4 7;2 492.1 30 1
1.0 : 1.95 — wow!
In a heavily used system that could make all sorts of savings.
We have created an environment, installed packages to connect, migrate and test, developed a schema and logic, prepared a test database with docker and finally optimised our usage with a stored procedure.
It should by this point be elementary that you don't pre-optimise. We wrote the logic in SQLAlchemy — which is portable — but then we broke portability by using returning and on_conflict_do_nothing and notify. But since our log function could call four statements and transaction control — six round trips to the server — using a stored procedure makes sense. But then you are committing yourself to a particular database.
SQLAlchemy is brilliant and yet we should not let that blind us to using vendor specific extensions or diving head first into complete adoption of the vendor's solution — just do it for a reason, wisely — surely that is elementary, dear Watson.