FlaskでUPSERT機能を実装する

こんにちは！
スマレジ テックファームのMichiです！

今回はFlaskでUPSERT機能を実装する方法を解説します。
UPSERTとは、保存対象のデータに対して、すでにDBに同じキーを持つデータが存在している場合はUPDATE、存在しない場合はINSERTする処理のことです。
Flaskでは標準でこのUPSERT機能を持っていないので、自分で実装する必要があります。（PostgreSQLの場合は標準機能で対応できるらしいが、MySQL/MariaDBは対応していないらしい...）

単独UPSERT

今回は例では、ユーザーが自分の勤怠情報を登録できるシステムを例にします。
user_idとtarget_dateカラムをキーにして、すでにDBに同じデータが存在する場合はUPDATE、存在しないはINSERTの処理を実行します。

def upsert(user_id, target_date, attendance_type):
    # テーブルにすでに同じuser_id, target_dateを持つデータが存在するかチェック
    attendance = Attendance.query.filter(
        Attendance.user_id == user_id,
        Attendance.target_date == target_date,
    ).first()

    # 存在する場合はattendance_typeを上書きする
    if attendance:
        attendance.attendance_type = attendance_type
    # 存在しない場合は新規のattendanceオブジェクトを作成して、session_addに加える
    else:
        attendance = Attendance(
            user_id=user_id,
            target_date=target_date,
            attendance_type=attendance_type,
        )
        db.session.add(attendance)

    # クエリの実行
    db.session.commit()
    return attendance

同じキーを持つデータが存在するか最初に問い合わせ、その結果次第でUPDATE OR INSERTの処理を分けます。 まあ、これは簡単。

一括UPSERT

def pre_upsert(user_id, target_date, attendance_type):
    # テーブルにすでに同じuser_id, target_dateを持つデータが存在するかチェック
    attendance = Attendance.query.filter(
        Attendance.user_id == user_id,
        Attendance.target_date == target_date,
    ).first()

    # 存在する場合はattendance_typeを上書きする
    if attendance:
        attendance.attendance_type = attendance_type
    # 存在しない場合は新規のattendanceオブジェクトを作成する
    else:
        attendance = Attendance(
            user_id=user_id,
            target_date=target_date,
            attendance_type=attendance_type,
        )

    return attendance


def bulk_upser(user_id, target_year, target_month, target_day, attendance_type,):
    # 一括UPSERTの対象となる日付を絞り込む
    last_day = monthrange(target_year, target_month)[1]
    date_list = [date(target_year, target_month, day) for day in range(1, last_day + 1)]
    target_dates = list(filter(lambda d: datetime.weekday(d) == target_day, date_list))

    # 対象日付の回数分pre_upsertをループして、結果を配列に格納する
    attendances = []
    for target_date in target_dates:
        attendances.append(
            pre_upsert(user_id, target_date, attendance_type)
        )

    # 一括インサートのクエリ発行（アップデート対象のオブジェクトは無視される）
    db.session.add_all(attendances)

    # クエリの実行
    commit()
    return attendances

次に複数の勤怠を一括でUPSERTしたいという場合です。
曜日ごとに一括で登録する機能を実装します。 ポイントは以下の2つです。

1. 単独UPSERTで使ったupsert関数はpre_upsertという名称に変更して、UPDATE or INSERT対象となるオブジェクトを返すだけの関数にする。
2. 実際のUPSERT処理はbulk_upsert関数の中で行う。db.session.add_all関数は、すでにプライマリーキーがテーブルに存在しているオブジェクトに対しては存在を無視して実行されるので、自動的にUPDAREとINSERTの処理を分けることができる。

まとめ

FlaskでUPSERT機能を実装する方法を解説しました。
一括UPSERTのUPDATE処理のほうはデータごとにクエリが発行されてしまい、大量のデータを更新するときは遅くなってしまうので注意です。
（一括UPDATEする方法はあるのか...？知っている人がいれば教えてください。）