KeyError calling `update()` after an `annotate()` and a `values()`

[Gav O'Connor]

Summary

We have noticed an issue when migrating from Django 5.1 to 5.2 in that we now get a KeyError where it previously worked without issue. I can't see anything in the changelog that would point towards this being an intended change.

The bug seems to happen when we try to call .update() on a QuerySet after we have called both .annotate() and .values() on it. A KeyError is raised with the name of the annotation.

Our real-world example is quite complex, but I have managed to simplify it somewhat in the example below.

Example

# models.py
from django.db import models

class Service(models.Model):
    name = models.CharField(max_length=255)
    slug = models.SlugField(max_length=255)

    def __str__(self):
        return self.name

class Order(models.Model):
    service_slug = models.CharField(max_length=255)

    def __str__(self):
        return self.id

class OrderLine(models.Model):
    order = models.ForeignKey(Order, on_delete=models.CASCADE)
    notes = models.CharField(max_length=255, blank=True, null=True)

    def __str__(self):
        return self.id

# tests.py
from django.test import TestCase
from django.db.models import OuterRef, Subquery
from .models import Service, Order, OrderLine

class MyTestCase(TestCase):
    def test_1(self):
        Service.objects.create(name="Green", slug="green")
        
        order = Order.objects.create(service_slug="green")
        OrderLine.objects.create(order=order)
        
        lines = OrderLine.objects.annotate(
            service=Subquery(
                Service.objects.filter(slug=OuterRef('order__service_slug')).values_list('name')[:1]
            )
        ).values(
            'id',
            'service',
        )
        
        lines.update(notes='foo')

Output

Django 5.1

❯ uv add django==5.1
❯ uv run manage.py test
Found 1 test(s).
.
----------------------------------------------------------------------
Ran 1 test in 0.001s

OK

Django 5.2

❯ uv add django==5.2
❯ uv run manage.py test
Found 1 test(s).
E
======================================================================
ERROR: test_1 (core.tests.MyTestCase.test_1)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/Users/gav/code/annotatebug/core/tests.py", line 21, in test_1
    lines.update(notes='foo')
    ~~~~~~~~~~~~^^^^^^^^^^^^^
  File "/Users/gav/code/annotatebug/.venv/lib/python3.13/site-packages/django/db/models/query.py", line 1258, in update
    rows = query.get_compiler(self.db).execute_sql(ROW_COUNT)
  File "/Users/gav/code/annotatebug/.venv/lib/python3.13/site-packages/django/db/models/sql/compiler.py", line 2059, in execute_sql
    row_count = super().execute_sql(result_type)
  File "/Users/gav/code/annotatebug/.venv/lib/python3.13/site-packages/django/db/models/sql/compiler.py", line 1609, in execute_sql
    sql, params = self.as_sql()
                  ~~~~~~~~~~~^^
  File "/Users/gav/code/annotatebug/.venv/lib/python3.13/site-packages/django/db/models/sql/compiler.py", line 1988, in as_sql
    self.pre_sql_setup()
    ~~~~~~~~~~~~~~~~~~^^
  File "/Users/gav/code/annotatebug/.venv/lib/python3.13/site-packages/django/db/models/sql/compiler.py", line 2110, in pre_sql_setup
    super().pre_sql_setup()
    ~~~~~~~~~~~~~~~~~~~~~^^
  File "/Users/gav/code/annotatebug/.venv/lib/python3.13/site-packages/django/db/models/sql/compiler.py", line 85, in pre_sql_setup
    self.setup_query(with_col_aliases=with_col_aliases)
    ~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/gav/code/annotatebug/.venv/lib/python3.13/site-packages/django/db/models/sql/compiler.py", line 74, in setup_query
    self.select, self.klass_info, self.annotation_col_map = self.get_select(
                                                            ~~~~~~~~~~~~~~~^
        with_col_aliases=with_col_aliases,
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    )
    ^
  File "/Users/gav/code/annotatebug/.venv/lib/python3.13/site-packages/django/db/models/sql/compiler.py", line 283, in get_select
    expression = self.query.annotations[expression]
                 ~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^
KeyError: 'service'

----------------------------------------------------------------------
Ran 1 test in 0.003s

FAILED (errors=1)

[Simon Charette]

Regression in 65ad4ade74dc9208b9d686a451cd6045df0c9c3a that isn't fixed by 543e17c4405dfdac4f18759fc78b190406d14239.

Not sure why we generate a SELECT clause in the first place for UPDATE queries but I'm pretty sure the issue relates to ​this line which naively clear the annotation mask as added to resolve #19513, #18580 in a84344bc539c66589c8d4fe30c6ceaecf8ba1af3.

[Sarah Boyce <42296566+sarahboyce@…>]

In 7f6a5fbe:

[5.2.x] Fixed #36360 -- Fixed QuerySet.update() crash when referring annotations through values().

The issue was only manifesting itself when also filtering againt a related
model as that forces the usage of a subquery because SQLUpdateCompiler doesn't
support the UPDATE FROM syntax yet.

Regression in 65ad4ade74dc9208b9d686a451cd6045df0c9c3a.

Refs #28900.

Thanks Gav O'Connor for the detailed report.

Backport of 8ef4e0bd423ac3764004c73c3d1098e7a51a2945 from main.

[Sarah Boyce]

In ​8ef4e0b:

Fixed #36360 -- Fixed QuerySet.update() crash when referring annotations through values().

The issue was only manifesting itself when also filtering againt a related
model as that forces the usage of a subquery because SQLUpdateCompiler doesn't
support the UPDATE FROM syntax yet.

Regression in 65ad4ade74dc9208b9d686a451cd6045df0c9c3a.

Refs #28900.

Thanks Gav O'Connor for the detailed report.