Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Automatically annotate Django models retrieved through a ManyToMany field

I have a Student model that's related to a Course model through a ManyToMany field:

class Student(Model):
    name = TextField()

class Course(Model):
    name = TextField()
    students = ManyToManyField(Student, through="StudentCourse")

class StudentCourse(Model):
    student = ForeignKey(Student)
    course = ForeignKey(Course)
    section_name = TextField()

How can I automatically annotate students retrieved through the Course.students many-to-many field with their section in that course?

For example, instead of having to add the explicit extra on each query:

>>> students = course.students.extra(select={
...     "section_name": "app_student_course.section_name",
... })
>>> print students[0].section_name
u'First Section'

I can just:

>>> students = course.students.all()
>>> print students[0].section_name
u'First Section'

Thanks!

like image 982
David Wolever Avatar asked Jul 07 '16 18:07

David Wolever


2 Answers

Is it possible to replace the relation manager as follows. From there you can do whatever you want with the queryset:

from django.db.models.query import F

class Student(Model):
    name = TextField()

class Course(Model):
    name = TextField()
    students = ManyToManyField(Student, through="StudentCourse")

class StudentCourse(Model):
    # Set related_names so that it is easy to refer to the relation
    # with the through model
    student = ForeignKey(Student, related_name='student_courses')
    course = ForeignKey(Course, related_name='student_courses')
    section_name = TextField()

# Create a new customized manager
class StudentRelatedWithCourseManager(
    Course.students.related_manager_cls
):

    def get_queryset(self):
        # Gets the queryset of related Students ... 
        qs = super(StudentRelatedWithCourseManager, self)\
            .get_queryset()

        # Annotate it before is handled by the ManyRelatedDescriptor
        return qs.annotate(
            section_name=F('student_courses__section_name')
        )

# Replace the stock manager with your custom manager
Course.students.related_manager_cls = \
    StudentRelatedWithCourseManager
like image 138
ajaest Avatar answered Oct 13 '22 23:10

ajaest


I would say the problem here is very similar with the one here, so my solution is almost identical to the one provided there.

I think prefetch_related with custom Prefetch is more suitable for such kind of problems instead of .annotations or .extra clauses. The benefit is that you get the whole related object instead of a single bit of it (so you can use more metadata), and there is zero performance drawback, only bigger memory footprint which can cause problems if its used for a large set of objects.

class StudenQuerySet(models.QuerySet):
    def prefetch_sections_for_course(self, course):
        return self.prefetch_related(models.Prefetch('studentcourse_set',
            queryset=StudentCourse.objects.filter(course=course),
            to_attr='course_sections'
        ))

class Student(Model):
    name = TextField()

    objects = StudenQuerySet.as_manager()

    @property
    def course_section_name(self):
        try:
            return self.course_sections[0].section_name
        except IndexError:
            #This student is not related with the prefetched course
            return None
        except AttributeError:
            raise AttributeError('You forgot to prefetch course sections')
            #or you can just
            return None

#usage
students = course.students.all().prefetch_sections_for_course(course)
for student in students:
    print student.course_section_name
like image 33
Todor Avatar answered Oct 14 '22 00:10

Todor