This tutorial shows you how build a simple Python application with CockroachDB using a PostgreSQL-compatible driver or ORM.
We have tested the Python psycopg2 driver and the SQLAlchemy ORM enough to claim beta-level support, so those are featured here. If you encounter problems, please open an issue with details to help us make progress toward full support.
Before you begin
- Install CockroachDB.
- Start up a secure or insecure local cluster.
- Choose the instructions that correspond to whether your cluster is secure or insecure:
Step 1. Install the psycopg2 driver
To install the Python psycopg2 driver, run the following command:
$ pip install psycopg2
For other ways to install psycopg2, see the official documentation.
Step 2. Create the maxroach user and bank database
Start the built-in SQL client:
$ cockroach sql --certs-dir=certs
In the SQL shell, issue the following statements to create the maxroach user and bank database:
> CREATE USER IF NOT EXISTS maxroach;
> CREATE DATABASE bank;
Give the maxroach user the necessary permissions:
> GRANT ALL ON DATABASE bank TO maxroach;
Exit the SQL shell:
> \q
Step 3. Generate a certificate for the maxroach user
Create a certificate and key for the maxroach user by running the following command. The code samples will run as this user.
$ cockroach cert create-client maxroach --certs-dir=certs --ca-key=my-safe-directory/ca.key
Step 4. Run the Python code
Now that you have a database and a user, you'll run the code shown below to:
- Create a table and insert some rows
- Read and update values as an atomic transaction
Basic statements
First, use the following code to connect as the maxroach user and execute some basic SQL statements, creating a table, inserting rows, and reading and printing the rows.
Download the basic-sample.py file, or create the file yourself and copy the code into it.
# Import the driver.
import psycopg2
# Connect to the "bank" database.
conn = psycopg2.connect(
database='bank',
user='maxroach',
sslmode='require',
sslrootcert='certs/ca.crt',
sslkey='certs/client.maxroach.key',
sslcert='certs/client.maxroach.crt',
port=26257,
host='localhost'
)
# Make each statement commit immediately.
conn.set_session(autocommit=True)
# Open a cursor to perform database operations.
cur = conn.cursor()
# Create the "accounts" table.
cur.execute("CREATE TABLE IF NOT EXISTS accounts (id INT PRIMARY KEY, balance INT)")
# Insert two rows into the "accounts" table.
cur.execute("INSERT INTO accounts (id, balance) VALUES (1, 1000), (2, 250)")
# Print out the balances.
cur.execute("SELECT id, balance FROM accounts")
rows = cur.fetchall()
print('Initial balances:')
for row in rows:
print([str(cell) for cell in row])
# Close the database connection.
cur.close()
conn.close()
Then run the code:
$ python basic-sample.py
The output should be:
Initial balances:
['1', '1000']
['2', '250']
Transaction (with retry logic)
Next, use the following code to again connect as the maxroach user but this time execute a batch of statements as an atomic transaction to transfer funds from one account to another, where all included statements are either committed or aborted.
Download the txn-sample.py file, or create the file yourself and copy the code into it.
# Import the driver.
import psycopg2
import psycopg2.errorcodes
# Connect to the cluster.
conn = psycopg2.connect(
database='bank',
user='maxroach',
sslmode='require',
sslrootcert='certs/ca.crt',
sslkey='certs/client.maxroach.key',
sslcert='certs/client.maxroach.crt',
port=26257,
host='localhost'
)
def onestmt(conn, sql):
with conn.cursor() as cur:
cur.execute(sql)
# Wrapper for a transaction.
# This automatically re-calls "op" with the open transaction as an argument
# as long as the database server asks for the transaction to be retried.
def run_transaction(conn, op):
with conn:
onestmt(conn, "SAVEPOINT cockroach_restart")
while True:
try:
# Attempt the work.
op(conn)
# If we reach this point, commit.
onestmt(conn, "RELEASE SAVEPOINT cockroach_restart")
break
except psycopg2.OperationalError as e:
if e.pgcode != psycopg2.errorcodes.SERIALIZATION_FAILURE:
# A non-retryable error; report this up the call stack.
raise e
# Signal the database that we'll retry.
onestmt(conn, "ROLLBACK TO SAVEPOINT cockroach_restart")
# The transaction we want to run.
def transfer_funds(txn, frm, to, amount):
with txn.cursor() as cur:
# Check the current balance.
cur.execute("SELECT balance FROM accounts WHERE id = " + str(frm))
from_balance = cur.fetchone()[0]
if from_balance < amount:
raise "Insufficient funds"
# Perform the transfer.
cur.execute("UPDATE accounts SET balance = balance - %s WHERE id = %s",
(amount, frm))
cur.execute("UPDATE accounts SET balance = balance + %s WHERE id = %s",
(amount, to))
# Execute the transaction.
run_transaction(conn, lambda conn: transfer_funds(conn, 1, 2, 100))
with conn:
with conn.cursor() as cur:
# Check account balances.
cur.execute("SELECT id, balance FROM accounts")
rows = cur.fetchall()
print('Balances after transfer:')
for row in rows:
print([str(cell) for cell in row])
# Close communication with the database.
conn.close()
Then run the code:
$ python txn-sample.py
The output should be:
Balances after transfer:
['1', '900']
['2', '350']
To verify that funds were transferred from one account to another, start the built-in SQL client:
$ cockroach sql --certs-dir=certs --database=bank
To check the account balances, issue the following statement:
> SELECT id, balance FROM accounts;
+----+---------+
| id | balance |
+----+---------+
| 1 | 900 |
| 2 | 350 |
+----+---------+
(2 rows)
Step 2. Create the maxroach user and bank database
Start the built-in SQL client:
$ cockroach sql --insecure
In the SQL shell, issue the following statements to create the maxroach user and bank database:
> CREATE USER IF NOT EXISTS maxroach;
> CREATE DATABASE bank;
Give the maxroach user the necessary permissions:
> GRANT ALL ON DATABASE bank TO maxroach;
Exit the SQL shell:
> \q
Step 3. Run the Python code
Now that you have a database and a user, you'll run the code shown below to:
- Create a table and insert some rows
- Read and update values as an atomic transaction
Basic statements
First, use the following code to connect as the maxroach user and execute some basic SQL statements, creating a table, inserting rows, and reading and printing the rows.
Download the basic-sample.py file, or create the file yourself and copy the code into it.
# Import the driver.
import psycopg2
# Connect to the "bank" database.
conn = psycopg2.connect(
database='bank',
user='maxroach',
sslmode='disable',
port=26257,
host='localhost'
)
# Make each statement commit immediately.
conn.set_session(autocommit=True)
# Open a cursor to perform database operations.
cur = conn.cursor()
# Create the "accounts" table.
cur.execute("CREATE TABLE IF NOT EXISTS accounts (id INT PRIMARY KEY, balance INT)")
# Insert two rows into the "accounts" table.
cur.execute("INSERT INTO accounts (id, balance) VALUES (1, 1000), (2, 250)")
# Print out the balances.
cur.execute("SELECT id, balance FROM accounts")
rows = cur.fetchall()
print('Initial balances:')
for row in rows:
print([str(cell) for cell in row])
# Close the database connection.
cur.close()
conn.close()
Then run the code:
$ python basic-sample.py
The output should be:
Initial balances:
['1', '1000']
['2', '250']
Transaction (with retry logic)
Next, use the following code to again connect as the maxroach user but this time execute a batch of statements as an atomic transaction to transfer funds from one account to another, where all included statements are either committed or aborted.
Download the txn-sample.py file, or create the file yourself and copy the code into it.
# Import the driver.
import psycopg2
import psycopg2.errorcodes
# Connect to the cluster.
conn = psycopg2.connect(
database='bank',
user='maxroach',
sslmode='disable',
port=26257,
host='localhost'
)
def onestmt(conn, sql):
with conn.cursor() as cur:
cur.execute(sql)
# Wrapper for a transaction.
# This automatically re-calls "op" with the open transaction as an argument
# as long as the database server asks for the transaction to be retried.
def run_transaction(conn, op):
with conn:
onestmt(conn, "SAVEPOINT cockroach_restart")
while True:
try:
# Attempt the work.
op(conn)
# If we reach this point, commit.
onestmt(conn, "RELEASE SAVEPOINT cockroach_restart")
break
except psycopg2.OperationalError as e:
if e.pgcode != psycopg2.errorcodes.SERIALIZATION_FAILURE:
# A non-retryable error; report this up the call stack.
raise e
# Signal the database that we'll retry.
onestmt(conn, "ROLLBACK TO SAVEPOINT cockroach_restart")
# The transaction we want to run.
def transfer_funds(txn, frm, to, amount):
with txn.cursor() as cur:
# Check the current balance.
cur.execute("SELECT balance FROM accounts WHERE id = " + str(frm))
from_balance = cur.fetchone()[0]
if from_balance < amount:
raise "Insufficient funds"
# Perform the transfer.
cur.execute("UPDATE accounts SET balance = balance - %s WHERE id = %s",
(amount, frm))
cur.execute("UPDATE accounts SET balance = balance + %s WHERE id = %s",
(amount, to))
# Execute the transaction.
run_transaction(conn, lambda conn: transfer_funds(conn, 1, 2, 100))
with conn:
with conn.cursor() as cur:
# Check account balances.
cur.execute("SELECT id, balance FROM accounts")
rows = cur.fetchall()
print('Balances after transfer:')
for row in rows:
print([str(cell) for cell in row])
# Close communication with the database.
conn.close()
Then run the code:
$ python txn-sample.py
The output should be:
Balances after transfer:
['1', '900']
['2', '350']
To verify that funds were transferred from one account to another, start the built-in SQL client:
$ cockroach sql --insecure --database=bank
To check the account balances, issue the following statement:
> SELECT id, balance FROM accounts;
+----+---------+
| id | balance |
+----+---------+
| 1 | 900 |
| 2 | 350 |
+----+---------+
(2 rows)
What's next?
Read more about using the Python psycopg2 driver.
You might also be interested in using a local cluster to explore the following CockroachDB benefits:
- Client Connection Parameters
- Data Replication
- Fault Tolerance & Recovery
- Automatic Rebalancing
- Cross-Cloud Migration
- Follow-the-Workload
- Automated Operations
Was this helpful?
