PythonでRedshiftにpsycopg2を使ってsqlクエリをdry-run実行

はじめに

Pythonでpsycopg2を使ってRedshiftにSQLクエリを実行する処理を実装していたのですが、dry-runで実行させる為にラッパークラスで実装したので、どのように実装したのか紹介したいと思います。

ラッパークラスの実装

psycopg2をラッパーしてdry-run機能を実装します。
クエリを実行する時にcommitはせず、rollbackさせることでdry-runを行います。

wrapper.py

import os
import psycopg2
import re

class DBWrapperClass:

    __is_dryrun = False # デフォルトはdry-run実行しません

    def __init__(self):
        self.__connection()

    def __connection(self):
        # DBの接続情報です適宜変更してください
        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)

    # with構文で引数を指定できるようにします
    def __call__(self, **kwargs):
        self.__is_dryrun = bool(kwargs['dryrun']) if 'dryrun' in kwargs else False
        return self

    # with構文でcursorを取得します
    def __enter__(self):
        self.__cursor = self.__conn.cursor()
        return self

    # with構文の終了時にcommitやrollbackを行い、cursorを閉じます
    def __exit__(self, exc_type, exc_value, traceback):
        # 例外が発生した場合はrollbackさせます
        needs_rollback = True if exc_value else False

        self.__commit(needs_rollback)
        self.__close_cursor()

        # with構文で関数のように呼び出さない場合、前回の値が引き継がれるので初期化します
        self.__is_dryrun = False

    def __commit(self, needs_rollback=False):
        if self.__conn:
            # dry-runかrollbackが必要な時
            if self.__is_dryrun or needs_rollback:
                self.__conn.rollback()
            else:
                self.__conn.commit()

    # cursorを閉じます
    def __close_cursor(self):
        if self.__conn:
            self.__cursor.close()


    # 接続を閉じます
    def close(self):
        if self.__conn:
            self.__conn.close()

    # クエリを実行します
    def execute(self, sql):
        # dry-runの場合、truncateは実行しません
        # truncateはrollbackできない為、この制御を入れています
        if not (self.__is_dryrun and re.match(r'^\s*truncate\s+', sql, flags=re.IGNORECASE)):
            self.__cursor.execute(sql)

psycopg2と同じようにwith構文でcursorの管理をできるようにする為、__enter__と__exit__の各メソッドにcursorに関する処理を実装しています。
また、with構文でdry-runの指定を設定できるように__call__メソッドを追加しています。

ラッパークラスでdry-runを実行

import wrapper

try:
    dbwrapper = wrapper.DBWrapperClass()
    with dbwrapper(dryrun=True) as db:
        # sqlをdry-runで実行
        db.execute('dry-runで実行したいsql')

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

実行側ではこれだけで行えます。
今回はできるだけシンプルにする為にdry-run用の引数しか設定していませんが、autocommit等の他の設定値も引数で設定できるようにするとより便利になります。

まとめ

psycopg2のラッパークラスを実装してdry-runの機能を追加してみました。
普段はphpでプログラムを組むことが多いのでpythonのwith構文がかなり便利だと思いました。
また、今回実装してみて思ったのがpythonにはマジックメソッドがかなり多く違う用途のクラスを実装する際はぜひ使ってみたいです。