Connecting to a PostgreSQL database with Python

30 November 2016 Estimated reading time: 2 minutes Python

I have used psycopg2 to connect the database and a little trick to switch the schema from development to production and viceversa. I have used an ini file to store informations about the database. You can store any other informations in the file and keeping them tidy using the ini sections. Let's start with my config.ini file:

[POSTGRESQL]host = localhostdbname = mydbusername = postgrespassword = postgresschema = dev

The configurations.py file. This file reads the informations from the ini file using ConfigParser and contains a method that returns the parser:

import osimport ConfigParserclass Configurations: def __init__(self): self.config = ConfigParser.ConfigParser() self.config.read( os.path.join(os.path.dirname(__file__),'config.ini') ) def getConfigParser(self): return self.config

With the database.py file, I create the connection to the database and some methods to fetch data and be able to use field names instead of tuples:

import psycopg2import psycopg2.extrasfrom configurations import Configurationsclass Database:def __init__(self):configurations = Configurations()self.configParser = configurations.getConfigParser()schema = self.configParser.get('POSTGRESQL', 'schema')self.conn = psycopg2.connect("host='{}' dbname='{}' user='{}' password='{}'". format( self.configParser.get('POSTGRESQL', 'host'), self.configParser.get('POSTGRESQL', 'dbname'), self.configParser.get('POSTGRESQL', 'username'), self.configParser.get('POSTGRESQL', 'password') ))cur = self.conn.cursor()if schema is not None and schema != '':cur.execute("SET search_path TO {}". format(schema) )def getDbConnection(self):return self.conndef fetchAll(self, query):cur = self._executeQuery(query)rows = cur.fetchall()return rowsdef fetchOne(self, query):cur = self._executeQuery(query)rows = cur.fetchone()return rows# Execute the query and return the cursor objectdef _executeQuery(self, query):cur = self.conn.cursor(cursor_factory = psycopg2.extras.DictCursor)cur.execute(query)return cur

Now you are free to create a getter class for a table of your database and access the field names of the database. I have created this simple example: I suppose we have a book table. The Getter class can be like this:

class BookGetter:def __init__(self, database):self.db = databasedef getBooks(self):query = "SELECT id, title, author FROM books"return self.db.fetchAll(query)def getBook(self, id):query = "SELECT id, title, author FROM books WHERE id = {}; ". format(id)return self.db.fetchOne(query)

Let's see how we can test and use the objects:

# -*- coding: utf-8 -*-from database import Databasefrom BookGetter import BookGetterdb = Database()bookGetter = BookGetter(db)bookList = bookGetter.getBooks()for book in bookList:print book['id'], '', book['title']

The next steps

I have built this examples quickly and I have no unit test. I know this is bad and I will update this post with unit test and a mock for the database class. We can also transform the database class into an abstract object and extend it when we write the other getter classes. I hope this post will be useful. Happy coding!