cockroachdb_crud_operations_with_psycopg3_and_retry.py

1import logging
2import os
3import time
4
5import psycopg
6from psycopg.errors import SerializationFailure, Error
7from psycopg.rows import dict_row
8
9def delete_accounts(conn):
10    with conn.cursor() as cur:
11        cur.execute("DELETE FROM accounts")
12        logging.debug("delete_accounts(): status message: %s", cur.statusmessage)
13    conn.commit()
14
15def print_balances(conn):
16    with conn.cursor() as cur:
17        print(f"Balances at {time.asctime()}:")
18        for row in cur.execute("SELECT id, balance FROM accounts"):
19            print("account id: {0}  balance: ${1:,.2f}".format(row[0], row[1]))
20
21def transfer_funds(conn, frm, to, amount):
22    with conn.cursor() as cur:
23        # Check the current balance.
24        cur.execute("SELECT balance FROM accounts WHERE id = %s", (frm,))
25        from_balance = cur.fetchone()[0]
26        if from_balance < amount:
27            raise RuntimeError(
28                f"Insufficient funds in account {frm} for transfer"
29            )
30
31        # Perform the transfer.
32        cur.execute(
33            "UPDATE accounts SET balance = balance - %s WHERE id = %s", (amount, frm)
34        )
35        cur.execute(
36            "UPDATE accounts SET balance = balance + %s WHERE id = %s", (amount, to)
37        )
38
39    conn.commit()
40    logging.debug("transfer_funds(): status message: %s", cur.statusmessage)
41
42def run_transaction(conn, op, max_retries=3):
43    """
44    Execute the operation within a transaction, with retries for serialization
45    failures.
46    """
47    for retry in range(1, max_retries + 1):
48        try:
49            op(conn)
50            return
51        except SerializationFailure as e:
52            # This is a retryable error.
53            logging.debug("run_transaction(): killed by serialization failure")
54            # Sleep a bit before retrying.
55            time.sleep(0.1 * retry)
56            continue
57        except Error as e:
58            logging.debug("run_transaction(): non-retryable error: %s", e)
59            raise e
60
61    raise RuntimeError(f"Transaction failed after {max_retries} retries")
62
63def main():
64    # Use the connection string provided by the CockroachDB console.
65    # Typically formatted as: postgresql://<user>:<password>@<host>:<port>/<database>?sslmode=verify-full
66    conn_string = os.expandvars(os.environ.get("DATABASE_URL"))
67
68    try:
69        with psycopg.connect(conn_string, application_name="$ docs_simplecrud_python") as conn:
70            # Create the "accounts" table.
71            with conn.cursor() as cur:
72                cur.execute(
73                    "CREATE TABLE IF NOT EXISTS accounts (id UUID PRIMARY KEY DEFAULT gen_random_uuid(), balance DECIMAL)"
74                )
75            conn.commit()
76
77            # Delete any existing accounts.
78            run_transaction(conn, delete_accounts)
79
80            # Insert some accounts.
81            with conn.cursor() as cur:
82                cur.execute(
83                    "INSERT INTO accounts (id, balance) VALUES (%s, %s), (%s, %s)",
84                    ("15754f73-2d2c-473d-9d48-3571d88a10d9", 1000, "9b266e79-8a8b-402e-b615-5205562d9899", 250),
85                )
86            conn.commit()
87
88            # Print the balances.
89            print_balances(conn)
90
91            # Transfer funds between accounts.
92            run_transaction(
93                conn,
94                lambda conn: transfer_funds(
95                    conn, "15754f73-2d2c-473d-9d48-3571d88a10d9", "9b266e79-8a8b-402e-b615-5205562d9899", 100
96                ),
97            )
98
99            # Print the balances.
100            print_balances(conn)
101
102    except Exception as e:
103        logging.fatal("database connection failed")
104        logging.fatal(e)
105        return
106
107if __name__ == "__main__":
108    main()