Skip to content

Commit

Permalink
Add test validating query performance with select_related + prefetch_…
Browse files Browse the repository at this point in the history
…related

This test passes after reverting the `CustomField` resolver change
introduced in
graphql-python/graphene-django#1315, but fails
with that resolver code present. For instance, adding back the resolver
code gives a test failure showing:

```
Failed: Expected to perform 2 queries but 11 were done
```

This should ensure there aren't regressions that prevent
query-optimization in the future.
  • Loading branch information
azuredeviil053 committed May 3, 2023
1 parent 0142c5f commit 85df0c0
Showing 1 changed file with 150 additions and 3 deletions.
153 changes: 150 additions & 3 deletions graphene_django/tests/test_fields.py
@@ -1,14 +1,19 @@
import datetime
from django.db.models import Count
import re
from django.db.models import Count, Prefetch

import pytest

from graphene import List, NonNull, ObjectType, Schema, String

from ..fields import DjangoListField
from ..types import DjangoObjectType
from .models import Article as ArticleModel
from .models import Reporter as ReporterModel
from .models import (
Article as ArticleModel,
Film as FilmModel,
FilmDetails as FilmDetailsModel,
Reporter as ReporterModel,
)


class TestDjangoListField:
Expand Down Expand Up @@ -500,3 +505,145 @@ class Query(ObjectType):

assert not result.errors
assert result.data == {"reporters": [{"firstName": "Tara"}]}

def test_select_related_and_prefetch_related_are_respected(
self, django_assert_num_queries
):
class Article(DjangoObjectType):
class Meta:
model = ArticleModel
fields = ("headline", "editor", "reporter")

class Film(DjangoObjectType):
class Meta:
model = FilmModel
fields = ("genre", "details")

class FilmDetail(DjangoObjectType):
class Meta:
model = FilmDetailsModel
fields = ("location",)

class Reporter(DjangoObjectType):
class Meta:
model = ReporterModel
fields = ("first_name", "articles", "films")

class Query(ObjectType):
articles = DjangoListField(Article)

@staticmethod
def resolve_articles(root, info):
# Optimize for querying associated editors and reporters, and the films and film
# details of those reporters. This is similar to what would happen using a library
# like https://github.com/tfoxy/graphene-django-optimizer for a query like the one
# below (albeit simplified and hardcoded here).
return ArticleModel.objects.select_related(
"editor", "reporter"
).prefetch_related(
Prefetch(
"reporter__films",
queryset=FilmModel.objects.select_related("details"),
),
)

schema = Schema(query=Query)

query = """
query {
articles {
headline
editor {
firstName
}
reporter {
firstName
films {
genre
details {
location
}
}
}
}
}
"""

r1 = ReporterModel.objects.create(first_name="Tara", last_name="West")
r2 = ReporterModel.objects.create(first_name="Debra", last_name="Payne")

ArticleModel.objects.create(
headline="Amazing news",
reporter=r1,
pub_date=datetime.date.today(),
pub_date_time=datetime.datetime.now(),
editor=r2,
)
ArticleModel.objects.create(
headline="Not so good news",
reporter=r2,
pub_date=datetime.date.today(),
pub_date_time=datetime.datetime.now(),
editor=r1,
)

film1 = FilmModel.objects.create(genre="ac")
film2 = FilmModel.objects.create(genre="ot")
film3 = FilmModel.objects.create(genre="do")
FilmDetailsModel.objects.create(location="Hollywood", film=film1)
FilmDetailsModel.objects.create(location="Antarctica", film=film3)
r1.films.add(film1, film2)
r2.films.add(film3)

# We expect 2 queries to be performed based on the above resolver definition: one for all
# articles joined with the reporters model (for associated editors and reporters), and one
# for the films prefetch (which includes its `select_related` JOIN logic in its queryset)
with django_assert_num_queries(2) as captured:
result = schema.execute(query)

assert not result.errors
assert result.data == {
"articles": [
{
"headline": "Amazing news",
"editor": {"firstName": "Debra"},
"reporter": {
"firstName": "Tara",
"films": [
{"genre": "AC", "details": {"location": "Hollywood"}},
{"genre": "OT", "details": None},
],
},
},
{
"headline": "Not so good news",
"editor": {"firstName": "Tara"},
"reporter": {
"firstName": "Debra",
"films": [
{"genre": "DO", "details": {"location": "Antarctica"}},
],
},
},
]
}

assert len(captured.captured_queries) == 2 # Sanity-check

# First we should have queried for all articles in a single query, joining on the reporters
# model (for the editors and reporters ForeignKeys)
assert re.match(
r'SELECT .* "tests_article" INNER JOIN "tests_reporter"',
captured.captured_queries[0]["sql"],
)

# Then we should have queried for all of the films of all reporters, joined with the film
# details for each film, using a single query
assert re.match(
r'SELECT .* FROM "tests_film" INNER JOIN "tests_film_reporters" .* LEFT OUTER JOIN "tests_filmdetails"',
captured.captured_queries[1]["sql"],
)

0 comments on commit 85df0c0

Please sign in to comment.