Pythontransactions_demo.py
Copy
1
import psycopg2
2
import psycopg2.extras
3
from contextlib import contextmanager
4
 
5
conn = psycopg2.connect("host=localhost dbname=learndb user=postgres")
6
conn.autocommit = False  # Manual transaction control
7
 
8
def setup():
9
    with conn.cursor() as cur:
10
        cur.execute("DROP TABLE IF EXISTS accounts;")
11
        cur.execute("""
12
            CREATE TABLE accounts (
13
                id      SERIAL PRIMARY KEY,
14
                name    TEXT NOT NULL,
15
                balance NUMERIC(12, 2) NOT NULL CHECK (balance >= 0)
16
            );
17
            INSERT INTO accounts (name, balance) VALUES ('Alice', 1000), ('Bob', 500);
18
        """)
19
    conn.commit()
20
 
21
# ACID bank transfer
22
def transfer(from_id: int, to_id: int, amount: float) -> None:
23
    with conn.cursor() as cur:
24
        try:
25
            # Lock both rows in consistent order to avoid deadlocks
26
            ids = sorted([from_id, to_id])
27
            cur.execute(
28
                "SELECT id, balance FROM accounts WHERE id = ANY(%s) ORDER BY id FOR UPDATE",
29
                (ids,)
30
            )
31
            rows = {row[0]: row[1] for row in cur.fetchall()}
32
 
33
            if rows[from_id] < amount:
34
                raise ValueError(f"Insufficient funds: {rows[from_id]} < {amount}")
35
 
36
            cur.execute("UPDATE accounts SET balance = balance - %s WHERE id = %s", (amount, from_id))
37
            cur.execute("UPDATE accounts SET balance = balance + %s WHERE id = %s", (amount, to_id))
38
            conn.commit()
39
            print(f"Transferred ${amount} from account {from_id} to {to_id}")
40
        except Exception as e:
41
            conn.rollback()
42
            print(f"Transfer failed: {e}")
43
            raise
44
 
45
def print_balances():
46
    with conn.cursor() as cur:
47
        cur.execute("SELECT name, balance FROM accounts ORDER BY id")
48
        for name, balance in cur.fetchall():
49
            print(f"  {name}: ${balance}")
50
 
51
setup()
52
print("Before transfer:")
53
print_balances()
54
 
55
transfer(1, 2, 250)  # Alice -> Bob
56
 
57
print("After transfer:")
58
print_balances()
59
 
60
try:
61
    transfer(2, 1, 1000)  # Should fail (insufficient funds)
62
except ValueError:
63
    pass
64
 
65
print("After failed transfer:")
66
print_balances()
67
 
68
conn.close()
69