graphene-django/graphene_django/optimization.py
2017-07-18 16:30:00 -05:00

83 lines
2.9 KiB
Python

from collections import namedtuple
from django.db.models import ForeignKey
from django.db.models.fields.reverse_related import ForeignObjectRel
from graphene.utils.str_converters import to_snake_case
from .registry import get_global_registry
from .utils import get_related_model
REGISTRY = get_global_registry()
SELECT = 'select'
PREFETCH = 'prefetch'
RelatedSelection = namedtuple('RelatedSelection', ['name', 'fetch_type'])
def model_fields_as_dict(model):
return dict((f.name, f) for f in model._meta.get_fields())
def get_related_fetches_for_model(model, graphql_ast):
model_fields = model_fields_as_dict(model)
selections = graphql_ast.selection_set.selections
graphene_obj_type = REGISTRY.get_type_for_model(model)
optimizations = {}
if graphene_obj_type and graphene_obj_type._meta.optimizations:
optimizations = graphene_obj_type._meta.optimizations
relateds = []
for selection in selections:
selection_name = to_snake_case(selection.name.value)
selection_field = model_fields.get(selection_name, None)
try:
related_model = get_related_model(selection_field)
except:
# This is not a ForeignKey or Relation, check manual optimizations
manual_optimizations = optimizations.get(selection_name)
if manual_optimizations:
for manual_select in manual_optimizations.get(SELECT, []):
relateds.append(RelatedSelection(manual_select, SELECT))
for manual_prefetch in manual_optimizations.get(PREFETCH, []):
relateds.append(RelatedSelection(manual_prefetch, PREFETCH))
continue
query_name = selection_field.name
if isinstance(selection_field, ForeignObjectRel):
query_name = selection_field.field.related_query_name()
nested_relateds = get_related_fetches_for_model(related_model, selection)
related_type = PREFETCH # default to prefetch, it's safer
if isinstance(selection_field, ForeignKey):
related_type = SELECT # we can only select for ForeignKeys
if nested_relateds:
for related in nested_relateds:
full_name = '{0}__{1}'.format(query_name, related.name)
nested_related_type = PREFETCH
if related_type == SELECT and related.fetch_type == SELECT:
nested_related_type = related_type
relateds.append(RelatedSelection(full_name, nested_related_type))
else:
relateds.append(RelatedSelection(query_name, related_type))
return relateds
def optimize_queryset(model, queryset, graphql_ast):
relateds = get_related_fetches_for_model(model, graphql_ast)
for related in relateds:
if related.fetch_type == SELECT:
queryset = queryset.select_related(related.name)
else:
queryset = queryset.prefetch_related(related.name)
return queryset