Django ORMのサブクエリを使う

この記事はiRidge Advent Calendar 2019の20日目の記事です。

今年担当した開発の中で、DjangoORMのSubqueryをはじめて使いました。やりたかったこと自体は非常にシンプルですが、DjangoのORMで実現するのには情報が少なくそこそこ苦戦したので、記事にしたためておきたいと思います。

やろうとしたこと

以下はモデルの例です。説明のために実際のモデル定義とは異なり、かなり単純化したものになっています。

class Shop(models.Model):
    identifier = models.CharField(max_length=255, unique=True)
    name = models.CharField(max_length=255)
    created_at = models.DateTimeField(auto_now_add=True)
    updated_at = models.DateTimeField(auto_now=True)
    is_deleted = models.BooleanField(default=False)

class Book(models.Model):
    shop = models.ForeignKey(Shop)
    uuid = models.CharField(max_length=36, unique=True, db_index=True,
                            blank=True)
    name = models.CharField(max_length=255)
    description = models.TextField(blank=True)
    created_at = models.DateTimeField(auto_now_add=True)
    updated_at = models.DateTimeField(auto_now=True)
    is_deleted = models.BooleanField(default=False, db_index=True)

class User(models.Model):
    shop = models.ForeignKey(Shop)
    device_user_id = models.CharField(max_length=255, db_index=True)
    created_at = models.DateTimeField(auto_now_add=True)
    updated_at = models.DateTimeField(auto_now=True)
    is_deleted = models.BooleanField(default=False)

class FavoriteBook(models.Model):
    book = models.ForeignKey(Book, db_index=True)
    user = models.ForeignKey(User, db_index=True)
    created_at = models.DateTimeField(auto_now_add=True)

既存のコードでは、UserとShopで絞り込んだBookの一覧を取得するquerysetが書かれていました。このquerysetに、あるユーザがBookをお気に入りしているか否かをFavoriteBookから情報を引っ張ってきて付与しようとしました。

実装

色々試行錯誤したのですが経緯は割愛します。
https://docs.djangoproject.com/en/3.0/ref/models/expressions/#subquery-expressions
を参考に以下のような2行のコードで実現できました。

queryset = ... # 既存のqueryset、UserとShopの条件で絞り込んだBookの一覧

# 各Bookに、あるユーザのお気に入り有無の情報を付加する
queryset_fav = FavoriteBook.objects.filter(user=user).filter(book=OuterRef('pk'))
queryset = queryset.annotate(favorite_book_id=Subquery(queryset_fav.values('id')[:1]))

OuterRef

サブクエリが外部のクエリのカラムを参照する時に使用します。
この場合、filter(book=OuterRef('pk'))でサブクエリのbook_idが元のquery_setのbook.idを参照するようになります。

SubQuery

作成しておいたqueryset_favをサブクエリとして使用します。
この場合、favorite_book_idのカラムが追加され、サブクエリの結果がセットされます。
より具体的に言えば、FaboriteBookに対象のUserとBookのレコードがある（＝対象ユーザにお気に入りされている）Bookにはfavorite_book.idが、そうでない（＝お気に入りされていない）Bookにはnullがセットされます。

このとき発行されるSQLは以下のようなイメージです、

SELECT `books`.`id`,
       `books`.`shop_id`,
       `books`.`uuid`,
       `books`.`name`,
       `books`.`description`,
       `books`.`created_at`,
       `books`.`updated_at`,
       `books`.`is_deleted`,
       (SELECT U0.`id`
        FROM `favorite_books` U0
        WHERE (U0.`user_id` = 999
          AND U0.`book_id` = (`books`.`id`))
        LIMIT 1) AS `favorite_book_id`
FROM books LEFT OUTER JOIN ...
WHERE ...
ORDER BY ...

おわりに

このようにまとめてみると非常にシンプルですが、実現したいことと実装すべきことが乖離しているような印象を個人的には受け、それが苦戦した要因かと思います。
本当はraw()で生のSQLを書けばそれほど悩まずに済み実装もよりわかりやすかったかもしれません。しかしながら、既存システムへの改修の場合そうもいかないケースはあると思います。今回のケースもSubQueryを使うことがベストな選択かはわかりませんが、選択肢の一つとして知っていてもよいかなと思います。