Fixed #31496 -- Fixed QuerySet.values()/values_list() crash on combined querysets ordered by annotations.

This commit is contained in:
David-Wobrock 2020-10-10 17:41:13 +02:00 committed by Mariusz Felisiak
parent 300def5c34
commit 464a4c0c59
3 changed files with 41 additions and 10 deletions

View File

@ -322,10 +322,15 @@ class SQLCompiler:
if col in self.query.annotations:
# References to an expression which is masked out of the SELECT
# clause.
expr = self.query.annotations[col]
if isinstance(expr, Value):
# output_field must be resolved for constants.
expr = Cast(expr, expr.output_field)
if self.query.combinator and self.select:
# Don't use the resolved annotation because other
# combinated queries might define it differently.
expr = F(col)
else:
expr = self.query.annotations[col]
if isinstance(expr, Value):
# output_field must be resolved for constants.
expr = Cast(expr, expr.output_field)
order_by.append((OrderBy(expr, descending=descending), False))
continue
@ -378,10 +383,14 @@ class SQLCompiler:
else:
if col_alias:
raise DatabaseError('ORDER BY term does not match any column in the result set.')
# Add column used in ORDER BY clause without an alias to
# the selected columns.
self.query.add_select_col(src)
resolved.set_source_expressions([RawSQL('%d' % len(self.query.select), ())])
# Add column used in ORDER BY clause to the selected
# columns and to each combined query.
order_by_idx = len(self.query.select) + 1
col_name = f'__orderbycol{order_by_idx}'
for q in self.query.combined_queries:
q.add_annotation(expr_src, col_name)
self.query.add_select_col(resolved, col_name)
resolved.set_source_expressions([RawSQL(f'{order_by_idx}', ())])
sql, params = self.compile(resolved)
# Don't add the same column twice, but the order direction is
# not taken into account so we strip it. When this entire method

View File

@ -1881,9 +1881,9 @@ class Query(BaseExpression):
self.select = ()
self.values_select = ()
def add_select_col(self, col):
def add_select_col(self, col, name):
self.select += col,
self.values_select += col.output_field.name,
self.values_select += name,
def set_select(self, cols):
self.default_cols = False

View File

@ -211,6 +211,28 @@ class QuerySetSetOperationTests(TestCase):
with self.subTest(qs=qs):
self.assertEqual(list(qs), expected_result)
def test_union_with_values_list_and_order_on_annotation(self):
qs1 = Number.objects.annotate(
annotation=Value(-1),
multiplier=F('annotation'),
).filter(num__gte=6)
qs2 = Number.objects.annotate(
annotation=Value(2),
multiplier=F('annotation'),
).filter(num__lte=5)
self.assertSequenceEqual(
qs1.union(qs2).order_by('annotation', 'num').values_list('num', flat=True),
[6, 7, 8, 9, 0, 1, 2, 3, 4, 5],
)
self.assertQuerysetEqual(
qs1.union(qs2).order_by(
F('annotation') * F('multiplier'),
'num',
).values('num'),
[6, 7, 8, 9, 0, 1, 2, 3, 4, 5],
operator.itemgetter('num'),
)
def test_count_union(self):
qs1 = Number.objects.filter(num__lte=1).values('num')
qs2 = Number.objects.filter(num__gte=2, num__lte=3).values('num')