INTPYTHON-527 Add Queryable Encryption support

Author: aclark4life

.evergreen/config.yml

@@ -90,6 +90,28 @@ buildvariants:
     tasks:
       - name: run-tests
 
+  - name: tests-7-noauth-nossl
+    display_name: Run Tests 7.0 NoAuth NoSSL
+    run_on: rhel87-small
+    expansions:
+      MONGODB_VERSION: "7.0"
+      TOPOLOGY: server
+      AUTH: "noauth"
+      SSL: "nossl"
+    tasks:
+      - name: run-tests
+
+  - name: tests-7-auth-ssl
+    display_name: Run Tests 7.0 Auth SSL
+    run_on: rhel87-small
+    expansions:
+      MONGODB_VERSION: "7.0"
+      TOPOLOGY: server
+      AUTH: "auth"
+      SSL: "ssl"
+    tasks:
+      - name: run-tests
+
   - name: tests-8-noauth-nossl
     display_name: Run Tests 8.0 NoAuth NoSSL
     run_on: rhel87-small

django_mongodb_backend/__init__.py

@@ -14,6 +14,7 @@
 from .indexes import register_indexes  # noqa: E402
 from .lookups import register_lookups  # noqa: E402
 from .query import register_nodes  # noqa: E402
+from .routers import register_routers  # noqa: E402
 
 __all__ = ["parse_uri"]
 
@@ -25,3 +26,4 @@
 register_indexes()
 register_lookups()
 register_nodes()
+register_routers()

django_mongodb_backend/base.py

@@ -229,4 +229,7 @@ def cursor(self):
 
     def get_database_version(self):
         """Return a tuple of the database's version."""
-        return tuple(self.connection.server_info()["versionArray"])
+        # Avoid using PyMongo to check the database version or require
+        # pymongocrypt>=1.14.2 which will contain a fix for the `buildInfo`
+        # command. https://jira.mongodb.org/browse/PYTHON-5429
+        return tuple(self.connection.admin.command("buildInfo")["versionArray"])

django_mongodb_backend/encryption.py

@@ -0,0 +1,24 @@
+# Queryable Encryption helper classes
+
+
+class EqualityQuery(dict):
+    def __init__(self, *, contention=None):
+        super().__init__(queryType="equality")
+        if contention is not None:
+            self["contention"] = contention
+
+
+class RangeQuery(dict):
+    def __init__(
+        self, *, contention=None, max=None, min=None, precision=None, sparsity=None, trimFactor=None
+    ):
+        super().__init__(queryType="range")
+        options = {
+            "contention": contention,
+            "max": max,
+            "min": min,
+            "precision": precision,
+            "sparsity": sparsity,
+            "trimFactor": trimFactor,
+        }
+        self.update({k: v for k, v in options.items() if v is not None})

django_mongodb_backend/features.py

@@ -569,9 +569,17 @@ def django_test_expected_failures(self):
         },
     }
 
+    @cached_property
+    def mongodb_version(self):
+        return self.connection.get_database_version()  # e.g., (6, 3, 0)
+
     @cached_property
     def is_mongodb_6_3(self):
-        return self.connection.get_database_version() >= (6, 3)
+        return self.mongodb_version >= (6, 3)
+
+    @cached_property
+    def is_mongodb_7_0(self):
+        return self.mongodb_version >= (7, 0)
 
     @cached_property
     def supports_atlas_search(self):
@@ -601,3 +609,18 @@ def _supports_transactions(self):
         hello = client.command("hello")
         # a replica set or a sharded cluster
         return "setName" in hello or hello.get("msg") == "isdbgrid"
+
+    @cached_property
+    def supports_queryable_encryption(self):
+        """
+        Queryable Encryption requires a MongoDB 7.0 or later replica set or sharded
+        cluster, as well as MonogDB Atlas or Enterprise.
+        """
+        self.connection.ensure_connection()
+        build_info = self.connection.connection.admin.command("buildInfo")
+        is_enterprise = "enterprise" in build_info.get("modules")
+        return (
+            (is_enterprise or self.supports_atlas_search)
+            and self._supports_transactions
+            and self.is_mongodb_7_0
+        )

django_mongodb_backend/fields/__init__.py

@@ -3,6 +3,28 @@
 from .duration import register_duration_field
 from .embedded_model import EmbeddedModelField
 from .embedded_model_array import EmbeddedModelArrayField
+from .encryption import (
+    EncryptedBigIntegerField,
+    EncryptedBinaryField,
+    EncryptedBooleanField,
+    EncryptedCharField,
+    EncryptedDateField,
+    EncryptedDateTimeField,
+    EncryptedDecimalField,
+    EncryptedEmailField,
+    EncryptedFieldMixin,
+    EncryptedFloatField,
+    EncryptedGenericIPAddressField,
+    EncryptedIntegerField,
+    EncryptedPositiveBigIntegerField,
+    EncryptedPositiveIntegerField,
+    EncryptedPositiveSmallIntegerField,
+    EncryptedSmallIntegerField,
+    EncryptedTextField,
+    EncryptedTimeField,
+    EncryptedURLField,
+    has_encrypted_fields,
+)
 from .json import register_json_field
 from .objectid import ObjectIdField
 from .polymorphic_embedded_model import PolymorphicEmbeddedModelField
@@ -12,10 +34,31 @@
     "ArrayField",
     "EmbeddedModelArrayField",
     "EmbeddedModelField",
+    "EncryptedBigIntegerField",
+    "EncryptedBinaryField",
+    "EncryptedBooleanField",
+    "EncryptedCharField",
+    "EncryptedDateField",
+    "EncryptedDateTimeField",
+    "EncryptedDecimalField",
+    "EncryptedEmailField",
+    "EncryptedFieldMixin",
+    "EncryptedFloatField",
+    "EncryptedGenericIPAddressField",
+    "EncryptedIntegerField",
+    "EncryptedPositiveBigIntegerField",
+    "EncryptedPositiveIntegerField",
+    "EncryptedPositiveSmallIntegerField",
+    "EncryptedSmallIntegerField",
+    "EncryptedTextField",
+    "EncryptedTimeField",
+    "EncryptedURLField",
     "ObjectIdAutoField",
     "ObjectIdField",
     "PolymorphicEmbeddedModelArrayField",
     "PolymorphicEmbeddedModelField",
+    "has_encrypted_fields",
+    "register_fields",
     "register_fields",
 ]
 

django_mongodb_backend/fields/encryption.py

@@ -0,0 +1,99 @@
+from django.db import models
+
+
+def has_encrypted_fields(model):
+    return any(getattr(field, "encrypted", False) for field in model._meta.fields)
+
+
+class EncryptedFieldMixin(models.Field):
+    encrypted = True
+
+    def __init__(self, *args, queries=None, **kwargs):
+        self.queries = queries
+        super().__init__(*args, **kwargs)
+
+    def deconstruct(self):
+        name, path, args, kwargs = super().deconstruct()
+
+        if self.queries is not None:
+            kwargs["queries"] = self.queries
+
+        if path.startswith("django_mongodb_backend.fields.encrypted_model"):
+            path = path.replace(
+                "django_mongodb_backend.fields.encrypted_model",
+                "django_mongodb_backend.fields",
+            )
+
+        return name, path, args, kwargs
+
+
+class EncryptedBigIntegerField(EncryptedFieldMixin, models.BigIntegerField):
+    pass
+
+
+class EncryptedBinaryField(EncryptedFieldMixin, models.BinaryField):
+    pass
+
+
+class EncryptedBooleanField(EncryptedFieldMixin, models.BooleanField):
+    pass
+
+
+class EncryptedCharField(EncryptedFieldMixin, models.CharField):
+    pass
+
+
+class EncryptedDateField(EncryptedFieldMixin, models.DateField):
+    pass
+
+
+class EncryptedDateTimeField(EncryptedFieldMixin, models.DateTimeField):
+    pass
+
+
+class EncryptedDecimalField(EncryptedFieldMixin, models.DecimalField):
+    pass
+
+
+class EncryptedEmailField(EncryptedFieldMixin, models.EmailField):
+    pass
+
+
+class EncryptedFloatField(EncryptedFieldMixin, models.FloatField):
+    pass
+
+
+class EncryptedGenericIPAddressField(EncryptedFieldMixin, models.GenericIPAddressField):
+    pass
+
+
+class EncryptedIntegerField(EncryptedFieldMixin, models.IntegerField):
+    pass
+
+
+class EncryptedPositiveBigIntegerField(EncryptedFieldMixin, models.PositiveBigIntegerField):
+    pass
+
+
+class EncryptedPositiveIntegerField(EncryptedFieldMixin, models.PositiveIntegerField):
+    pass
+
+
+class EncryptedPositiveSmallIntegerField(EncryptedFieldMixin, models.PositiveSmallIntegerField):
+    pass
+
+
+class EncryptedSmallIntegerField(EncryptedFieldMixin, models.SmallIntegerField):
+    pass
+
+
+class EncryptedTimeField(EncryptedFieldMixin, models.TimeField):
+    pass
+
+
+class EncryptedTextField(EncryptedFieldMixin, models.TextField):
+    pass
+
+
+class EncryptedURLField(EncryptedFieldMixin, models.URLField):
+    pass

django_mongodb_backend/management/commands/get_encrypted_fields_map.py

@@ -0,0 +1,59 @@
+from bson import json_util
+from django.apps import apps
+from django.core.management.base import BaseCommand
+from django.db import DEFAULT_DB_ALIAS, connections, router
+from pymongo.encryption import ClientEncryption
+
+
+class Command(BaseCommand):
+    help = "Generate a `schema_map` of encrypted fields for all encrypted"
+    " models in the database for use with `AutoEncryptionOpts` in"
+    " client configuration."
+
+    def add_arguments(self, parser):
+        parser.add_argument(
+            "--database",
+            default=DEFAULT_DB_ALIAS,
+            help="Specify the database to use for generating the encrypted"
+            "fields map. Defaults to the 'default' database.",
+        )
+        parser.add_argument(
+            "--kms-provider",
+            default="local",
+            help="Specify the KMS provider to use for encryption. Defaults to 'local'.",
+        )
+
+    def handle(self, *args, **options):
+        db = options["database"]
+        kms_provider = options["kms_provider"]
+        connection = connections[db]
+        schema_map = json_util.dumps(
+            self.get_encrypted_fields_map(connection, kms_provider), indent=2
+        )
+        self.stdout.write(schema_map)
+
+    def get_client_encryption(self, connection):
+        client = connection.connection
+        options = client._options.auto_encryption_opts
+        key_vault_namespace = options._key_vault_namespace
+        kms_providers = options._kms_providers
+        return ClientEncryption(kms_providers, key_vault_namespace, client, client.codec_options)
+
+    def get_encrypted_fields_map(self, connection, kms_provider):
+        schema_map = {}
+        for app_config in apps.get_app_configs():
+            for model in router.get_migratable_models(
+                app_config, connection.settings_dict["NAME"], include_auto_created=False
+            ):
+                if getattr(model, "encrypted", False):
+                    fields = connection.schema_editor()._get_encrypted_fields_map(model)
+                    ce = self.get_client_encryption(connection)
+                    master_key = connection.settings_dict.get("KMS_CREDENTIALS").get(kms_provider)
+                    for field in fields["fields"]:
+                        data_key = ce.create_data_key(
+                            kms_provider=kms_provider,
+                            master_key=master_key,
+                        )
+                        field["keyId"] = data_key
+                    schema_map[model._meta.db_table] = fields
+        return schema_map

django_mongodb_backend/routers.py

@@ -1,6 +1,8 @@
 from django.apps import apps
+from django.core.exceptions import ImproperlyConfigured
+from django.db.utils import ConnectionRouter
 
-from django_mongodb_backend.models import EmbeddedModel
+from .fields import has_encrypted_fields
 
 
 class MongoRouter:
@@ -9,10 +11,32 @@ def allow_migrate(self, db, app_label, model_name=None, **hints):
         EmbeddedModels don't have their own collection and must be ignored by
         dumpdata.
         """
+        from django_mongodb_backend.models import EmbeddedModel  # noqa: PLC0415
+
         if not model_name:
             return None
         try:
             model = apps.get_model(app_label, model_name)
         except LookupError:
             return None
+
         return False if issubclass(model, EmbeddedModel) else None
+
+
+def kms_provider(self, model, *args, **kwargs):
+    for router in self.routers:
+        func = getattr(router, "kms_provider", None)
+        if func and callable(func):
+            result = func(model, *args, **kwargs)
+            if result is not None:
+                return result
+    if has_encrypted_fields(model):
+        raise ImproperlyConfigured("No kms_provider found in database router.")
+    return None
+
+
+def register_routers():
+    """
+    Patch the ConnectionRouter to use the custom kms_provider method.
+    """
+    ConnectionRouter.kms_provider = kms_provider