add ArrayField

Author: timgraham

.pre-commit-config.yaml

@@ -44,6 +44,7 @@ repos:
  hooks:
  - id: rstcheck
  additional_dependencies: [sphinx]
+ args: ["--ignore-directives=fieldlookup,setting", "--ignore-roles=lookup,setting"]
 
 # We use the Python version instead of the original version which seems to require Docker
 # https://github.com/koalaman/shellcheck-precommit

django_mongodb_backend/features.py

@@ -80,6 +80,19 @@ class DatabaseFeatures(BaseDatabaseFeatures):
  "auth_tests.test_views.LoginTest.test_login_session_without_hash_session_key",
  # GenericRelation.value_to_string() assumes integer pk.
  "contenttypes_tests.test_fields.GenericRelationTests.test_value_to_string",
+ # contains with Exists() doesn't work:
+ # https://github.com/mongodb-labs/django-mongodb/issues/204
+ "model_fields_.test_arrayfield.QueryingTests.test_contains_subquery",
+ # overlap with values() returns no results:
+ # https://github.com/mongodb-labs/django-mongodb/issues/209
+ "model_fields_.test_arrayfield.QueryingTests.test_overlap_values",
+ # icontains doesn't work on ArrayField:
+ # Unsupported conversion from array to string in $convert
+ "model_fields_.test_arrayfield.QueryingTests.test_icontains",
+ # ArrayField's contained_by lookup crashes with Exists: "both operands "
+ # of $setIsSubset must be arrays. Second argument is of type: null"
+ # https://jira.mongodb.org/browse/SERVER-99186
+ "model_fields_.test_arrayfield.QueryingTests.test_contained_by_subquery",
  }
  # $bitAnd, #bitOr, and $bitXor are new in MongoDB 6.3.
  _django_test_expected_failures_bitwise = {

django_mongodb_backend/fields/__init__.py

@@ -1,9 +1,10 @@
+from .array import ArrayField
 from .auto import ObjectIdAutoField
 from .duration import register_duration_field
 from .json import register_json_field
 from .objectid import ObjectIdField
 
-__all__ = ["register_fields", "ObjectIdAutoField", "ObjectIdField"]
+__all__ = ["register_fields", "ArrayField", "ObjectIdAutoField", "ObjectIdField"]
 
 
 def register_fields():

django_mongodb_backend/fields/array.py

@@ -0,0 +1,350 @@
+import json
+
+from django.contrib.postgres.validators import ArrayMaxLengthValidator
+from django.core import checks, exceptions
+from django.db.models import DecimalField, Field, Func, IntegerField, Transform, Value
+from django.db.models.fields.mixins import CheckFieldDefaultMixin
+from django.db.models.lookups import Exact, FieldGetDbPrepValueMixin, In, Lookup
+from django.utils.translation import gettext_lazy as _
+
+from ..forms import SimpleArrayField
+from ..query_utils import process_lhs, process_rhs
+from ..utils import prefix_validation_error
+
+__all__ = ["ArrayField"]
+
+
+class AttributeSetter:
+ def __init__(self, name, value):
+ setattr(self, name, value)
+
+
+class ArrayField(CheckFieldDefaultMixin, Field):
+ empty_strings_allowed = False
+ default_error_messages = {
+ "item_invalid": _("Item %(nth)s in the array did not validate:"),
+ "nested_array_mismatch": _("Nested arrays must have the same length."),
+ }
+ _default_hint = ("list", "[]")
+
+ def __init__(self, base_field, size=None, **kwargs):
+ self.base_field = base_field
+ self.size = size
+ if self.size:
+ self.default_validators = [
+ *self.default_validators,
+ ArrayMaxLengthValidator(self.size),
+ ]
+ # For performance, only add a from_db_value() method if the base field
+ # implements it.
+ if hasattr(self.base_field, "from_db_value"):
+ self.from_db_value = self._from_db_value
+ super().__init__(**kwargs)
+
+ @property
+ def model(self):
+ try:
+ return self.__dict__["model"]
+ except KeyError:
+ raise AttributeError(
+ "'%s' object has no attribute 'model'" % self.__class__.__name__
+ ) from None
+
+ @model.setter
+ def model(self, model):
+ self.__dict__["model"] = model
+ self.base_field.model = model
+
+ @classmethod
+ def _choices_is_value(cls, value):
+ return isinstance(value, list | tuple) or super()._choices_is_value(value)
+
+ def check(self, **kwargs):
+ errors = super().check(**kwargs)
+ if self.base_field.remote_field:
+ errors.append(
+ checks.Error(
+ "Base field for array cannot be a related field.",
+ obj=self,
+ id="django_mongodb_backend.array.E002",
+ )
+ )
+ else:
+ base_checks = self.base_field.check()
+ if base_checks:
+ error_messages = "\n ".join(
+ f"{base_check.msg} ({base_check.id})"
+ for base_check in base_checks
+ if isinstance(base_check, checks.Error)
+ )
+ if error_messages:
+ errors.append(
+ checks.Error(
+ f"Base field for array has errors:\n {error_messages}",
+ obj=self,
+ id="django_mongodb_backend.array.E001",
+ )
+ )
+ warning_messages = "\n ".join(
+ f"{base_check.msg} ({base_check.id})"
+ for base_check in base_checks
+ if isinstance(base_check, checks.Warning)
+ )
+ if warning_messages:
+ errors.append(
+ checks.Warning(
+ f"Base field for array has warnings:\n {warning_messages}",
+ obj=self,
+ id="django_mongodb_backend.array.W004",
+ )
+ )
+ return errors
+
+ def set_attributes_from_name(self, name):
+ super().set_attributes_from_name(name)
+ self.base_field.set_attributes_from_name(name)
+
+ @property
+ def description(self):
+ return f"Array of {self.base_field.description}"
+
+ def db_type(self, connection):
+ return "array"
+
+ def get_db_prep_value(self, value, connection, prepared=False):
+ if isinstance(value, list | tuple):
+ # Workaround for https://code.djangoproject.com/ticket/35982
+ # (fixed in Django 5.2).
+ if isinstance(self.base_field, DecimalField):
+ return [self.base_field.get_db_prep_save(i, connection) for i in value]
+ return [self.base_field.get_db_prep_value(i, connection, prepared=False) for i in value]
+ return value
+
+ def deconstruct(self):
+ name, path, args, kwargs = super().deconstruct()
+ if path == "django_mongodb_backend.fields.array.ArrayField":
+ path = "django_mongodb_backend.fields.ArrayField"
+ kwargs.update(
+ {
+ "base_field": self.base_field.clone(),
+ "size": self.size,
+ }
+ )
+ return name, path, args, kwargs
+
+ def to_python(self, value):
+ if isinstance(value, str):
+ # Assume value is being deserialized.
+ vals = json.loads(value)
+ value = [self.base_field.to_python(val) for val in vals]
+ return value
+
+ def _from_db_value(self, value, expression, connection):
+ if value is None:
+ return value
+ return [self.base_field.from_db_value(item, expression, connection) for item in value]
+
+ def value_to_string(self, obj):
+ values = []
+ vals = self.value_from_object(obj)
+ base_field = self.base_field
+
+ for val in vals:
+ if val is None:
+ values.append(None)
+ else:
+ obj = AttributeSetter(base_field.attname, val)
+ values.append(base_field.value_to_string(obj))
+ return json.dumps(values)
+
+ def get_transform(self, name):
+ transform = super().get_transform(name)
+ if transform:
+ return transform
+ if "_" not in name:
+ try:
+ index = int(name)
+ except ValueError:
+ pass
+ else:
+ return IndexTransformFactory(index, self.base_field)
+ try:
+ start, end = name.split("_")
+ start = int(start)
+ end = int(end)
+ except ValueError:
+ pass
+ else:
+ return SliceTransformFactory(start, end)
+
+ def validate(self, value, model_instance):
+ super().validate(value, model_instance)
+ for index, part in enumerate(value):
+ try:
+ self.base_field.validate(part, model_instance)
+ except exceptions.ValidationError as error:
+ raise prefix_validation_error(
+ error,
+ prefix=self.error_messages["item_invalid"],
+ code="item_invalid",
+ params={"nth": index + 1},
+ ) from None
+ if isinstance(self.base_field, ArrayField) and len({len(i) for i in value}) > 1:
+ raise exceptions.ValidationError(
+ self.error_messages["nested_array_mismatch"],
+ code="nested_array_mismatch",
+ )
+
+ def run_validators(self, value):
+ super().run_validators(value)
+ for index, part in enumerate(value):
+ try:
+ self.base_field.run_validators(part)
+ except exceptions.ValidationError as error:
+ raise prefix_validation_error(
+ error,
+ prefix=self.error_messages["item_invalid"],
+ code="item_invalid",
+ params={"nth": index + 1},
+ ) from None
+
+ def formfield(self, **kwargs):
+ return super().formfield(
+ **{
+ "form_class": SimpleArrayField,
+ "base_field": self.base_field.formfield(),
+ "max_length": self.size,
+ **kwargs,
+ }
+ )
+
+
+class Array(Func):
+ def as_mql(self, compiler, connection):
+ return [expr.as_mql(compiler, connection) for expr in self.get_source_expressions()]
+
+
+class ArrayRHSMixin:
+ def __init__(self, lhs, rhs):
+ if isinstance(rhs, tuple | list):
+ expressions = []
+ for value in rhs:
+ if not hasattr(value, "resolve_expression"):
+ field = lhs.output_field
+ value = Value(field.base_field.get_prep_value(value))
+ expressions.append(value)
+ rhs = Array(*expressions)
+ super().__init__(lhs, rhs)
+
+
+@ArrayField.register_lookup
+class ArrayContains(ArrayRHSMixin, FieldGetDbPrepValueMixin, Lookup):
+ lookup_name = "contains"
+
+ def as_mql(self, compiler, connection):
+ lhs_mql = process_lhs(self, compiler, connection)
+ value = process_rhs(self, compiler, connection)
+ return {"$and": [{"$ne": [lhs_mql, None]}, {"$setIsSubset": [value, lhs_mql]}]}
+
+
+@ArrayField.register_lookup
+class ArrayContainedBy(ArrayRHSMixin, FieldGetDbPrepValueMixin, Lookup):
+ lookup_name = "contained_by"
+
+ def as_mql(self, compiler, connection):
+ lhs_mql = process_lhs(self, compiler, connection)
+ value = process_rhs(self, compiler, connection)
+ return {
+ "$and": [
+ {"$ne": [lhs_mql, None]},
+ {"$ne": [value, None]},
+ {"$setIsSubset": [lhs_mql, value]},
+ ]
+ }
+
+
+@ArrayField.register_lookup
+class ArrayExact(ArrayRHSMixin, Exact):
+ pass
+
+
+@ArrayField.register_lookup
+class ArrayOverlap(ArrayRHSMixin, FieldGetDbPrepValueMixin, Lookup):
+ lookup_name = "overlap"
+
+ def as_mql(self, compiler, connection):
+ lhs_mql = process_lhs(self, compiler, connection)
+ value = process_rhs(self, compiler, connection)
+ return {
+ "$and": [{"$ne": [lhs_mql, None]}, {"$size": {"$setIntersection": [value, lhs_mql]}}]
+ }
+
+
+@ArrayField.register_lookup
+class ArrayLenTransform(Transform):
+ lookup_name = "len"
+ output_field = IntegerField()
+
+ def as_mql(self, compiler, connection):
+ lhs_mql = process_lhs(self, compiler, connection)
+ return {"$cond": {"if": {"$eq": [lhs_mql, None]}, "then": None, "else": {"$size": lhs_mql}}}
+
+
+@ArrayField.register_lookup
+class ArrayInLookup(In):
+ def get_prep_lookup(self):
+ values = super().get_prep_lookup()
+ if hasattr(values, "resolve_expression"):
+ return values
+ # process_rhs() expects hashable values, so convert lists to tuples.
+ prepared_values = []
+ for value in values:
+ if hasattr(value, "resolve_expression"):
+ prepared_values.append(value)
+ else:
+ prepared_values.append(tuple(value))
+ return prepared_values
+
+
+class IndexTransform(Transform):
+ def __init__(self, index, base_field, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+ self.index = index
+ self.base_field = base_field
+
+ def as_mql(self, compiler, connection):
+ lhs_mql = process_lhs(self, compiler, connection)
+ return {"$arrayElemAt": [lhs_mql, self.index]}
+
+ @property
+ def output_field(self):
+ return self.base_field
+
+
+class IndexTransformFactory:
+ def __init__(self, index, base_field):
+ self.index = index
+ self.base_field = base_field
+
+ def __call__(self, *args, **kwargs):
+ return IndexTransform(self.index, self.base_field, *args, **kwargs)
+
+
+class SliceTransform(Transform):
+ def __init__(self, start, end, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+ self.start = start
+ self.end = end
+
+ def as_mql(self, compiler, connection):
+ lhs_mql = process_lhs(self, compiler, connection)
+ return {"$slice": [lhs_mql, self.start, self.end]}
+
+
+class SliceTransformFactory:
+ def __init__(self, start, end):
+ self.start = start
+ self.end = end
+
+ def __call__(self, *args, **kwargs):
+ return SliceTransform(self.start, self.end, *args, **kwargs)