Fix Trunc database function with tzinfo parameter

Author: timgraham

django_mongodb_backend/features.py

@@ -58,10 +58,6 @@ class DatabaseFeatures(BaseDatabaseFeatures):
         "model_fields.test_jsonfield.TestQuerying.test_icontains",
         # MongoDB gives ROUND(365, -1)=360 instead of 370 like other databases.
         "db_functions.math.test_round.RoundTests.test_integer_with_negative_precision",
-        # Truncating in another timezone doesn't work becauase MongoDB converts
-        # the result back to UTC.
-        "db_functions.datetime.test_extract_trunc.DateFunctionWithTimeZoneTests.test_trunc_func_with_timezone",
-        "db_functions.datetime.test_extract_trunc.DateFunctionWithTimeZoneTests.test_trunc_timezone_applied_before_truncation",
         # Unexpected alias_refcount in alias_map.
         "queries.tests.Queries1Tests.test_order_by_tables",
         # The $sum aggregation returns 0 instead of None for null.

django_mongodb_backend/functions.py

@@ -1,4 +1,8 @@
+from datetime import datetime
+
+from django.conf import settings
 from django.db import NotSupportedError
+from django.db.models import DateField, DateTimeField, TimeField
 from django.db.models.expressions import Func
 from django.db.models.functions.comparison import Cast, Coalesce, Greatest, Least, NullIf
 from django.db.models.functions.datetime import (
@@ -195,6 +199,33 @@ def trunc(self, compiler, connection):
     return {"$dateTrunc": lhs_mql}
 
 
+def trunc_convert_value(self, value, expression, connection):
+    if connection.vendor == "mongodb":
+        # A custom TruncBase.convert_value() for MongoDB.
+        if value is None:
+            return None
+        convert_to_tz = settings.USE_TZ and self.get_tzname() != "UTC"
+        if isinstance(self.output_field, DateTimeField):
+            if convert_to_tz:
+                # Unlike other databases, MongoDB returns the value in UTC,
+                # so rather than setting the time zone equal to self.tzinfo,
+                # the value must be converted to tzinfo.
+                value = value.astimezone(self.tzinfo)
+        elif isinstance(value, datetime):
+            if isinstance(self.output_field, DateField):
+                if convert_to_tz:
+                    value = value.astimezone(self.tzinfo)
+                # Truncate for Trunc(..., output_field=DateField)
+                value = value.date()
+            elif isinstance(self.output_field, TimeField):
+                if convert_to_tz:
+                    value = value.astimezone(self.tzinfo)
+                # Truncate for Trunc(..., output_field=TimeField)
+                value = value.time()
+        return value
+    return self.convert_value(value, expression, connection)
+
+
 def trunc_date(self, compiler, connection):
     # Cast to date rather than truncate to date.
     lhs_mql = process_lhs(self, compiler, connection)
@@ -254,6 +285,7 @@ def register_functions():
     Substr.as_mql = substr
     Trim.as_mql = trim("trim")
     TruncBase.as_mql = trunc
+    TruncBase.convert_value = trunc_convert_value
     TruncDate.as_mql = trunc_date
     TruncTime.as_mql = trunc_time
     Upper.as_mql = preserve_null("toUpper")

django_mongodb_backend/operations.py

@@ -10,7 +10,7 @@
 from django.db.backends.base.operations import BaseDatabaseOperations
 from django.db.models import TextField
 from django.db.models.expressions import Combinable, Expression
-from django.db.models.functions import Cast
+from django.db.models.functions import Cast, Trunc
 from django.utils import timezone
 from django.utils.regex_helper import _lazy_re_compile
 
@@ -97,7 +97,11 @@ def get_db_converters(self, expression):
                 ]
             )
         elif internal_type == "DateField":
-            converters.append(self.convert_datefield_value)
+            # Trunc(... output_field="DateField") values must remain datetime
+            # until Trunc.convert_value() so they can be converted from UTC
+            # before truncation.
+            if not isinstance(expression, Trunc):
+                converters.append(self.convert_datefield_value)
         elif internal_type == "DateTimeField":
             if settings.USE_TZ:
                 converters.append(self.convert_datetimefield_value)
@@ -106,7 +110,11 @@ def get_db_converters(self, expression):
         elif internal_type == "JSONField":
             converters.append(self.convert_jsonfield_value)
         elif internal_type == "TimeField":
-            converters.append(self.convert_timefield_value)
+            # Trunc(... output_field="TimeField") values must remain datetime
+            # until Trunc.convert_value() so they can be converted from UTC
+            # before truncation.
+            if not isinstance(expression, Trunc):
+                converters.append(self.convert_timefield_value)
         elif internal_type == "UUIDField":
             converters.append(self.convert_uuidfield_value)
         return converters

docs/source/releases/5.1.x.rst

@@ -10,6 +10,8 @@ Django MongoDB Backend 5.1.x
 - Added support for a field's custom lookups and transforms in
   ``EmbeddedModelField``, e.g. ``ArrayField``’s ``contains``,
   ``contained__by``, ``len``, etc.
+- Fixed the results of queries that use the ``tzinfo`` parameter of the
+  ``Trunc`` database functions.
 
 .. _django-mongodb-backend-5.1.0-beta-2:
 

docs/source/topics/known-issues.rst

@@ -74,10 +74,6 @@ Database functions
     :class:`~django.db.models.functions.SHA512`
   - :class:`~django.db.models.functions.Sign`
 
-- The ``tzinfo`` parameter of the :class:`~django.db.models.functions.Trunc`
-  database functions doesn't work properly because MongoDB converts the result
-  back to UTC.
-
 Transaction management
 ======================