Build a Python App with CockroachDB

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

Step 1. Install the psycopg2 driver

To install the Python psycopg2 driver, run the following command:

  1. $ 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:

  1. $ cockroach sql --certs-dir=certs

In the SQL shell, issue the following statements to create the maxroach user and bank database:

  1. > CREATE USER IF NOT EXISTS maxroach;
  1. > CREATE DATABASE bank;

Give the maxroach user the necessary permissions:

  1. > GRANT ALL ON DATABASE bank TO maxroach;

Exit the SQL shell:

  1. > \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.

  1. $ 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.

  1. # Import the driver.
  2. import psycopg2
  4. # Connect to the "bank" database.
  5. conn = psycopg2.connect(
  6.     database='bank',
  7.     user='maxroach',
  8.     sslmode='require',
  9.     sslrootcert='certs/ca.crt',
  10.     sslkey='certs/client.maxroach.key',
  11.     sslcert='certs/client.maxroach.crt',
  12.     port=26257,
  13.     host='localhost'
  14. )
  16. # Make each statement commit immediately.
  17. conn.set_session(autocommit=True)
  19. # Open a cursor to perform database operations.
  20. cur = conn.cursor()
  22. # Create the "accounts" table.
  23. cur.execute("CREATE TABLE IF NOT EXISTS accounts (id INT PRIMARY KEY, balance INT)")
  25. # Insert two rows into the "accounts" table.
  26. cur.execute("INSERT INTO accounts (id, balance) VALUES (1, 1000), (2, 250)")
  28. # Print out the balances.
  29. cur.execute("SELECT id, balance FROM accounts")
  30. rows = cur.fetchall()
  31. print('Initial balances:')
  32. for row in rows:
  33.     print([str(cell) for cell in row])
  35. # Close the database connection.
  36. cur.close()
  37. conn.close()

Then run the code:

  1. $ python basic-sample.py

The output should be:

  1. Initial balances:
  2. ['1', '1000']
  3. ['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.

Note:
CockroachDB may require the client to retry a transaction in case of read/write contention. CockroachDB provides a generic retry function that runs inside a transaction and retries it as needed. You can copy and paste the retry function from here into your code.

  1. # Import the driver.
  2. import psycopg2
  3. import psycopg2.errorcodes
  5. # Connect to the cluster.
  6. conn = psycopg2.connect(
  7.     database='bank',
  8.     user='maxroach',
  9.     sslmode='require',
  10.     sslrootcert='certs/ca.crt',
  11.     sslkey='certs/client.maxroach.key',
  12.     sslcert='certs/client.maxroach.crt',
  13.     port=26257,
  14.     host='localhost'
  15. )
  17. def onestmt(conn, sql):
  18.     with conn.cursor() as cur:
  19.         cur.execute(sql)
  22. # Wrapper for a transaction.
  23. # This automatically re-calls "op" with the open transaction as an argument
  24. # as long as the database server asks for the transaction to be retried.
  25. def run_transaction(conn, op):
  26.     with conn:
  27.         onestmt(conn, "SAVEPOINT cockroach_restart")
  28.         while True:
  29.             try:
  30.                 # Attempt the work.
  31.                 op(conn)
  33.                 # If we reach this point, commit.
  34.                 onestmt(conn, "RELEASE SAVEPOINT cockroach_restart")
  35.                 break
  37.             except psycopg2.OperationalError as e:
  38.                 if e.pgcode != psycopg2.errorcodes.SERIALIZATION_FAILURE:
  39.                     # A non-retryable error; report this up the call stack.
  40.                     raise e
  41.                 # Signal the database that we'll retry.
  42.                 onestmt(conn, "ROLLBACK TO SAVEPOINT cockroach_restart")
  45. # The transaction we want to run.
  46. def transfer_funds(txn, frm, to, amount):
  47.     with txn.cursor() as cur:
  49.         # Check the current balance.
  50.         cur.execute("SELECT balance FROM accounts WHERE id = " + str(frm))
  51.         from_balance = cur.fetchone()[0]
  52.         if from_balance < amount:
  53.             raise "Insufficient funds"
  55.         # Perform the transfer.
  56.         cur.execute("UPDATE accounts SET balance = balance - %s WHERE id = %s",
  57.                     (amount, frm))
  58.         cur.execute("UPDATE accounts SET balance = balance + %s WHERE id = %s",
  59.                     (amount, to))
  62. # Execute the transaction.
  63. run_transaction(conn, lambda conn: transfer_funds(conn, 1, 2, 100))
  66. with conn:
  67.     with conn.cursor() as cur:
  68.         # Check account balances.
  69.         cur.execute("SELECT id, balance FROM accounts")
  70.         rows = cur.fetchall()
  71.         print('Balances after transfer:')
  72.         for row in rows:
  73.             print([str(cell) for cell in row])
  75. # Close communication with the database.
  76. conn.close()

Then run the code:

  1. $ python txn-sample.py

The output should be:

  1. Balances after transfer:
  2. ['1', '900']
  3. ['2', '350']

To verify that funds were transferred from one account to another, start the built-in SQL client:

  1. $ cockroach sql --certs-dir=certs --database=bank

To check the account balances, issue the following statement:

  1. > SELECT id, balance FROM accounts;
  1. +----+---------+
  2. | id | balance |
  3. +----+---------+
  4. |  1 |     900 |
  5. |  2 |     350 |
  6. +----+---------+
  7. (2 rows)

Step 2. Create the maxroach user and bank database

Start the built-in SQL client:

  1. $ cockroach sql --insecure

In the SQL shell, issue the following statements to create the maxroach user and bank database:

  1. > CREATE USER IF NOT EXISTS maxroach;
  1. > CREATE DATABASE bank;

Give the maxroach user the necessary permissions:

  1. > GRANT ALL ON DATABASE bank TO maxroach;

Exit the SQL shell:

  1. > \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.

  1. # Import the driver.
  2. import psycopg2
  4. # Connect to the "bank" database.
  5. conn = psycopg2.connect(
  6.     database='bank',
  7.     user='maxroach',
  8.     sslmode='disable',
  9.     port=26257,
  10.     host='localhost'
  11. )
  13. # Make each statement commit immediately.
  14. conn.set_session(autocommit=True)
  16. # Open a cursor to perform database operations.
  17. cur = conn.cursor()
  19. # Create the "accounts" table.
  20. cur.execute("CREATE TABLE IF NOT EXISTS accounts (id INT PRIMARY KEY, balance INT)")
  22. # Insert two rows into the "accounts" table.
  23. cur.execute("INSERT INTO accounts (id, balance) VALUES (1, 1000), (2, 250)")
  25. # Print out the balances.
  26. cur.execute("SELECT id, balance FROM accounts")
  27. rows = cur.fetchall()
  28. print('Initial balances:')
  29. for row in rows:
  30.     print([str(cell) for cell in row])
  32. # Close the database connection.
  33. cur.close()
  34. conn.close()

Then run the code:

  1. $ python basic-sample.py

The output should be:

  1. Initial balances:
  2. ['1', '1000']
  3. ['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.

Note:
CockroachDB may require the client to retry a transaction in case of read/write contention. CockroachDB provides a generic retry function that runs inside a transaction and retries it as needed. You can copy and paste the retry function from here into your code.

  1. # Import the driver.
  2. import psycopg2
  3. import psycopg2.errorcodes
  5. # Connect to the cluster.
  6. conn = psycopg2.connect(
  7.     database='bank',
  8.     user='maxroach',
  9.     sslmode='disable',
  10.     port=26257,
  11.     host='localhost'
  12. )
  14. def onestmt(conn, sql):
  15.     with conn.cursor() as cur:
  16.         cur.execute(sql)
  19. # Wrapper for a transaction.
  20. # This automatically re-calls "op" with the open transaction as an argument
  21. # as long as the database server asks for the transaction to be retried.
  22. def run_transaction(conn, op):
  23.     with conn:
  24.         onestmt(conn, "SAVEPOINT cockroach_restart")
  25.         while True:
  26.             try:
  27.                 # Attempt the work.
  28.                 op(conn)
  30.                 # If we reach this point, commit.
  31.                 onestmt(conn, "RELEASE SAVEPOINT cockroach_restart")
  32.                 break
  34.             except psycopg2.OperationalError as e:
  35.                 if e.pgcode != psycopg2.errorcodes.SERIALIZATION_FAILURE:
  36.                     # A non-retryable error; report this up the call stack.
  37.                     raise e
  38.                 # Signal the database that we'll retry.
  39.                 onestmt(conn, "ROLLBACK TO SAVEPOINT cockroach_restart")
  42. # The transaction we want to run.
  43. def transfer_funds(txn, frm, to, amount):
  44.     with txn.cursor() as cur:
  46.         # Check the current balance.
  47.         cur.execute("SELECT balance FROM accounts WHERE id = " + str(frm))
  48.         from_balance = cur.fetchone()[0]
  49.         if from_balance < amount:
  50.             raise "Insufficient funds"
  52.         # Perform the transfer.
  53.         cur.execute("UPDATE accounts SET balance = balance - %s WHERE id = %s",
  54.                     (amount, frm))
  55.         cur.execute("UPDATE accounts SET balance = balance + %s WHERE id = %s",
  56.                     (amount, to))
  59. # Execute the transaction.
  60. run_transaction(conn, lambda conn: transfer_funds(conn, 1, 2, 100))
  63. with conn:
  64.     with conn.cursor() as cur:
  65.         # Check account balances.
  66.         cur.execute("SELECT id, balance FROM accounts")
  67.         rows = cur.fetchall()
  68.         print('Balances after transfer:')
  69.         for row in rows:
  70.             print([str(cell) for cell in row])
  72. # Close communication with the database.
  73. conn.close()

Then run the code:

  1. $ python txn-sample.py

The output should be:

  1. Balances after transfer:
  2. ['1', '900']
  3. ['2', '350']

To verify that funds were transferred from one account to another, start the built-in SQL client:

  1. $ cockroach sql --insecure --database=bank

To check the account balances, issue the following statement:

  1. > SELECT id, balance FROM accounts;
  1. +----+---------+
  2. | id | balance |
  3. +----+---------+
  4. |  1 |     900 |
  5. |  2 |     350 |
  6. +----+---------+
  7. (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:

Was this page helpful?
YesNo