Convert django RawQuerySet to Queryset

mishbah picture mishbah · Jan 16, 2018 · Viewed 12k times · Source

I have 2 Django models, ModelA with an ArrayField that is used to store a large list of primary key values (possibly 50k+ list)

class ModelA(models.Model):
    pk_values = ArrayField(models.IntegerField())

class CustomManager(manager.Manager):

    def get_for_index(self, index_id):
        qs = self.get_queryset()
        obj = ModelA.objects.get(pk=index_id)
        return qs.filter(id__in=obj.pk_values)

class ModelB(models.Model):
    # [...] some fields

    objects = CustomManager()

This works:

qs = ModelB.objects.get_for_index(index_id=1)

However, this would be super slow where "pk_values" is a large list.

So I tried doing raw SQL queries:

class CustomManager(manager.Manager):
    def get_for_index(self, index_id):
        qs = self.get_queryset()
        sql = "SELECT * FROM myapp_model_b JOIN myapp_model_a ON myapp_model_b.id = ANY(myapp_model_a.pk_values) WHERE myapp_model_a.id = '%s'" % index_id
        return qs.raw(sql)

But this returns a django.db.models.query.RawQuerySet instance.

But with this, I cant do things like queryset.values() afterwards.

How can I convert this to a normal Django queryset?

Is there a better way of doing this?

Docs:

Answer

emulbreh picture emulbreh · Jan 17, 2018

You can use a RawSQL expression:

ModelB.objects.filter(id__in=RawSQL(
    'SELECT unnest(a.pk_values) FROM app_modela a WHERE a.id = %s',
    [index_id]
))

Alternatively you can reproduce the exact query you have in your question with extra():

ModelB.objects.extra(
    tables=['foo_modela'],
    where=[
        '"app_modelb"."id" = ANY("app_modela"."pk_values")',
        '"app_modela"."id" = %s',
    ],
    params=[index_id],
)