Python3

デコレータでトランザクション処理を試しに実装してみた

More than 1 year has passed since last update.

OverView

Pythonのデコレータがまだイマイチ馴染みきれていないので、何か実際に書いてみようと思った。
DBのトランザクション処理がちょうど良さそうだと思ったので、試しに実装してみた。

DBの処理フロー

成功時 begin -> [何らかの処理] -> flush -> commit -> end
テスト時 begin -> [何らかの処理] -> flush -> rollback -> end
失敗時 begin -> [何らかのエラー] -> rollback -> end

環境

python: 3.6

構成

簡単に以下の2ファイル

- app.py
- utils.py

メイン関数

先にメインの関数であるmain_funcのソースから。デコレータtransactionは後述する。
main_funcはprintしているだけだが、本来はinsertしたりupdateしたりいろいろ考えられる。
今回は動作検証が目的なので、successの値に応じて例外を投げている。
この例ではデコレータの引数productionがハードコーディングされているが、実際はアプリケーション起動時とテスト起動時で読み分けることを想定。

from utils import transaction

@transaction(production=True)
def main_func(success):
    if success:
        print("DBで処理する")
    else:
        raise


if __name__ == '__main__':
    main_func(success=True)

デコレータ

def begin():
    print("開始しました")

def end():
    print("終了しました")

def flush():
    print("flushしました")

def commit():
    print("commitしました")

def rollback():
    print("rollbackしました")

def transaction(production):
    def decorator(f):
        def wrapper(*args, **kwargs):
            begin()
            try:
                f(*args, **kwargs)
                flush()

                if production:
                    commit()
                else:
                    rollback()
            except Exception:
                rollback()
                raise MyError("失敗しました")
            finally:
                end()
        return wrapper
    return decorator


class MyError(Exception):
    def __init__(self, message):
        self.message = message

実際にテスト書くときは、flushとrollbackの間にテストコードが入るので、setupの中にbegin、teardownの中にrollbackとendが入る。

まとめ

多少デコレータがわかった。
もっと良い方法など知ってる方いたら教えてください。