Dry-run sql query using psycopg2 on Redshift in Python

Introduction

I implemented the process to execute SQL query in Redshift using psycopg2 in Python, but I implemented it in a wrapper class to execute it with dry-run, so I will introduce how it was implemented I think that I want to do it.

Wrapper class implementation

Wrap psycopg2 to implement dry-run functionality. When executing a query, do not commit, but rollback to perform dry-run.

wrapper.py

import os
import psycopg2
import re

class DBWrapperClass:

    __is_dryrun = False #The default is dry-run do not run

    def __init__(self):
        self.__connection()

    def __connection(self):
        #DB connection information Please change accordingly
        dsn = {
            "host" : os.environ.get('DB_HOST'),
            "port" : os.environ.get('DB_PORT'),
            "database" : os.environ.get('DB_NAME'),
            "user" : os.environ.get('DB_USER'),
            "password" : os.environ.get('DB_PASS'),
        }

        self.__conn = psycopg2.connect(**dsn)

    #Allows you to specify arguments with with syntax
    def __call__(self, **kwargs):
        self.__is_dryrun = bool(kwargs['dryrun']) if 'dryrun' in kwargs else False
        return self

    #Get cursor with with syntax
    def __enter__(self):
        self.__cursor = self.__conn.cursor()
        return self

    #Commit or rollback at the end of the with syntax and close the cursor
    def __exit__(self, exc_type, exc_value, traceback):
        #Roll back if an exception occurs
        needs_rollback = True if exc_value else False

        self.__commit(needs_rollback)
        self.__close_cursor()

        #If you do not call it like a function with with syntax, the previous value will be inherited, so initialize it.
        self.__is_dryrun = False

    def __commit(self, needs_rollback=False):
        if self.__conn:
            # dry-When you need a run or rollback
            if self.__is_dryrun or needs_rollback:
                self.__conn.rollback()
            else:
                self.__conn.commit()

    #close cursor
    def __close_cursor(self):
        if self.__conn:
            self.__cursor.close()


    #Close the connection
    def close(self):
        if self.__conn:
            self.__conn.close()

    #Execute the query
    def execute(self, sql):
        # dry-For run, truncate is not executed
        #Since truncate cannot roll back, this control is included.
        if not (self.__is_dryrun and re.match(r'^\s*truncate\s+', sql, flags=re.IGNORECASE)):
            self.__cursor.execute(sql)

In order to be able to manage cursor with with syntax like psycopg2, processing related to cursor is implemented in each method of __enter__ and __exit__. We've also added a __call__ method so that you can set the dry-run specification in the with syntax.

Run dry-run in wrapper class

import wrapper

try:
    dbwrapper = wrapper.DBWrapperClass()
    with dbwrapper(dryrun=True) as db:
        #dry sql-Run with run
        db.execute('dry-sql you want to run with run')

finally:
    if 'dbwrapper' in locals():
        dbwrapper.close()

On the execution side, this is all you need to do. This time, only the argument for dry-run is set to make it as simple as possible, but it will be more convenient if other setting values ​​such as autocommit can also be set with the argument.

Summary

I implemented the wrapper class of psycopg2 and added the function of dry-run. I usually write programs in php, so I thought the with syntax of python was quite convenient. Also, what I thought about implementing this time is that python has quite a lot of magic methods and I definitely want to use it when implementing classes for different purposes.