Execute an SQL statement for PostgreSQL using Python

Introduction to SQL statements and psycopg2

The psycopg2 Python adapter for PostgreSQL is a lightweight distribution that lets you to execute SQL statements, using cursors, and make concurrent, multi-threaded programs with Python scripts.

Execute an SQL statement for PostgreSQL using Python

This article will demonstrate how you can execute an SQL statement for PostgreSQL using a Python script that can accept command-line arguments for concatenating SQL statement strings. The sql.SQL object can then be passed to a psycopg2 cursor object’s execute() method. The code in this article uses Python’s sys.argv list to check for any arguments passed to that it can concatenate SQL statement strings using them.

Prerequisites for using the psycopg2 PostgreSQL Python adapter

  • Make sure PIP3 is installed so that you can install Python 3 package. Use the pip3 -V command to verify that it’s installed by getting its version number.

  • Install the psycopg2 distribution for Python 3 using its PIP package manager with the following command: pip3 install psycopg2.

  • You can verify that PostgreSQL is installed on the machine or server by having it return the version number for the psql interface with this command: psql -V.

Setup the PostgreSQL database and grant a user permissions

Make sure that the PostgreSQL server is running, and that you have access to one of its databases. You can enter into the psql command-line interface with the following command:

sudo su - postgres

Once inside of the psql interface you can create a user, database, and grant privileges using the following SQL statements:

CREATE USER some_user WITH PASSWORD 'somePass';
CREATE DATABASE "some_database";
GRANT ALL PRIVILEGES ON ALL TABLES IN SCHEMA public TO some_user;

Create a Python script for the PostgreSQL application

The last step in setting up the application is to create a new Python script. Use the touch command in a UNIX terminal to create the script in your project directory:

touch execute.py

NOTE: The code in this article has been designed for Python 3 and has not been tested with Python 2.7

Import the Python libraries needed to execute SQL statements in PostgreSQL

Python libraries and modules have to imported into a script (using the import statement) before they can be used. The following code will import the necessary Python libraries to connect to PostgreSQL and execute SQL statements:

# import the psycopg2 database adapter for PostgreSQL
from psycopg2 import connect, sql

# import the sys library for arguments
import sys

NOTE: The sys library will get the system’s command-line arguments that have been passed to the Python script using a terminal or command prompt window.

Connect to PostgreSQL using the psycopg2 Python adapter

Declare a connection object instance of the psycopg2 adapter library inside of a try-except indentation block with the following code:

# instantiate a cursor object from the connection
try:

    # declare a new PostgreSQL connection object
    conn = connect (
        dbname = "some_database",
        user = "objectrocket",
        host = "localhost",
        password = "mypass"
    )

NOTE: Make sure that the database and user specified in the method call exist, and that the use has the appropriate permissions for the SQL statements, or else Python will return an exception.

Instantiate a psycopg2 cursor object using the PostgreSQL connection

The following code will attempt to declare a cursor instance for psycopg2 using the cursor() method call:

    # attempt to create a cursor object
    cursor = conn.cursor()

except Exception as err:
   
    # set the cursor to None if exception
    cursor = None
    print ("\npsycopg2 error:", err)

If there’s an error while creating the connection or cursor objects, then the code above will set the cursor’s value to None, and print the exception response to the terminal.

Check the PostgreSQL connection

If the cursor is not set to None, and the connection was successful, then Python will proceed to parse the arguments passed to it. This next bit of code evaluates the cursor object, and the length of the argument list, before making any SQL statements:

# check if the connection was valid
if cursor != None:
    print ("\nconnection successful:", conn, "\n")

    # get the table name from args
    if len(sys.argv) <= 1:
        print ("Please pass a SQL string as an argument")
        quit()

NOTE: As you can see in the code above, Python will print a message and quit() the script if the user failed to pass an argument at the time of script execution.

The ‘sys.argv’ list of command-line arguments in Python

The sys.argv string list in Python stores the script’s name as the first element ([0]), and you can pass an argument (with spaces in it) if you enclose it in quotation marks. The following example code is what the list would look like if the user passed two arguments:

# user executes: python3 script_name.py param_1 "param 2"
sys.argv == ["script_name.py", "param_1", "param 2"]

Screenshot of IDLE for Python 3 accessing the sys argv list of string arguments

Parse the sys arguments passed to the Python script

Python will parse the arguments if the length of sys.argv list has something passed to it. The following code will get all of the table’s records if only one argument is passed, but will get just the records from a specific column in the table if two arguments are passed when the script is run:

    # if the sys.argv list has more than one element
    else:

        # first sys arg is always the script name
        if len(sys.argv) == 2:
            # join the sys arguments
            table_name = "".join(sys.argv[1])

            # default is all cols if no arg is passed
            cols = "*"
        else:
            # table name will be second param
            table_name = "".join(sys.argv[1])

            # column name is third param
            cols = "".join(sys.argv[2:])

NOTE: The default parameter for the columns will be the * wildcard if the user only passes the table name. This will allow for all the records, in all of the columns, to be stored in the cursor.

Concatenate a SQL string for PostgreSQL using the Python arguments

Putting curly braces ({}) inside of a string, while calling the string’s format() method, allows you to pass string values into a string. The following code inserts the columns and table names into the SQL string while concatenating it:

        # concatenate a SQL string from table name and columns
        sql_statement = "SELECT {} FROM {};".format(
            cols,
            table_name
        )
        print ("sql_statement:", sql_statement)

Declare a sql.SQL() psycopg2 object from the SQL statement string

This next code creates a sql.SQL object from the SQL statement string using the psycopg2 library’s sql.SQL() method call:

        try:
            # have sql.SQL() return a sql.SQL object
            sql_object = sql.SQL(
                # pass SQL string to sql.SQL()
                sql_statement
            ).format(
                # pass the identifier to the Identifier() method
                sql.Identifier( table_name )
            )

NOTE: It’s recommended that you pass a sql.SQL object to the cursor, instead of just concatenating a string, in order to help prevent SQL injection attacks.

Execute the sql.SQL() object and store the PostgreSQL records

The code is now ready to call the cursor object’s execute() method. The following code will pass the sql.SQL object to the method call in order to store the records in the cursor:

            # pass the psycopg2.sql.SQL object to execute() method
            cursor.execute( sql_object )

            # print message if no exceptions were raised
            print ("cursor.execute() FINISHED")

Fetch all of the PostgreSQL records from the psycopg2 cursor

The last step is to call the cursor’s fetchall() method in order to have it return a list of tuples containing the records. Here’s how you can do that and enumerate() over the list of records:

            # use the fetchall() method to return a list of all the data
            print ("cursor.fetchall():")
            table_data = cursor.fetchall()

            # enumerate over the list of tuple rows
            for num, row in enumerate(table_data):
                print ("row:", row)
                print (type(row))
                print ("\n")

Rollback the PostgreSQL transaction if there’s an exception

This script only execute one SQL string every time you run it, but, if you decide to modify the code so that more than one are executed, you should call the connection object’s rollback() method anytime an exception is returned before another transaction can be executed:

        except Exception as err:
            print ("cursor.execute() ERROR:", err)

            # rollback the statement
            conn.rollback()

Commit the PostgreSQL transactions

If you modify this example code, so that it’s designed to make changes to a PostgreSQL table instead of just returning records, then you’ll have to call the connection object’s commit() method to commit any SQL transaction to the table data:

        # commit transactions to PostgreSQL
        conn.commit()

Close the PostgreSQL cursor and connection objects

The last bit of code in the script will close the psycopg2 cursor and connection objects in order to prevent memory leaks:

        # close the cursor object to avoid memory leaks
        cursor.close()

        # close the connection object
        conn.close()

Pass SQL parameters to the Python script when you run it

Save the code in your Python script, and then you can execute it using the python3 command in a terminal or command prompt window. Here’s an example of how you can run the script and have it get table row data from the some_table PostgreSQL table:

python3 execute.py some_table

Screenshot of passing an argument to Python when running the psycopg2 script

Or you can have it get a specific column of data by passing a second argument:

python3 execute.py some_table some_col

Screenshot of passing two arguments to Python when running psycopg2 script

If the table doesn’t exist then the Python script will return a response that looks something like this:

sql_statement: SELECT * FROM not_there;
cursor.execute() ERROR: relation "not_there" does not exist
LINE 1: SELECT * FROM not_there;

Conclusion

The code in this article shows how you can easily execute SQL statements in Python using the psycopg2 adapter by connecting to PostgreSQL in the Python script and concatenating SQL strings. Check out the complete code for the script below.

Just the Code

#!/usr/bin/python3
# -*- coding: utf-8 -*-

# import the psycopg2 database adapter for PostgreSQL
from psycopg2 import connect, sql

# import the sys library for arguments
import sys

# instantiate a cursor object from the connection
try:

    # declare a new PostgreSQL connection object
    conn = connect (
        dbname = "some_database",
        user = "objectrocket",
        host = "localhost",
        password = "mypass"
    )

    # attempt to create a cursor object
    cursor = conn.cursor()

except Exception as err:
   
    # set the cursor to None if exception
    cursor = None
    print ("\npsycopg2 error:", err)

# check if the connection was valid
if cursor != None:
    print ("\nconnection successful:", conn, "\n")

    # get the table name from args
    if len(sys.argv) <= 1:
        print ("Please pass a SQL string as an argument")
        quit()
    else:

        # first sys arg is always the script name
        if len(sys.argv) == 2:
            # join the sys arguments ()
            table_name = "".join(sys.argv[1])

            # default is all cols if no arg is passed
            cols = "*"
        else:
            # table name will be second param
            table_name = "".join(sys.argv[1])

            # column name is third param
            cols = "".join(sys.argv[2:])

        # concatenate a SQL string from table name and columns
        sql_statement = "SELECT {} FROM {};".format(
            cols,
            table_name
        )
        print ("sql_statement:", sql_statement)

        try:
            # have sql.SQL() return a sql.SQL object
            sql_object = sql.SQL(
                # pass SQL string to sql.SQL()
                sql_statement
            ).format(
                # pass the identifier to the Identifier() method
                sql.Identifier( table_name )
            )

            # pass the psycopg2.sql.SQL object to execute() method
            cursor.execute( sql_object )

            # print message if no exceptions were raised
            print ("cursor.execute() FINISHED")

            # use the fetchall() method to return a list of all the data
            print ("cursor.fetchall():")
            table_data = cursor.fetchall()

            # enumerate over the list of tuple rows
            for num, row in enumerate(table_data):
                print ("row:", row)
                print (type(row))
                print ("\n")

        except Exception as err:
            print ("cursor.execute() ERROR:", err)

            # rollback the statement
            conn.rollback()

        # commit transactions to PostgreSQL
        conn.commit()

        # close the cursor object to avoid memory leaks
        cursor.close()

        # close the connection object
        conn.close()

Pilot the ObjectRocket Platform Free!

Try Fully-Managed CockroachDB, Elasticsearch, MongoDB, PostgreSQL (Beta) or Redis.

Get Started

Keep in the know!

Subscribe to our emails and we’ll let you know what’s going on at ObjectRocket. We hate spam and make it easy to unsubscribe.