Annotations returning pure PK value instead of Django HashID

22 views Asked by At

In my application the user can assume one of 3 different roles.

These users can be assigned to programs, and they can be assigned to 3 different fields, each of those exclusevely to a role.

In my API I'm trying to query all my users and annotate the programs they are in:

def get_queryset(self):
    queryset = (
        User.objects
        .select_related('profile')
        .prefetch_related('managed_programs', 'supervised_programs', 'case_managed_programs')
        .annotate(
            programs=Case(
                When(
                    role=RoleChoices.PROGRAM_MANAGER,
                    then=ArraySubquery(
                        Program.objects.filter(program_managers=OuterRef('pk'), is_active=True)
                        .values('id', 'name')
                        .annotate(
                            data=Func(
                                Value('{"id": '),
                                F('id'),
                                Value(', "name": "'),
                                F('name'),
                                Value('"}'),
                                function='concat',
                                output_field=CharField(),
                            )
                        )
                        .values('data')
                    ),
                ),
                When(
                    role=RoleChoices.SUPERVISOR,
                    then=ArraySubquery(
                        Program.objects.filter(supervisors=OuterRef('pk'), is_active=True)
                        .values('id', 'name')
                        .annotate(
                            data=Func(
                                Value('{"id": '),
                                F('id'),
                                Value(', "name": "'),
                                F('name'),
                                Value('"}'),
                                function='concat',
                                output_field=CharField(),
                            )
                        )
                        .values('data')
                    ),
                ),
                When(
                    role=RoleChoices.CASE_MANAGER,
                    then=ArraySubquery(
                        Program.objects.filter(case_managers=OuterRef('pk'), is_active=True)
                        .values('id', 'name')
                        .annotate(
                            data=Func(
                                Value('{"id": '),
                                F('id'),
                                Value(', "name": "'),
                                F('name'),
                                Value('"}'),
                                function='concat',
                                output_field=CharField(),
                            )
                        )
                        .values('data')
                    ),
                ),
                default=Value(list()),
                output_field=ArrayField(base_field=JSONField()),
            )
        )
        .order_by('id')
    )
    return queryset

This works (almost) flawlessly and gives me only 5 DB hits, perfect, or not...

The problem is that I'm using Django HashID fields for the Program PK, and this query returns the pure integer value for each Program.

I've tried a more "normal" approach, by getting the data using a SerializerMethodField:

@staticmethod
def get_programs(obj):
    role_attr = {
        RoleChoices.PROGRAM_MANAGER: 'managed_programs',
        RoleChoices.SUPERVISOR: 'supervised_programs',
        RoleChoices.CASE_MANAGER: 'case_managed_programs',
    }
    try:
        programs =  getattr(obj, role_attr[obj.role], None).values_list('id', 'name')
        return [{'id': str(id), 'name': name} for id, name in programs]
    except (AttributeError, KeyError):
        return []

This gives me the result I need, but the query quantity skyrockets. It seems that it's not taking advantage of the prefetch_related, but I don't understand how is this possible, considering I'm using the same queryset.

So, I have two options here:

  • Use the annotations but having the HashID returning, instead of integer PK
  • Have the SerializerMethodField reuse the prefetched data, instead of requerying

Is there a way to accomplish any of those?

EDIT:

A small heads-up, I've decided to use the first approach and Hash the ID manually inside the serializer

programs = serializers.SerializerMethodField()

@staticmethod
def get_programs(obj):
    return [
        {"id": str(Hashid(value=program['id'], salt=settings.SECRET_KEY, min_length=13)), "name": program['name']} for program in obj.programs
    ]

For now it works, but I'd be more satisfied if there's a more direct way to accomplish this.

0

There are 0 answers