sqlalchemy_cockroachdb_crud_with_transaction_retry.py

1import uuid
2from sqlalchemy import create_engine, Column, Integer, String
3from sqlalchemy.orm import sessionmaker, declarative_base
4from sqlalchemy_cockroachdb import run_transaction
5
6Base = declarative_base()
7
8# The Account class corresponds to the "accounts" table.
9class Account(Base):
10    __tablename__ = 'accounts'
11    id = Column(Integer, primary_key=True)
12    balance = Column(Integer)
13
14# Create an engine to communicate with the database.
15# Note: For CockroachDB, use the 'cockroachdb://' dialect prefix.
16engine = create_engine('cockroachdb://root@localhost:26257/defaultdb?sslmode=disable')
17
18# Create the table in the database.
19Base.metadata.create_all(engine)
20
21# Use run_transaction to handle retry logic automatically.
22def create_accounts(session):
23    session.add_all([
24        Account(id=1, balance=1000),
25        Account(id=2, balance=250)
26    ])
27
28run_transaction(sessionmaker(bind=engine), create_accounts)
29
30# Print the balances.
31def print_balances(session):
32    accounts = session.query(Account).all()
33    print("Balances at start:")
34    for account in accounts:
35        print(f"Account {account.id}: {account.balance}")
36
37run_transaction(sessionmaker(bind=engine), print_balances)