Skip to content

Instantly share code, notes, and snippets.

@kalda341
Created February 4, 2019 01:34
Show Gist options
  • Select an option

  • Save kalda341/0d66570d2cf8130c11a36ec041674e78 to your computer and use it in GitHub Desktop.

Select an option

Save kalda341/0d66570d2cf8130c11a36ec041674e78 to your computer and use it in GitHub Desktop.
class SerializerQuerysetMixin():
excluded_prefetches = []
@classmethod
def get_queryset(cls, request):
"""
This is the base queryset - will be use for resources that are even just referenced
"""
return cls.Meta.model.objects.all().order_by('id') # .distinct('id')
@classmethod
def get_resource_queryset(cls, request):
"""
This is the queryset that will be used when we're actually including the serialized form of the
resource
"""
return cls.get_queryset(request)
@classmethod
def get_prefetched_queryset(cls, request, includes=[], from_serializer=None, depth=0):
queryset = cls.get_resource_queryset(request)
included_serializers = get_included_serializers(cls)
for relation_key, serializer in included_serializers.items():
# Give us an escape hatch for problematic relations
if relation_key in cls.excluded_prefetches:
LOGGER.info('Skipping prefetch for %s as it has been excluded', relation_key)
continue
if serializer == from_serializer:
LOGGER.warning('Skipping reverse direction')
# TODO: Make this less destructive
continue
if relation_key in includes:
related_queryset = cls.get_included_prefetch(
request=request,
relation_key=relation_key,
serializer=serializer,
sub_includes=includes[relation_key],
depth=depth,
)
else:
related_queryset = cls.get_regular_prefetch(
request=request,
relation_key=relation_key,
serializer=serializer,
depth=depth
)
queryset = queryset.prefetch_related(
# One day we may not need to flatten them!
*cls.flatten_prefetches(Prefetch(relation_key, queryset=related_queryset))
)
return queryset
@classmethod
def flatten_prefetches(cls, prefetch):
# Absolute hack to pull prefetches back to the top level (because Django doesn't handle nested prefetches
# very well)
# Hopefully one day this can be removed
flattened = []
queryset = prefetch.queryset
# We only handle one layer. This could be made recursive if we ever needed it to be
for nested_prefetch in queryset._prefetch_related_lookups:
if not isinstance(nested_prefetch, Prefetch):
nested_prefetch = Prefetch(nested_prefetch)
else:
# Let's be kind and not modify our input parameters
nested_prefetch = copy.copy(nested_prefetch)
nested_prefetch.add_prefix(prefetch.prefetch_through)
flattened.append(nested_prefetch)
flattened.insert(0, Prefetch(
lookup=prefetch.prefetch_through,
# Clear the prefetches that we have covered
queryset=prefetch.queryset.prefetch_related(None),
to_attr=prefetch.to_attr,
))
return flattened
@classmethod
def get_regular_prefetch(cls, request, relation_key, serializer, depth):
LOGGER.info('\t' * depth + 'Including regular prefetch for %s', relation_key)
return serializer.get_queryset(request=request)
@classmethod
def get_included_prefetch(cls, request, relation_key, serializer, sub_includes, depth):
LOGGER.info('\t' * depth + 'Including nested resource prefetch for %s', relation_key)
return serializer.get_prefetched_queryset(
request=request,
includes=sub_includes,
depth=depth + 1,
from_serializer=serializer
)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment