add support for datetime database functions

Author: timgraham

.github/workflows/test-python.yml

@@ -63,7 +63,7 @@ jobs:
       - name: Start MongoDB
         uses: supercharge/[email protected]
         with:
-          mongodb-version: 4.4
+          mongodb-version: 5.0
       - name: Run tests
         run: >
           python3 django_repo/tests/runtests.py --settings mongodb_settings -v 2
@@ -74,6 +74,7 @@ jobs:
           bulk_create
           dates
           datetimes
+          db_functions.datetime
           db_functions.math
           empty
           defer

README.md

@@ -62,6 +62,16 @@ DATABASES = {
 
 - `DateTimeField` doesn't support microsecond precision.
 
+- The following database functions aren't supported:
+    - `ExtractQuarter`
+    - `Now`
+    - `Sign`
+    - `TruncDate`
+    - `TruncTime`
+
+- The `tzinfo` parameter of the `Trunc` database functions doesn't work
+  properly because MongoDB converts the result back to UTC.
+
 ## Troubleshooting
 
 TODO

django_mongodb/features.py

@@ -19,19 +19,8 @@ class DatabaseFeatures(BaseDatabaseFeatures):
         # Database defaults not supported: bson.errors.InvalidDocument:
         # cannot encode object: <django.db.models.expressions.DatabaseDefault
         "basic.tests.ModelInstanceCreationTests.test_save_primary_with_db_default",
-        # Date lookups aren't implemented: https://github.com/mongodb-labs/django-mongodb/issues/9
-        # (e.g. ExtractWeekDay is not supported.)
-        "basic.tests.ModelLookupTest.test_does_not_exist",
-        "basic.tests.ModelLookupTest.test_equal_lookup",
-        "basic.tests.ModelLookupTest.test_rich_lookup",
+        # Query for chained lookups not generated correctly.
         "lookup.tests.LookupTests.test_chain_date_time_lookups",
-        "lookup.test_timefield.TimeFieldLookupTests.test_hour_lookups",
-        "lookup.test_timefield.TimeFieldLookupTests.test_minute_lookups",
-        "lookup.test_timefield.TimeFieldLookupTests.test_second_lookups",
-        "timezones.tests.LegacyDatabaseTests.test_query_datetime_lookups",
-        "timezones.tests.NewDatabaseTests.test_query_convert_timezones",
-        "timezones.tests.NewDatabaseTests.test_query_datetime_lookups",
-        "timezones.tests.NewDatabaseTests.test_query_datetime_lookups_in_other_timezone",
         # 'NulledTransform' object has no attribute 'as_mql'.
         "lookup.tests.LookupTests.test_exact_none_transform",
         # "Save with update_fields did not affect any rows."
@@ -62,6 +51,10 @@ class DatabaseFeatures(BaseDatabaseFeatures):
         "db_functions.math.test_log.LogTests.test_decimal",
         # 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",
     }
 
     django_test_skips = {
@@ -171,6 +164,8 @@ class DatabaseFeatures(BaseDatabaseFeatures):
             "lookup.tests.LookupQueryingTests.test_filter_lookup_lhs",
             # Subquery not supported.
             "annotations.tests.NonAggregateAnnotationTestCase.test_empty_queryset_annotation",
+            "db_functions.datetime.test_extract_trunc.DateFunctionTests.test_extract_outerref",
+            "db_functions.datetime.test_extract_trunc.DateFunctionTests.test_trunc_subquery_with_parameters",
             "lookup.tests.LookupQueryingTests.test_filter_subquery_lhs",
             # ExpressionWrapper not supported.
             "annotations.tests.NonAggregateAnnotationTestCase.test_combined_expression_annotation_with_aggregation",
@@ -206,6 +201,11 @@ class DatabaseFeatures(BaseDatabaseFeatures):
             "annotations.tests.AliasTests.test_order_by_alias",
             "annotations.tests.NonAggregateAnnotationTestCase.test_order_by_aggregate",
             "annotations.tests.NonAggregateAnnotationTestCase.test_order_by_annotation",
+            # annotate().filter().count() gives incorrect results.
+            "db_functions.datetime.test_extract_trunc.DateFunctionTests.test_extract_year_exact_lookup",
+            # Year lookup + lt/gt crashes: 'dict' object has no attribute 'startswith'
+            "db_functions.datetime.test_extract_trunc.DateFunctionTests.test_extract_year_greaterthan_lookup",
+            "db_functions.datetime.test_extract_trunc.DateFunctionTests.test_extract_year_lessthan_lookup",
         },
         "Count doesn't work in QuerySet.annotate()": {
             "annotations.tests.AliasTests.test_alias_annotate_with_aggregation",
@@ -312,9 +312,25 @@ class DatabaseFeatures(BaseDatabaseFeatures):
             "db_functions.math.test_sqrt.SqrtTests.test_transform",
             "db_functions.math.test_tan.TanTests.test_transform",
         },
-        "MongoDB does not support Sign.": {
+        "MongoDB does not support this database function.": {
+            "db_functions.datetime.test_now.NowTests",
             "db_functions.math.test_sign.SignTests",
         },
+        "ExtractQuarter database function not supported.": {
+            "db_functions.datetime.test_extract_trunc.DateFunctionTests.test_extract_quarter_func",
+            "db_functions.datetime.test_extract_trunc.DateFunctionTests.test_extract_quarter_func_boundaries",
+        },
+        "TruncDate database function not supported.": {
+            "db_functions.datetime.test_extract_trunc.DateFunctionTests.test_trunc_date_func",
+            "db_functions.datetime.test_extract_trunc.DateFunctionTests.test_trunc_date_none",
+            "db_functions.datetime.test_extract_trunc.DateFunctionTests.test_trunc_lookup_name_sql_injection",
+            "timezones.tests.NewDatabaseTests.test_query_convert_timezones",
+        },
+        "TruncTime database function not supported.": {
+            "db_functions.datetime.test_extract_trunc.DateFunctionTests.test_trunc_time_comparison",
+            "db_functions.datetime.test_extract_trunc.DateFunctionTests.test_trunc_time_func",
+            "db_functions.datetime.test_extract_trunc.DateFunctionTests.test_trunc_time_none",
+        },
         "MongoDB can't annotate ($project) a function like PI().": {
             "db_functions.math.test_pi.PiTests.test",
         },

django_mongodb/functions.py

@@ -1,6 +1,19 @@
 from django.db import NotSupportedError
 from django.db.models.expressions import Func
-from django.db.models.functions.datetime import Extract
+from django.db.models.functions.datetime import (
+    Extract,
+    ExtractDay,
+    ExtractHour,
+    ExtractIsoWeekDay,
+    ExtractIsoYear,
+    ExtractMinute,
+    ExtractMonth,
+    ExtractSecond,
+    ExtractWeek,
+    ExtractWeekDay,
+    ExtractYear,
+    TruncBase,
+)
 from django.db.models.functions.math import Ceil, Cot, Degrees, Log, Power, Radians, Random, Round
 from django.db.models.functions.text import Upper
 
@@ -14,6 +27,18 @@
     Random: "rand",
     Upper: "toUpper",
 }
+EXTRACT_OPERATORS = {
+    ExtractDay.lookup_name: "dayOfMonth",
+    ExtractHour.lookup_name: "hour",
+    ExtractIsoWeekDay.lookup_name: "isoDayOfWeek",
+    ExtractIsoYear.lookup_name: "isoWeekYear",
+    ExtractMinute.lookup_name: "minute",
+    ExtractMonth.lookup_name: "month",
+    ExtractSecond.lookup_name: "second",
+    ExtractWeek.lookup_name: "isoWeek",
+    ExtractWeekDay.lookup_name: "dayOfWeek",
+    ExtractYear.lookup_name: "year",
+}
 
 
 def cot(self, compiler, connection):
@@ -23,15 +48,12 @@ def cot(self, compiler, connection):
 
 def extract(self, compiler, connection):
     lhs_mql = process_lhs(self, compiler, connection)
-    if self.lookup_name == "week":
-        operator = "$week"
-    elif self.lookup_name == "month":
-        operator = "$month"
-    elif self.lookup_name == "year":
-        operator = "$year"
-    else:
+    operator = EXTRACT_OPERATORS.get(self.lookup_name)
+    if operator is None:
         raise NotSupportedError("%s is not supported." % self.__class__.__name__)
-    return {operator: lhs_mql}
+    if timezone := self.get_tzname():
+        lhs_mql = {"date": lhs_mql, "timezone": timezone}
+    return {f"${operator}": lhs_mql}
 
 
 def func(self, compiler, connection):
@@ -53,9 +75,18 @@ def round_(self, compiler, connection):
     return {"$round": [expr.as_mql(compiler, connection) for expr in self.get_source_expressions()]}
 
 
+def trunc(self, compiler, connection):
+    lhs_mql = process_lhs(self, compiler, connection)
+    lhs_mql = {"date": lhs_mql, "unit": self.kind, "startOfWeek": "mon"}
+    if timezone := self.get_tzname():
+        lhs_mql["timezone"] = timezone
+    return {"$dateTrunc": lhs_mql}
+
+
 def register_functions():
     Cot.as_mql_agg = cot
-    Extract.as_mql = extract
+    Extract.as_mql = Extract.as_mql_agg = extract
     Func.as_mql_agg = func
     Log.as_mql_agg = log
     Round.as_mql_agg = round_
+    TruncBase.as_mql = TruncBase.as_mql_agg = trunc

django_mongodb/operations.py

@@ -5,6 +5,7 @@
 from django.conf import settings
 from django.db.backends.base.operations import BaseDatabaseOperations
 from django.utils import timezone
+from django.utils.regex_helper import _lazy_re_compile
 
 
 class DatabaseOperations(BaseDatabaseOperations):
@@ -133,3 +134,43 @@ def _prep_lookup_value(self, value, field, field_kind, lookup):
         if field_kind == "DecimalField":
             value = self.adapt_decimalfield_value(value, field.max_digits, field.decimal_places)
         return value
+
+    """Django uses these methods to generate SQL queries before it generates MQL queries."""
+
+    # EXTRACT format cannot be passed in parameters.
+    _extract_format_re = _lazy_re_compile(r"[A-Z_]+")
+
+    def date_extract_sql(self, lookup_type, sql, params):
+        if lookup_type == "week_day":
+            # For consistency across backends, we return Sunday=1, Saturday=7.
+            return f"EXTRACT(DOW FROM {sql}) + 1", params
+        if lookup_type == "iso_week_day":
+            return f"EXTRACT(ISODOW FROM {sql})", params
+        if lookup_type == "iso_year":
+            return f"EXTRACT(ISOYEAR FROM {sql})", params
+
+        lookup_type = lookup_type.upper()
+        if not self._extract_format_re.fullmatch(lookup_type):
+            raise ValueError(f"Invalid lookup type: {lookup_type!r}")
+        return f"EXTRACT({lookup_type} FROM {sql})", params
+
+    def datetime_extract_sql(self, lookup_type, sql, params, tzname):
+        if lookup_type == "second":
+            # Truncate fractional seconds.
+            return f"EXTRACT(SECOND FROM DATE_TRUNC(%s, {sql}))", ("second", *params)
+        return self.date_extract_sql(lookup_type, sql, params)
+
+    def datetime_trunc_sql(self, lookup_type, sql, params, tzname):
+        return f"DATE_TRUNC(%s, {sql})", (lookup_type, *params)
+
+    def date_trunc_sql(self, lookup_type, sql, params, tzname=None):
+        return f"DATE_TRUNC(%s, {sql})", (lookup_type, *params)
+
+    def datetime_cast_date_sql(self, sql, params, tzname):
+        return f"({sql})::date", params
+
+    def datetime_cast_time_sql(self, sql, params, tzname):
+        return f"({sql})::time", params
+
+    def time_trunc_sql(self, lookup_type, sql, params, tzname=None):
+        return f"DATE_TRUNC(%s, {sql})::time", (lookup_type, *params)