Refs #30188 -- Prevented double annotation of subquery when aggregated over.

Thanks Can Sarıgöl for the suggested trimming approach.
This commit is contained in:
Simon Charette 2019-03-08 15:54:03 -05:00 committed by Tim Graham
parent bdc07f176e
commit d1e9c25162
2 changed files with 23 additions and 15 deletions

View File

@ -383,12 +383,16 @@ class Query(BaseExpression):
new_expr, col_cnt = self.rewrite_cols(expr, col_cnt) new_expr, col_cnt = self.rewrite_cols(expr, col_cnt)
new_exprs.append(new_expr) new_exprs.append(new_expr)
elif isinstance(expr, (Col, Subquery)) or (expr.contains_aggregate and not expr.is_summary): elif isinstance(expr, (Col, Subquery)) or (expr.contains_aggregate and not expr.is_summary):
# Reference to column. Make sure the referenced column # Reference to column, subquery, or another aggregate. Make
# is selected. # sure the expression is selected and reuse its alias if so.
col_cnt += 1 for col_alias, selected_annotation in self.annotation_select.items():
col_alias = '__col%d' % col_cnt if selected_annotation == expr:
self.annotations[col_alias] = expr break
self.append_annotation_mask([col_alias]) else:
col_cnt += 1
col_alias = '__col%d' % col_cnt
self.annotations[col_alias] = expr
self.append_annotation_mask([col_alias])
new_exprs.append(Ref(col_alias, expr)) new_exprs.append(Ref(col_alias, expr))
else: else:
# Some other expression not referencing database values # Some other expression not referencing database values

View File

@ -553,16 +553,20 @@ class BasicExpressionsTests(TestCase):
@skipUnlessDBFeature('supports_subqueries_in_group_by') @skipUnlessDBFeature('supports_subqueries_in_group_by')
def test_aggregate_subquery_annotation(self): def test_aggregate_subquery_annotation(self):
aggregate = Company.objects.annotate( with self.assertNumQueries(1) as ctx:
ceo_salary=Subquery( aggregate = Company.objects.annotate(
Employee.objects.filter( ceo_salary=Subquery(
id=OuterRef('ceo_id'), Employee.objects.filter(
).values('salary') id=OuterRef('ceo_id'),
), ).values('salary')
).aggregate( ),
ceo_salary_gt_20=Count('pk', filter=Q(ceo_salary__gt=20)), ).aggregate(
) ceo_salary_gt_20=Count('pk', filter=Q(ceo_salary__gt=20)),
)
self.assertEqual(aggregate, {'ceo_salary_gt_20': 1}) self.assertEqual(aggregate, {'ceo_salary_gt_20': 1})
# Aggregation over a subquery annotation doesn't annotate the subquery
# twice in the inner query.
self.assertLessEqual(ctx.captured_queries[0]['sql'].count('SELECT'), 4,)
def test_explicit_output_field(self): def test_explicit_output_field(self):
class FuncA(Func): class FuncA(Func):