Connecting to a Postgresql Database with Python

2020-07-18 13:50:05 | #sysadmin

Tested On

  • Linux Ubuntu 20.04

This tutorial will teach you how to program a reusable Python 3 module for connecting to a Postgresql database.

While we could certainly leverage an ORM, such as SqlAlchemy, our goal here is to create a lightweight class for simple database connections and queries. This will come in handy for when we want to create a small microservice, with one function to perform, that we want to deploy to AWS Lambda.

Setting Up A Project Skeleton

Run the following commands on Linux to create the project skeleton and activate the virtual environment:

cd ~
mkdir python-postgresql-connections
cd python-postgresql-connections
virtualenv -p python3 venv
source venv/bin/activate
touch database.py
touch main.py

If you're on a different operating system, just make sure your skeleton looks like the following example and that you're running in the virtual environment.

▾ python-postgresql-connections/
  ▸ venv/
  database.py
  main.py

Full Code Example for the Postgresql Database Connection Utility

The following example connects to a Postgresql DB, in under 75 lines of Python 3 code.

Filename: database.py

import functools
import psycopg2
import psycopg2.extras  # optional
psycopg2.extras.register_uuid()  # optional


def require_connection(f):
    @functools.wraps(f)
    def wrapped(self, *args, **kwargs):
        if not self.connected():
            self.connect()
        return f(self, *args, **kwargs)
    return wrapped


class PsqlClient:

    def __init__(self, **kwargs):
        self.conn = None
        self.database = kwargs['database']
        self.user = kwargs['user']
        self.password = kwargs['password']
        self.host = kwargs['host']
        self.port = kwargs['port']

    def connected(self) -> bool:
        return self.conn and self.conn.closed == 0

    def connect(self):
        self.close()
        self.conn = psycopg2.connect(
            database=self.database,
            user=self.user,
            password=self.password,
            host=self.host,
            port=int(self.port)
        )

    @require_connection
    def execute(self, sql, data=None):
        cur = self.conn.cursor()

        try:
            cur.execute(sql, data)
            self.conn.commit()
            return cur
        except Exception:
            self.conn.rollback()
            raise

    @require_connection
    def rollback(self):
        self.conn.rollback()

    def close(self):
        if self.connected():
            self.conn.close()

        self.conn = None

Explanation Of The database.py Code

Lines 1-4: imports the required modules from psycopg2, a Postgresql database adapter. Lines 3 and 4 are responsible for UUID support and can be removed if you aren't storing UUIDs in your database. We also import functools to make it easier to define decorators.

The require_connection decorator defined in lines 7-13 ensures a connection is made prior to any queries.

And finally, we have our PsqlClient class, in lines 16-59, which instantiates with database credentials, and defines functions to check for connections, to make connections, execute queries, roll back in case a query fails, and close the database connection.

Full Code Example for the main.py File

This is the file we'll use to create a reusable instance of the PsqlClient.

Filename: main.py

import os
from database import PsqlClient

db = PsqlClient(
    database=os.environ['POSTGRES_DB'],
    user=os.environ['POSTGRES_USER'],
    password=os.environ['POSTGRES_PASSWORD'],
    host=os.environ['POSTGRES_HOST'],
    port=int(os.environ['POSTGRES_PORT'])
)

if __name__ == '__main__':
  try:
      cur = db.execute("SELECT * FROM table")
      result = cur.fetchall()
      if result.rowcount:
        print(result)
      else:
        print('No records found')
  except Exception as e:
      print(e)

How to Run the Program

Make sure you are still inside your virtual machine by running source venv/bin/activate from inside the project root folder. Then run python3 main.py to execute the program.

Explanation of the main.py Code

Lines 1-2: imports the required dependencies, including the PsqlClient class from our database.py module.

Lines 4-10: instantiates the client instance, passing in the credentials defined in environment variables. We recommend environment variables because they keep sensitive credentials out of the code. If you're having trouble connecting because you're not familiar with environment variables, for testing purposes, you can substitute os.environ['POSTGRES_*****'] with whatever hardcoded values. But make sure not to commit or deploy the code with the hardcoded values still embedded.

Lines 12-21: executes a SELECT query to fetch records from the database and outputs errors, if any. Note: you must substitute table with the name of a table you've already defined. A table and its columns must be defined, and records must exist before you can fetch said records. To learn more about table creation, please read this article.

If you'd like an example of how to insert a record using the PsqlClient class, please refer to the following. Note: You will have to substitute your own column names and values. The main purpose of this code sample is to demonstrate how to safely pass values into the query while preventing SQL injection.

try:
    db.execute('INSERT INTO table \
        (column1, column2, column3) \
        VALUES (%(column1)s, %(column2)s, %(column3)s)', (
        {
            'column1': 'value1',
            'column2': 'value2',
            'column3': 'value3'
        })
    )
except Exception as e:
    print(e)

If you'd like an example of how to update a record using the PsqlClient class, please refer to the following:

try:
    cur = db.execute("UPDATE users SET column1 = %(column1)s, column2 = %(column2)s WHERE id = %(id)s", {'column1': 'new value1', 'column2': 'new value2', 'id': 1})
    if not cur.rowcount:
      print('An error occurred while updating')
except Exception as e:
    print(e)

Comments

You must log in to comment. Don't have an account? Sign up for free.

Subscribe to Our Newsletter

Would you like to receive free whitepapers and other IT news? Just leave your email address below. You may opt out at any time.



Tell Us About Your Project









Contact Us

Do you have a specific IT problem that needs solving or just have a general IT question? Use the contact form to get in touch with us and an IT professional will be with you, momentarily.

Hire Us

We offer web development, enterprise software development, QA & testing, google analytics, domains and hosting, databases, security, IT consulting, and other IT-related services.

Free IT Tutorials

Head over to our tutorials section to learn all about working with various IT solutions.

Contact