【python】【sls】【lambda】デコレータで入力チェックする

pythonのデコレータ使ってますか？？

 

現場で、serverless frameworkでデプロイしたAPI Gateway ＋ lambdaなメソッドで

HTTPヘッダに特定の値が設定されているか？をチェックする必要が出てきました。

そこで今回はそのチェックをデコレータでやる方法を採用してみました。

その内容をメモメモ。

 

 

 

 

デコレータ本体

デコレータそのものはこんな感じに書く。

import functools


def header_check(arg_hoge):
    """ヘッダーチェックデコレータ"""

    def wrapper(org_function):
        @functools.wraps(org_function)
        def decorate_method(*args, **kwargs):

            # ヘッダーを取得
            event = args[0]
            header = event["headers"]

            # ヘッダーからチェック対象を取得
            hoge = header.get("hoge")

            # チェック対象が設定されていない、または指定された値ではない場合エラーとする
            error_message = (
                "header error [required value not found or invalid value]"
            )
            if not hoge or hoge!= arg_hoge:
                return {
                    "statusCode": "500",
                    "body": error_message,
                }

            return org_function(*args, **kwargs)

        return decorate_method

    return wrapper

 

classで書く方法もあるみたいですが、

個人的にはメソッドで定義しちゃうほうが、その関数名がそのまま

デコレータとして使えるのでよいかなーと思います。

 

wrapperで呼び元の関数を引数に受けて、

functool.wrapsデコレータで実際に呼び元関数を包んでしまうんですね。

 

実行イメージはこんな感じ。

org_function = decorate_method(org_function)

 

呼び元の引数の値は*argsに格納されているので、そこから取得可能。

また、デコレータ自体に渡された引数は、header_checkの引数をそのまま使います。

 

個人的にちょっとハマったのは、

org_functionを実行するとき、呼び元が何かをreturnする場合には、

return org_functionとしてあげないといけなかったってところです。

 

 

色々なサンプル記事を見ましたが、どれもvoidなメソッドに

デコレートしててうまく動作しない原因を探すのに時間がかかってしまいました。

 

 

 

 

 

 

デコレータを使う

使い方はこんな感じ。

from typing import Dict
from decorators import header_check


@header_check(org_hoge="1")
def main(event: Dict, context):
    """lambdaの処理"""

 

 

こうするとorg_hogeの値とヘッダーに格納されたhogeの値が

一致していなければ、HTTP 500でエラーとして返却され、

正しければ、mainの関数が処理されるということになります。

 

 

まとめ

デコレータ自体は結構簡単に作れたので、

事前処理とか、共通化できそうなところがあれば積極的に作っていきたいなー。

 

 

それでは！！

 

 

にほんブログ村

人気ブログランキング