mirror of
https://github.com/graphql-python/graphene-django.git
synced 2025-12-08 10:44:28 +03:00
95 lines
3.3 KiB
Python
95 lines
3.3 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 find_model_selections(ast):
|
|
selections = ast.selection_set.selections
|
|
|
|
for selection in selections:
|
|
if selection.name.value == 'edges':
|
|
for sub_selection in selection.selection_set.selections:
|
|
if sub_selection.name.value == 'node':
|
|
return sub_selection.selection_set.selections
|
|
|
|
return selections
|
|
|
|
|
|
def get_related_fetches_for_model(model, graphql_ast):
|
|
model_fields = model_fields_as_dict(model)
|
|
selections = find_model_selections(graphql_ast)
|
|
|
|
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
|