Add query type classes and update test

Rather than require users to add JSON to their field definitions let's
provide a Python API via QueryTypes class that outputs
JSON with values provided in field definition.

E.g.

query_types = QueryTypes()

query_types.equality.contention = 1
queries = [query_types.equality]

Author: aclark4life

django_mongodb_backend/encryption.py

@@ -12,6 +12,58 @@
 KMS_PROVIDER = "local"  # e.g., "aws", "azure", "gcp", "kmip", or "local"
 
 
+class EqualityQuery:
+    """
+    Represents an encrypted equality query for encrypted fields in MongoDB's
+    Queryable Encryption.
+    """
+
+    def __init__(self, contention=None):
+        self.queryType = "equality"
+        self.contention = contention
+
+    def to_dict(self):
+        query_type = {"queryType": self.queryType}
+        if self.contention is not None:
+            query_type["contention"] = self.contention
+        return [query_type]
+
+
+class RangeQuery:
+    """Represents an encrypted range query configuration for encrypted fields in
+    MongoDB's Queryable Encryption.
+    """
+
+    def __init__(self, sparsity=None, precision=None, trimFactor=None):
+        self.queryType = "range"
+        self.sparsity = sparsity
+        self.precision = precision
+        self.trimFactor = trimFactor
+
+    def to_dict(self):
+        query_type = {"queryType": self.queryType}
+        if self.sparsity is not None:
+            query_type["sparsity"] = self.sparsity
+        if self.precision is not None:
+            query_type["precision"] = self.precision
+        if self.trimFactor is not None:
+            query_type["trimFactor"] = self.trimFactor
+        return query_type
+
+
+class QueryTypes:
+    """
+    Factory class for creating query type configurations for
+    MongoDB Queryable Encryption.
+    """
+
+    def equality(self, *, contention=None):
+        return EqualityQuery(contention=contention)
+
+    def range(self, *, sparsity=None, precision=None, trimFactor=None):
+        return RangeQuery(sparsity=sparsity, precision=precision, trimFactor=trimFactor)
+
+
 def get_auto_encryption_opts(
     key_vault_namespace=None, crypt_shared_lib_path=None, kms_providers=None
 ):

django_mongodb_backend/schema.py

@@ -446,15 +446,14 @@ def _create_collection(self, model):
     def _get_encrypted_fields_map(self, model):
         conn = self.connection
         fields = model._meta.fields
+
         return {
             "fields": [
                 {
                     "path": field.column,
                     "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]}
+                        {"queries": field.queries[0].to_dict()}
                         if getattr(field, "queries", None)
                         else {}
                     ),

tests/encryption_/models.py

@@ -1,12 +1,17 @@
 from django.db import models
 
+from django_mongodb_backend.encryption import QueryTypes
 from django_mongodb_backend.fields import EncryptedCharField
 from django_mongodb_backend.models import EncryptedModel
 
+# Query types for encrypted fields with optional parameters
+query_types = QueryTypes()
+queries = [query_types.equality(contention=1), query_types.range(sparsity=2, precision=3)]
+
 
 class Person(EncryptedModel):
     name = models.CharField("name", max_length=100)
-    ssn = EncryptedCharField("ssn", max_length=11, queries=["equality"])
+    ssn = EncryptedCharField("ssn", max_length=11, queries=queries)
 
     class Meta:
         required_db_features = {"supports_queryable_encryption"}

tests/encryption_/tests.py

@@ -15,7 +15,11 @@ def test_encrypted_fields_map(self):
         """ """
         expected = {
             "fields": [
-                {"path": "ssn", "bsonType": "string", "queries": [{"queryType": "equality"}]}
+                {
+                    "path": "ssn",
+                    "bsonType": "string",
+                    "queries": [{"contention": 1, "queryType": "equality"}],
+                }
             ]
         }
         with connection.schema_editor() as editor: