Fix renamed annotations relying on group by clauses not working

Co-authored-by: Hannes Engelhardt <[email protected]>

Author: Photonios

psqlextra/sql.py

@@ -1,4 +1,5 @@
 from collections import OrderedDict
+from collections.abc import Iterable
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 import django
@@ -7,6 +8,7 @@
 from django.db import connections, models
 from django.db.models import Expression, sql
 from django.db.models.constants import LOOKUP_SEP
+from django.db.models.expressions import Ref
 
 from .compiler import PostgresInsertOnConflictCompiler
 from .compiler import SQLUpdateCompiler as PostgresUpdateCompiler
@@ -74,6 +76,14 @@ def rename_annotations(self, annotations) -> None:
                     self.annotation_select_mask.remove(old_name)
                     self.annotation_select_mask.append(new_name)
 
+        if isinstance(self.group_by, Iterable):
+            for statement in self.group_by:
+                if not isinstance(statement, Ref):
+                    continue
+
+                if statement.refs in annotations:  # type: ignore[attr-defined]
+                    statement.refs = annotations[statement.refs]  # type: ignore[attr-defined]
+
         self.annotations.clear()
         self.annotations.update(new_annotations)
 

settings.py

@@ -24,3 +24,6 @@
     'psqlextra',
     'tests',
 )
+
+USE_TZ = True
+TIME_ZONE = 'UTC'

tests/test_query.py

@@ -1,5 +1,8 @@
+from datetime import datetime, timezone
+
 from django.db import connection, models
-from django.db.models import Case, F, Q, Value, When
+from django.db.models import Case, F, Min, Q, Value, When
+from django.db.models.functions.datetime import TruncSecond
 from django.test.utils import CaptureQueriesContext, override_settings
 
 from psqlextra.expressions import HStoreRef
@@ -96,6 +99,40 @@ def test_query_annotate_in_expression():
     assert result.is_he_henk == "really henk"
 
 
+def test_query_annotate_group_by():
+    """Tests whether annotations with GROUP BY clauses are properly renamed
+    when the annotation overwrites a field name."""
+
+    model = get_fake_model(
+        {
+            "name": models.TextField(),
+            "timestamp": models.DateTimeField(null=False),
+            "value": models.IntegerField(),
+        }
+    )
+
+    timestamp = datetime(2024, 1, 1, 0, 0, 0, 0, tzinfo=timezone.utc)
+
+    model.objects.create(name="me", timestamp=timestamp, value=1)
+
+    result = (
+        model.objects.values("name")
+        .annotate(
+            timestamp=TruncSecond("timestamp", tzinfo=timezone.utc),
+            value=Min("value"),
+        )
+        .values_list(
+            "name",
+            "value",
+            "timestamp",
+        )
+        .order_by("name")
+        .first()
+    )
+
+    assert result == ("me", 1, timestamp)
+
+
 def test_query_hstore_value_update_f_ref():
     """Tests whether F(..) expressions can be used in hstore values when
     performing update queries."""