Fix Trunc database function with tzinfo parameter

Author: timgraham

django_mongodb_backend/features.py

@@ -57,10 +57,6 @@ class DatabaseFeatures(BaseDatabaseFeatures):
         # Pattern lookups that use regexMatch don't work on JSONField:
         # Unsupported conversion from array to string in $convert
         "model_fields.test_jsonfield.TestQuerying.test_icontains",
-        # 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 import JSONArray
 from django.db.models.functions.comparison import Cast, Coalesce, Greatest, Least, NullIf
@@ -196,6 +200,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.
+        convert_to_tz = settings.USE_TZ and self.get_tzname() != "UTC"
+        if isinstance(self.output_field, DateTimeField):
+            if not convert_to_tz:
+                pass
+            elif value is not None:
+                # 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 value is None:
+                pass
+            elif isinstance(self.output_field, DateField):
+                if convert_to_tz:
+                    value = value.astimezone(self.tzinfo)
+                value = value.date()
+            elif isinstance(self.output_field, TimeField):
+                if convert_to_tz:
+                    value = value.astimezone(self.tzinfo)
+                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)
@@ -256,6 +287,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/topics/known-issues.rst

@@ -75,10 +75,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
 ======================