generic foreignkeyのsub relationをprefetchする方法

はじめに

generic foreignkey 自体のprefetchはできる。しかし、その更に先のrelationをprefetchすることができない。これをどうにかしようと言う苦肉の策を考えてみたという話。

言い訳

djangoのgeneric foreignkey関連のコードを読んでみたところ綺麗にできる方法は無さそうだった。実行時のprefetchの条件を上手く受け渡す方法が存在しなさそうだったので。仕方がないので thread localなcontext objectを作りそこでprefetchの条件を指定できるようにする。

概要

以下の様な形のモデルになっているとする。

Feed -- generic foreign key --> cotent = {A,B,C}

A -- 1:N --> xs = {X}
B -- 1:N --> ys = {Y}
C
X
Y

Feedというモデルが有りこれがgenericなrelationを持っており、A,B,Cのいずれかを保持する。また、AはXモデルとBはYモデルと1:Nの関係になっている。今回はFeedのqueryを取ってくる際にA,B,Cだけでなく、Aに結びつくX,Bに結びつくYも一緒に取ってくるようにしたい。

例えば以下のようなqueryを実行するとする。xs,ysの取得に関してはN+1になってしまう。

def use(content):
    if hasattr(content, "xs"):
        return [content, list(content.xs.all())]
    elif hasattr(content, "ys"):
        return [content, list(content.ys.all())]
    else:
        return [content]

content_list = []
for feed in Feed.objects.all().prefetch_related("content"):
    # content :: {A, B, C} はprefetchされるがそのsub relationであるxs,ysがprefetchされない
    content_list.append(use(feed.content))

この時以下のようにprefetchを指定するとエラーになる。

Feed.objects.all().prefetch_related("content", "content__xs", "content__ys")

generic foreignkeyのsub relationをprefetchする方法

試行錯誤を行なったdjangoのversionは1.9.5だった。

素朴な方法

基本的な方針としては以下のようになる。

GenericForeignKeyのget_content_typeが各relation(ここではA,B,C)を取ってくる際の始端となるオブジェクトになっている。これを書き換え、各relationを取ってくるquerysetを生成するところにprefetch_relatedを追加する処理を加えてあげる。

なので、django.contrib.contenttypes.fields.GenericForeignKey と　django.contrib.contenttypes.models.ContentType を自分で定義したサブクラスに置き換える。

class MyContentType(ContentType):
    class Meta:
        proxy = True

    def get_all_objects_for_this_type(self, **kwargs):
        qs = super().get_all_objects_for_this_type(**kwargs)
        return self.attach_prefetch(qs)

    def attach_prefetch(self, qs):
        # ここでprefetch_relatedの設定をする
        model = qs.model
        if issubclass(model, A):
            return qs.prefetch_related("xs")
        elif issubclass(model, B):
            return qs.prefetch_related("ys")
        else:
            return qs

class MyGenericForeignKey(GenericForeignKey):
    def get_content_type(self, *args, **kwargs):
        ct = super().get_content_type(*args, **kwargs)
        ct.__class__ = MyContentType  # これはmethodの挙動を書き換えるための雑な方法。真面目な方法ではない。
        return ct

各relationの取得に、MyContentTypeのget_all_objects_for_this_type()が使われる。これをフックするために、GenericForeignKeyからはMyContentTypeが使われるように get_content_type() を書き換える。内部の実装的にcontent typeのインスタンスの取得はキャッシュされているのでまじめに使う場合には注意が必要。

このようにすると以下のようなquery中でもX,Yを含めてprefetchしてくれるようになる。

content_list = []
for feed in Feed.objects.all().prefetch_related("content"):
    # sub relationであるxs, ysもprefetchされる
    content_list.append(use(feed.content))

ただし、上のようにした場合には、xs,ysのprefetchがデフォルトの動作になってしまう点が問題になる。できれば実行時にprefetchの条件を指定したい。

実行時にprefetchの条件をしていするための苦肉の策

冒頭の方にも書いたがdjangoの現状のコードセットではつらい。なので苦肉の策としてthread local objectにprefetchの条件を格納させることにする。まじめに実装するなら、異なる条件が二重に重なる場合なども考えなければいけなそうではあるけれど。そこまではやっていない。(本来は上手くqueryのhintを指定するカタチで情報を付加できたら良いのだけれど)

import threading
import contextlib

class GFKPrefetchContext:
    def __init__(self):
        self._context = threading.local()
        self._context.attach_prefetch = lambda qs: qs

    @contextlib.contextmanager
    def activate_prefetch(self, fn):
        oldvalue = self._context.attach_prefetch
        self._context.attach_prefetch = fn
        yield
        self._context.attach_prefetch = oldvalue

    def attach_prefetch(self, qs):
        return self._context.attach_prefetch(qs)

gfk_prefetch_context = GFKPrefetchContext()


class ContentTypeWithPrefetch(ContentType):
    class Meta:
        proxy = True

    def get_all_objects_for_this_type(self, **kwargs):
        qs = super().get_all_objects_for_this_type(**kwargs)
        return gfk_prefetch_context.attach_prefetch(qs)


class MyGenericForeignKey(GenericForeignKey):
    def get_content_type(self, *args, **kwargs):
        ct = super().get_content_type(*args, **kwargs)
        ct.__class__ = ContentTypeWithPrefetch
        return ct

以下の様にして使う。

def attach_prefetch(qs):
    model = qs.model
    if issubclass(model, A):
        return qs.prefetch_related("xs")
    elif issubclass(model, B):
        return qs.prefetch_related("ys")
    else:
        return qs

content_list = []
with gfk_prefetch_context.activate_prefetch(attach_prefetch):
    for feed in Feed.objects.all().prefetch_related("content"):
        content_list.append(use(feed.content))

# もちろんsub relationのprefetchを効かせたくなければcontextを指定しなければ良い
for feed in Feed.objects.all().prefetch_related("content"):
    content_list.append(use(feed.content))

補足

想定していたモデルの定義は以下のようなものだった。

class A(models.Model):
    name = models.CharField(max_length=32, default="", blank=False)

    class Meta:
        db_table = "a"

class B(models.Model):
    name = models.CharField(max_length=32, default="", blank=False)

    class Meta:
        db_table = "b"

class C(models.Model):
    name = models.CharField(max_length=32, default="", blank=False)

    class Meta:
        db_table = "c"

class X(models.Model):
    name = models.CharField(max_length=32, default="", blank=False)
    a = models.ForeignKey(A, related_name="xs")

    class Meta:
        db_table = "x"

class Y(models.Model):
    name = models.CharField(max_length=32, default="", blank=False)
    b = models.ForeignKey(B, related_name="ys")

    class Meta:
        db_table = "y"

class Feed(models.Model):
    class Meta:
        db_table = "feed"
        unique_together = ("content_type", "object_id")

    object_id = models.PositiveIntegerField()
    content_type = models.ForeignKey(ContentType)
    content = MyGenericForeignKey('content_type', 'object_id')