GeeWee / django-auto-prefetching

Automatic prefetching for Django
MIT License
231 stars 18 forks source link

FEATURE REQUEST: user can customize prefetch fields in ViewSet #40

Closed siteshen closed 2 years ago

siteshen commented 2 years ago

Reference Implementation:

class AutoPrefetchViewSetMixin:
    auto_prefetch_excluded_fields = set()
    auto_prefetch_extra_select_fields = set()
    auto_prefetch_extra_prefetch_fields = set()

    def get_auto_prefetch_excluded_fields(self):
        return self.auto_prefetch_excluded_fields

    def get_auto_prefetch_extra_select_fields(self):
        return self.auto_prefetch_extra_select_fields

    def get_auto_prefetch_extra_prefetch_fields(self):
        return self.auto_prefetch_extra_prefetch_fields

    def get_queryset(self):
        serializer = self.get_serializer()
        qs = super().get_queryset()

        kwargs = {
            "excluded_fields": self.get_auto_prefetch_excluded_fields(),
            "extra_select_fiedls": self.get_auto_prefetch_extra_select_fields(),
            "extra_prefetch_fields": self.get_auto_prefetch_extra_prefetch_fields(),
        }
        return prefetch(qs, serializer, **kwargs)

def prefetch(
    queryset,
    serializer: Type[ModelSerializer],
    *,
    excluded_fields=set(),
    extra_select_fields=set(),
    extra_prefetch_fields=set(),
):
    select_related, prefetch_related = _prefetch(serializer)
    select_related = select_related + extra_select_fields - excluded_fields
    prefetch_related = prefetch_related + extra_prefetch_fields - excluded_fields
    try:
        if select_related:
            queryset = queryset.select_related(*select_related)
        if prefetch_related:
            queryset = queryset.prefetch_related(*prefetch_related)
        return queryset
    except FieldError as e:
        raise ValueError(
            f"Calculated wrong field in select_related. Do you have a nested serializer for a ForeignKey where "
            f"you've forgotten to specify many=True? Original error: {e}"
        )
GeeWee commented 2 years ago

I'd be down for getting this in, but I'm pretty pressed for time. Would you want to submit a PR with the changes?

siteshen commented 2 years ago

I'd be down for getting this in, but I'm pretty pressed for time. Would you want to submit a PR with the changes?

Pull request created #41