implement date lookups

Author: timgraham

django_mongodb/features.py

@@ -7,6 +7,7 @@ class DatabaseFeatures(BaseDatabaseFeatures):
     supports_ignore_conflicts = False
     # Not implemented: https://github.com/mongodb-labs/django-mongodb/issues/8
     supports_json_field = False
+    supports_microsecond_precision = False
     # MongoDB stores datetimes in UTC.
     supports_timezones = False
     # Not implemented: https://github.com/mongodb-labs/django-mongodb/issues/7
@@ -17,20 +18,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",
-        "basic.tests.ModelTest.test_year_lookup_edge_case",
-        "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."

django_mongodb/functions.py

@@ -1,21 +1,52 @@
 from django.db import NotSupportedError
-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 .query_utils import process_lhs
 
+ExtractDay.mongo_operator = "$dayOfMonth"
+ExtractHour.mongo_operator = "$hour"
+ExtractIsoWeekDay.mongo_operator = "$isoDayOfWeek"
+ExtractIsoYear.mongo_operator = "$isoWeekYear"
+ExtractMinute.mongo_operator = "$minute"
+ExtractMonth.mongo_operator = "$month"
+ExtractSecond.mongo_operator = "$second"
+ExtractWeek.mongo_operator = "$week"
+ExtractWeekDay.mongo_operator = "$dayOfWeek"
+ExtractYear.mongo_operator = "$year"
+
 
 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 = self.mongo_operator
+    except AttributeError:
+        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}
 
 
+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():
-    Extract.as_mql = extract
+    Extract.as_mql = Extract.as_mql_agg = extract
+    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):
@@ -37,6 +38,35 @@ 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 get_db_converters(self, expression):
         converters = super().get_db_converters(expression)
         internal_type = expression.output_field.get_internal_type()