test_psycopg2_isolation_level.py

import psycopg2


dsn = "postgresql://user:password@localhost:5432/"
db_name = "test_db"

conn = psycopg2.connect(dsn)
try:
    conn.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
    with conn.cursor() as cur:
        try:
            cur.execute(f"DROP database {db_name};")
        except psycopg2.errors.InvalidCatalogName:
            pass
        cur.execute(f"CREATE database {db_name};")
finally:
    conn.close()

dsn = dsn + db_name

with psycopg2.connect(dsn) as conn1:
    with conn1.cursor() as cur:
        cur.execute(
            """
            CREATE TABLE users (
                id          integer PRIMARY KEY GENERATED BY DEFAULT AS IDENTITY,
                email       varchar(100) NOT NULL,
                balance     integer NOT NULL
            );
            """
        )
        cur.execute("INSERT INTO users (email, balance) VALUES (%s, %s)", ("[email protected]", 100))

level = psycopg2.extensions.ISOLATION_LEVEL_READ_COMMITTED
# level = psycopg2.extensions.ISOLATION_LEVEL_REPEATABLE_READ
# level = psycopg2.extensions.ISOLATION_LEVEL_SERIALIZABLE

with psycopg2.connect(dsn) as conn1:
    conn1.set_isolation_level(level)
    with conn1.cursor() as cur1:
        cur1.execute("SELECT balance FROM users WHERE email=%s;", ("[email protected]",))
        user1_balance = cur1.fetchone()[0]

        with psycopg2.connect(dsn) as conn2:
            conn2.set_isolation_level(level)
            with conn2.cursor() as cur2:
                cur2.execute("UPDATE users SET balance = %s WHERE email=%s;", (0, "[email protected]"))

        # psycopg2.errors.SerializationFailure is raised if the isolation level is ISOLATION_LEVEL_REPEATABLE_READ or ISOLATION_LEVEL_SERIALIZABLE
        cur1.execute("UPDATE users SET balance = %s WHERE email=%s;", (user1_balance + 100, "[email protected]"))