PostgreSQL Python Tutorial: Transaction

Summary: in this tutorial, you will learn how to handle PostgreSQL transactions in Python using psycopg database adapter.

Introduction to the transaction in psycopg

In psycopg, the connection class is responsible for handling transactions. When you issue the first SQL statement to the PostgreSQL database using a cursor object, psycopg creates a new transaction.

From that moment, psycopg executes all the subsequent statements in the same transaction. If any statement fails, psycopg will abort the transaction.

The connection class has two methods for ending a transaction: commit() and rollback(). If you want to commit all changes to the PostgreSQL database permanently, you call the commit() method. And in case you want to cancel the changes, you call the rollback() method. Closing the connection object or destroying it using the del will also result in an implicit rollback.

It is important to notice that a simple SELECT statement will start a transaction that may result in undesirable effects such as table bloat and locks. Therefore, if you are developing a long-living application, you should call the commit() or rollback() method before leaving the connection unused for a long time.

Alternatively, you can set the autocommit attribute of the connection object to True. This ensures that psycopg executes every statement and commits it immediately.

The autocommit mode is also useful when you execute statements required to execute outside a transaction such as CREATE DATABASE and VACUUM.

The following shows a typical pattern for handling a transaction in psycopg:

#!/usr/bin/python

import psycopg2

conn = None
try:
    conn = psycopg2.connect(dsn)
    cur = conn.cursor()
    # execute 1st statement
    cur.execute(statement_1)
    # execute 2nd statement
    cur.execute(statement_1)
    # commit the transaction
    conn.commit()
    # close the database communication
    cur.close()
except psycopg2.DatabaseError as error:
    print(error)
finally:
    if conn is not None:
        conn.close()

Postgres Python transaction example

We will use the parts and vendor_parts tables that we created in the creating table tutorial for the demonstration.

Suppose you need to add a new part and assign the vendors who supply the part at the same time. To do this, first, you insert a new row into the parts table and get the part id. Then, you insert rows into the vendor_parts table. The following add_part() function demonstrates the idea:

#!/usr/bin/python
import psycopg2
from config import config

def add_part(part_name, vendor_list):
    # statement for inserting a new row into the parts table
    insert_part = "INSERT INTO parts(part_name) VALUES(%s) RETURNING part_id;"
    # statement for inserting a new row into the vendor_parts table
    assign_vendor = "INSERT INTO vendor_parts(vendor_id,part_id) VALUES(%s,%s)"

    conn = None
    try:
        params = config()
        conn = psycopg2.connect(**params)
        cur = conn.cursor()
        # insert a new part
        cur.execute(insert_part, (part_name,))
        # get the part id
        part_id = cur.fetchone()[0]
        # assign parts provided by vendors
        for vendor_id in vendor_list:
            cur.execute(assign_vendor, (vendor_id, part_id))

        # commit changes
        conn.commit()
    except (Exception, psycopg2.DatabaseError) as error:
        print(error)
    finally:
        if conn is not None:
            conn.close()

To test the add_part() function, we call it to insert some parts and assign them to the respective vendors as follows:

if __name__ == '__main__':
    add_part('SIM Tray', (1, 2))
    add_part('Speaker', (3, 4))
    add_part('Vibrator', (5, 6))
    add_part('Antenna', (6, 7))
    add_part('Home Button', (1, 5))
    add_part('LTE Modem', (1, 5))

Let’s query the parts and vendor_parts table to confirm the transaction.

suppliers=# select * from parts;
 part_id |  part_name
---------+-------------
       1 | SIM Tray
       2 | Speaker
       3 | Vibrator
       4 | Antenna
       5 | Home Button
       6 | LTE Modem
(6 rows)

suppliers=# select * from vendor_parts;
 vendor_id | part_id
-----------+---------
         1 |       1
         2 |       1
         3 |       2
         4 |       2
         5 |       3
         6 |       3
         6 |       4
         7 |       4
         1 |       5
         5 |       5
         1 |       6
         5 |       6
(12 rows)

As you see, we have successfully inserted the data into both parts and vendor_parts tables.

Let’s insert another part, but this time, we use an invalid vendor id purposefully for the demonstration purpose. The program should not add a new part without assigning it to a vendor.

# no rows inserted into the parts and vendor_parts tables
add_part('Power Amplifier', (99,))

An exception occurred.

insert or update on table "vendor_parts" violates foreign key constraint "vendor_parts_vendor_id_fkey"
DETAIL:  Key (vendor_id)=(99) is not present in table "vendors".

We query data from the parts and vendor_parts tables again, there is no new data, therefore, the function works as expected.

Transactions using the with statement

Starting from psycopg 2.5, the connection and cursor are Context Managers and therefore you can use them with the with statement:

with psycopg2.connect(dsn) as conn:
    with conn.cursor() as cur:
        cur.execute(sql)

Psycopg commits the transaction if no exception occurs within the with block, and otherwise it rolls back the transaction.

Unlike other context manager objects, exiting the with block does not close the connection but only terminates the transaction. As a result, you can use the same connection object in the subsequent with statement in another transaction as follows:

conn = psycopg2.connect(dsn)

# transaction 1
with conn:
    with conn.cursor() as cur:
        cur.execute(sql)

# transaction 2
with conn:
    with conn.cursor() as cur:
        cur.execute(sql)

conn.close()

In this tutorial, you learned how to use the psycopg transaction and how to use transactions to insert data in the PostgreSQL database.