Add queries key to encrypted fields map

Author: aclark4life

django_mongodb_backend/fields/encryption.py

@@ -3,3 +3,8 @@
 
 class EncryptedCharField(models.CharField):
     encrypted = True
+    queries = []
+
+    def __init__(self, *args, **kwargs):
+        self.queries = kwargs.pop("queries", [])
+        super().__init__(*args, **kwargs)

django_mongodb_backend/schema.py

@@ -448,7 +448,17 @@ def _get_encrypted_fields_map(self, model):
         fields = model._meta.fields
         return {
             "fields": [
-                {"path": field.name, "bsonType": field.db_type(conn)}
+                {
+                    "path": field.name,
+                    "bsonType": field.db_type(conn),
+                    # Specify queries in the field definition as a list of query
+                    # types e.g. queries=["equality", "range"]
+                    **(
+                        {"queries": [{"queryType": query} for query in field.queries]}
+                        if hasattr(field, "queries") and field.queries
+                        else {}
+                    ),
+                }
                 for field in fields
                 if getattr(field, "encrypted", False)
             ]

tests/encryption_/models.py

@@ -6,7 +6,8 @@
 
 class Person(EncryptedModel):
     name = models.CharField("name", max_length=100)
-    ssn = EncryptedCharField("ssn", max_length=11)
+    ssn = EncryptedCharField("ssn", max_length=11, queries=["equality"])
+    ssn2 = EncryptedCharField("ssn", max_length=11, queries=["equality"])
 
     def __str__(self):
         return self.name