How To Write A Simple Postgres JSON Wrapper With Python
By: Tate Galbraith
Tags: Postgres, Code, Programming, Python, Technology

Postgres is a legendary database. It has been around for a while and is well established with great support. This database is fast, efficient and used by a wide range of companies ranging from startups to international conglomerates. One reason Postgres is so popular is because there are a ton of extensions, plugins and connectors for it. You can connect via so many different programming languages and there are a number of top-notch ORMs available too.

Depending on your specific needs and project size, you might not need to setup a ton of extra abstraction layers or frameworks in order to handle queries into Postgres. Making a simple wrapper for different queries and commonly used functions is straightforward and effective in languages like Python.

If you’re prototyping a tiny application or working with a highly constrained embedded system then there can be great benefit in writing a light wrapper. Getting JSON results out of PG is super easy using the extremely popular psycopg2 module. This low-level database connection module for Python allows you to hook up to Postgres, execute SQL and perform other database functions. Let’s explore how to setup a database wrapper in Python using this module to make executing queries fast and easy.

The database

For our example Postgres server we’re going to use a pre-populated database full of test data. This gives us a bunch of sample data to query against. After all what fun is an empty database.

We’ll be using Docker to setup the database container so we don’t pollute our host operating system with test installations. If you’ve never used Docker before or need a quick refresher Docker — Beginner’s Guide — Part 1: Images & Containers is a fantastic guide by Sebastian Eschweiler to help you get up and running.

Check out postgres-dataset on Docker Hub (shout out to aa8y for creating this incredibly handy tool) to pull the latest image and deploy it using the commands below. You’ll want to change the first port from 5432 to something else if you’re already running a Postgres instance on your machine.

docker pull aa8y/postgres-dataset

docker run -d -p 5432:5432 --name pg-ds-fake aa8y/postgres-dataset:latest

Give the database a few minutes to populate with the test data and then go ahead and attach to the container using the following command:

docker exec -it pg-ds-fake /bin/bash

We will need to make one tweak to allow access from the host machine into the container’s PG instance. While connected to the container, run the following:

echo "host all all 0.0.0.0/0 trust" > /var/lib/postgresql/data/pg_hba.conf pg_ctl reload

The first line will allow all TCP/IP connections to Postgres and remove all other authentication. This is strictly for testing to make things easier and not deal with passwords and authentication from the host. Under no circumstances should you ever add this configuration to a production instance of Postgres. The last command reloads Postgres so that it picks up the new configuration.

Now you should be able to connect to Postgres using psql from the host machine by running (ensure you exit the container first):

psql -U postgres -h localhost

If you don’t have psql installed you can reference the guide available here to install. Try out a test query on the world database:

\c world

SELECT * FROM country LIMIT 5;

The wrapper

Now that we’ve got the database up and running with some test data to look at we can begin setting up our wrapper in Python. Let’s open up a new Python file and start editing. We can just call this wrapper.py for now and add the following:

#!/usr/bin/env python3 import simplejson as json

import psycopg2

from psycopg2.extras import RealDictCursor class PostgresWrapper:

def __init__(self, **args):

self.user = args.get('user', 'postgres')

self.port = args.get('port', 5432)

self.dbname = args.get('dbname', 'world')

self.host = args.get('host', 'localhost')

self.connection = None

First, we’re pulling in simplejson instead of the standard json library because when we query for Decimal values they will only be serialized correctly using a more updated library with support for dumping this datatype.

Next, we bring in the psycopg2 library and the RealDictCursor factory. This factory will let us easily serialize query results directly to a Dict type and then subsequently dump that to JSON. We will also define a basic PostgresWrapper class to hold all of our features and supply some default initialization values. Now let’s add some simple connection functions:

def connect(self):

pg_conn = psycopg2.connect(

user=self.user,

port=self.port,

dbname=self.dbname,

host=self.host

)

self.connection = pg_conn def get_json_cursor(self):

return self.connection.cursor(cursor_factory=RealDictCursor)

Inside the connect() function we’re using psycopg2 to connect directly to the Postgres database we setup previously. This connection also specifies which database to connect to inside Postgres. For this example we’ll be using the world data which contains things like cities and countries. The next function uses the previous connection to setup a new cursor. The cursor is similar to the interactive psql prompt that you type SQL statements into, except this one is non-interactive and can be programmatically controlled. This cursor also uses the RealDictCursor specification to return Dict types which are easily returned as a JSON string later.

Finally, let’s setup functions to actually fetch some data:

@staticmethod

def execute_and_fetch(cursor, query):

cursor.execute(query)

res = cursor.fetchall()

cursor.close()

return res def get_json_response(self, query):

cursor = self.get_json_cursor()

response = self.execute_and_fetch(cursor, query)

return json.dumps(response) def get_countries(self):

query = "SELECT * FROM country LIMIT 2;"

print(self.get_json_response(query))

The execute_and_fetch() function uses the cursor we setup previously to execute a SQL query, close the cursor connection and then return the results. It is important to close the cursor when the action is complete because otherwise each call to execute_and_fetch() will instantiate multiple cursors and consume more resources in the database. The get_json_response() function is a helper function that handles instantiating the cursor and passing both the cursor and query to execute_and_fetch() . Afterwards, this function returns the response as a JSON string.

The get_countries() function will leverage get_json_response() and pass in the query for countries. When we call this function we receive the final JSON blob of queried rows from the database.

Putting it all together

Let’s recap what the whole wrapper class should look like:

#!/usr/bin/env python3 import simplejson as json

import psycopg2

from psycopg2.extras import RealDictCursor class PostgresWrapper:

def __init__(self, **args):

self.user = args.get('user', 'postgres')

self.port = args.get('port', 5432)

self.dbname = args.get('dbname', 'world')

self.host = args.get('host', 'localhost')

self.connection = None def connect(self):

pg_conn = psycopg2.connect(

user=self.user,

port=self.port,

dbname=self.dbname,

host=self.host

)

self.connection = pg_conn def get_json_cursor(self):

return self.connection.cursor(cursor_factory=RealDictCursor) @staticmethod

def execute_and_fetch(cursor, query):

cursor.execute(query)

res = cursor.fetchall()

cursor.close()

return res def get_json_response(self, query):

cursor = self.get_json_cursor()

response = self.execute_and_fetch(cursor, query)

return json.dumps(response) def get_countries(self):

query = "SELECT * FROM country LIMIT 2;"

print(self.get_json_response(query)) dbconn = PostgresWrapper()

dbconn.connect()

dbconn.get_countries()

Near the bottom of the file, we handle instantiating the PostgresWrapper , connecting and then calling get_countries() to print some sample JSON data.

You should be able to simply run the file and produce the results:

./wrapper.py

Now you should see JSON output similar to this:

[{"code": "AFG", "name": "Afghanistan", "continent": "Asia", "region": "Southern and Central Asia", ...}]

If everything worked correctly, you’ll end up with a very simple wrapper class for Postgres that can be tweaked to output JSON blobs of different queries. This can be incorporated into an existing Python application or even used as a rudimentary backend for an API. Keep in mind this is a simple example class, you could extend this class quite a bit and make performing queries even more flexible.

Pairing Postgres with the power of Python and great modules you can simplify and abstract the complexity of connecting to and getting data out of databases. For an even more polished and elegant approach, consider exploring an ORM module called SQLAlchemy. This allows you to interact with database results as actual Python objects and write even less database-specific code.