implement date lookups

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

django_mongodb/features.py

@@ -7,6 +7,8 @@ class DatabaseFeatures(BaseDatabaseFeatures):
     supports_ignore_conflicts = False
     # Not implemented: https://github.com/mongodb-labs/django-mongodb/issues/8
     supports_json_field = False
+    # BSON Date type doesn't support microsecond precision
+    supports_microsecond_precision = False
     # MongoDB stores datetimes in UTC.
     supports_timezones = False
     # Not implemented: https://github.com/mongodb-labs/django-mongodb/issues/7
@@ -17,19 +19,10 @@ 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",
-        "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",
+        # TruncDate isn't implemented.
         "timezones.tests.NewDatabaseTests.test_query_convert_timezones",
-        "timezones.tests.NewDatabaseTests.test_query_datetime_lookups",
-        "timezones.tests.NewDatabaseTests.test_query_datetime_lookups_in_other_timezone",
+        # Query for chained looksup not generated correctly.
+        "lookup.tests.LookupTests.test_chain_date_time_lookups",
         # 'NulledTransform' object has no attribute 'as_mql'.
         "lookup.tests.LookupTests.test_exact_none_transform",
         # "Save with update_fields did not affect any rows."
@@ -60,6 +53,8 @@ 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",
+        # Wrong result truncating in other time zone.
+        "db_functions.datetime.test_extract_trunc.DateFunctionWithTimeZoneTests.test_trunc_timezone_applied_before_truncation",
     }
 
     django_test_skips = {
@@ -169,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",
@@ -204,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",
@@ -292,14 +294,22 @@ class DatabaseFeatures(BaseDatabaseFeatures):
             "timezones.tests.NewDatabaseTests.test_cursor_explicit_time_zone",
             "timezones.tests.NewDatabaseTests.test_raw_sql",
         },
-        "BSON Date type doesn't support microsecond precision.": {
-            "basic.tests.ModelRefreshTests.test_refresh_unsaved",
-            "basic.tests.ModelTest.test_microsecond_precision",
-            "timezones.tests.LegacyDatabaseTests.test_auto_now_and_auto_now_add",
-            "timezones.tests.LegacyDatabaseTests.test_aware_datetime_in_local_timezone_with_microsecond",
-            "timezones.tests.LegacyDatabaseTests.test_naive_datetime_with_microsecond",
-            "timezones.tests.NewDatabaseTests.test_aware_datetime_in_local_timezone_with_microsecond",
-            "timezones.tests.NewDatabaseTests.test_naive_datetime_with_microsecond",
+        "ExtractQuarter 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 (datetime to date) 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",
+        },
+        "TruncTime (datetime to time) 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",
+        },
+        "No Now() database function in MongoDB.": {
+            "db_functions.datetime.test_now.NowTests",
         },
         "Transform not supported.": {
             "db_functions.math.test_abs.AbsTests.test_transform",

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
 
@@ -9,6 +22,16 @@
 MONGO_OPERATORS = {
     Ceil: "ceil",
     Degrees: "radiansToDegrees",
+    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",
     Power: "pow",
     Radians: "degreesToRadians",
     Random: "rand",
@@ -23,14 +46,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:
-        raise NotSupportedError("%s is not supported." % self.__class__.__name__)
+    try:
+        operator = f"${MONGO_OPERATORS[self.lookup_name]}"
+    except KeyError:
+        raise NotSupportedError("%s is not supported." % self.__class__.__name__) from None
+    if timezone := self.get_tzname():
+        lhs_mql = {"date": lhs_mql, "timezone": timezone}
     return {operator: lhs_mql}
 
 
@@ -53,9 +74,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):
@@ -43,6 +44,44 @@ def adapt_timefield_value(self, value):
             raise ValueError("MongoDB backend does not support timezone-aware times.")
         return datetime.datetime.combine(datetime.datetime.min.date(), value)
 
+    # 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)
+
     def get_db_converters(self, expression):
         converters = super().get_db_converters(expression)
         internal_type = expression.output_field.get_internal_type()