INTPYTHON-527 Add Queryable Encryption support

Author: aclark4life

.github/workflows/encrypted_settings.py

@@ -0,0 +1,40 @@
+# Settings for django_mongodb_backend/tests when encryption is supported.
+import os
+
+from mongodb_settings import *  # noqa: F403
+from pymongo.encryption import AutoEncryptionOpts
+
+DATABASES["encrypted"] = {  # noqa: F405
+    "ENGINE": "django_mongodb_backend",
+    "NAME": "djangotests_encrypted",
+    "OPTIONS": {
+        "auto_encryption_opts": AutoEncryptionOpts(
+            key_vault_namespace="djangotests_encrypted.__keyVault",
+            kms_providers={"local": {"key": os.urandom(96)}},
+        ),
+        "directConnection": True,
+    },
+    "KMS_CREDENTIALS": {},
+}
+
+
+class EncryptedRouter:
+    def db_for_read(self, model, **hints):
+        if model._meta.app_label == "encryption_":
+            return "encrypted"
+        return None
+
+    db_for_write = db_for_read
+
+    def allow_migrate(self, db, app_label, model_name=None, **hints):
+        # The encryption_ app's models are only created in the encrypted
+        # database.
+        if app_label == "encryption_":
+            return db == "encrypted"
+        # Don't create other app's models in the encrypted database.
+        if db == "encrypted":
+            return False
+        return None
+
+
+DATABASE_ROUTERS.append(EncryptedRouter())  # noqa: F405

.github/workflows/mongodb_settings.py

@@ -1,4 +1,5 @@
-# Settings for django_mongodb_backend/tests.
+# Settings for django_mongodb_backend/tests when encryption isn't supported.
 from django_settings import *  # noqa: F403
 
+DATABASES["encrypted"] = {}  # noqa: F405
 DATABASE_ROUTERS = ["django_mongodb_backend.routers.MongoRouter"]

.github/workflows/test-python-atlas.yml

@@ -28,7 +28,7 @@ jobs:
       - name: install django-mongodb-backend
         run: |
           pip3 install --upgrade pip
-          pip3 install -e .
+          pip3 install -e .[encryption]
       - name: Checkout Django
         uses: actions/checkout@v5
         with:
@@ -51,8 +51,30 @@ jobs:
         run: cp .github/workflows/runtests.py django_repo/tests/runtests_.py
       - name: Start local Atlas
         working-directory: .
-        run: bash .github/workflows/start_local_atlas.sh mongodb/mongodb-atlas-local:7
+        run: bash .github/workflows/start_local_atlas.sh mongodb/mongodb-atlas-local:8.0.15
+      - name: Install mongosh
+        run: |
+          wget -q https://downloads.mongodb.com/compass/mongosh-2.2.10-linux-x64.tgz
+          tar -xzf mongosh-*-linux-x64.tgz
+          sudo cp mongosh-*-linux-x64/bin/mongosh /usr/local/bin/
+          mongosh --version
+      - name: Install mongocryptd from Enterprise tarball
+        run: |
+          curl -sSL -o mongodb-enterprise.tgz "https://downloads.mongodb.com/linux/mongodb-linux-x86_64-enterprise-ubuntu2204-8.0.15.tgz"
+          tar -xzf mongodb-enterprise.tgz
+          sudo cp mongodb-linux-x86_64-enterprise-ubuntu2204-8.0.15/bin/mongocryptd /usr/local/bin/
+      - name: Start mongocryptd
+        run: |
+          nohup mongocryptd --logpath=/tmp/mongocryptd.log &
+      - name: Verify MongoDB installation
+        run: |
+          mongosh --eval 'db.runCommand({ connectionStatus: 1 })'
+      - name: Verify mongocryptd is running
+        run: |
+          pgrep mongocryptd
       - name: Run tests
         run: python3 django_repo/tests/runtests_.py
     permissions:
         contents: read
+    env:
+      DJANGO_SETTINGS_MODULE: "encrypted_settings"

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

@@ -11,6 +11,7 @@
 from django.utils.functional import cached_property
 from pymongo.collection import Collection
 from pymongo.driver_info import DriverInfo
+from pymongo.encryption import ClientEncryption
 from pymongo.mongo_client import MongoClient
 from pymongo.uri_parser import parse_uri
 
@@ -241,6 +242,16 @@ def get_database(self):
             return OperationDebugWrapper(self)
         return self.database
 
+    @cached_property
+    def client_encryption(self):
+        auto_encryption_opts = self.connection._options.auto_encryption_opts
+        return ClientEncryption(
+            auto_encryption_opts._kms_providers,
+            auto_encryption_opts._key_vault_namespace,
+            self.connection,
+            self.connection.codec_options,
+        )
+
     @cached_property
     def database(self):
         """Connect to the database the first time it's accessed."""

django_mongodb_backend/creation.py

@@ -1,12 +1,20 @@
 from django.conf import settings
-from django.db.backends.base.creation import BaseDatabaseCreation
+from django.db.backends.base.creation import TEST_DATABASE_PREFIX, BaseDatabaseCreation
 
 
 class DatabaseCreation(BaseDatabaseCreation):
     def _execute_create_test_db(self, cursor, parameters, keepdb=False):
         # Close the connection (which may point to the non-test database) so
         # that a new connection to the test database can be established later.
         self.connection.close_pool()
+        # Use a test _key_vault_namespace. This assumes the key vault database
+        # is the same as the encrypted database so that _destroy_test_db() can
+        # reset the collection by dropping it.
+        opts = self.connection.settings_dict["OPTIONS"].get("auto_encryption_opts")
+        if opts:
+            self.connection.settings_dict["OPTIONS"][
+                "auto_encryption_opts"
+            ]._key_vault_namespace = TEST_DATABASE_PREFIX + opts._key_vault_namespace
         if not keepdb:
             self._destroy_test_db(parameters["dbname"], verbosity=0)
 
@@ -24,3 +32,9 @@ def destroy_test_db(self, old_database_name=None, verbosity=1, keepdb=False, suf
         super().destroy_test_db(old_database_name, verbosity, keepdb, suffix)
         # Close the connection to the test database.
         self.connection.close_pool()
+        # Restore the original _key_vault_namespace.
+        opts = self.connection.settings_dict["OPTIONS"].get("auto_encryption_opts")
+        if opts:
+            self.connection.settings_dict["OPTIONS"][
+                "auto_encryption_opts"
+            ]._key_vault_namespace = opts._key_vault_namespace[len(TEST_DATABASE_PREFIX) :]

django_mongodb_backend/features.py

@@ -588,9 +588,21 @@ def django_test_skips(self):
         skips.update(self._django_test_skips)
         return skips
 
+    @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 is_mongodb_8_0(self):
+        return self.mongodb_version >= (8, 0)
 
     @cached_property
     def supports_atlas_search(self):
@@ -620,3 +632,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 8.0 or later replica set or sharded
+        cluster, as well as MongoDB 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_8_0
+        )

django_mongodb_backend/fields/__init__.py

@@ -3,6 +3,33 @@
 from .duration import register_duration_field
 from .embedded_model import EmbeddedModelField
 from .embedded_model_array import EmbeddedModelArrayField
+from .encryption import (
+    EncryptedArrayField,
+    EncryptedBigIntegerField,
+    EncryptedBinaryField,
+    EncryptedBooleanField,
+    EncryptedCharField,
+    EncryptedDateField,
+    EncryptedDateTimeField,
+    EncryptedDecimalField,
+    EncryptedDurationField,
+    EncryptedEmailField,
+    EncryptedEmbeddedModelArrayField,
+    EncryptedEmbeddedModelField,
+    EncryptedFieldMixin,
+    EncryptedFloatField,
+    EncryptedGenericIPAddressField,
+    EncryptedIntegerField,
+    EncryptedObjectIdField,
+    EncryptedPositiveBigIntegerField,
+    EncryptedPositiveIntegerField,
+    EncryptedPositiveSmallIntegerField,
+    EncryptedSmallIntegerField,
+    EncryptedTextField,
+    EncryptedTimeField,
+    EncryptedURLField,
+    EncryptedUUIDField,
+)
 from .json import register_json_field
 from .objectid import ObjectIdField
 from .polymorphic_embedded_model import PolymorphicEmbeddedModelField
@@ -12,6 +39,31 @@
     "ArrayField",
     "EmbeddedModelArrayField",
     "EmbeddedModelField",
+    "EncryptedArrayField",
+    "EncryptedBigIntegerField",
+    "EncryptedBinaryField",
+    "EncryptedBooleanField",
+    "EncryptedCharField",
+    "EncryptedDateField",
+    "EncryptedDateTimeField",
+    "EncryptedDecimalField",
+    "EncryptedDurationField",
+    "EncryptedEmailField",
+    "EncryptedEmbeddedModelArrayField",
+    "EncryptedEmbeddedModelField",
+    "EncryptedFieldMixin",
+    "EncryptedFloatField",
+    "EncryptedGenericIPAddressField",
+    "EncryptedIntegerField",
+    "EncryptedObjectIdField",
+    "EncryptedPositiveBigIntegerField",
+    "EncryptedPositiveIntegerField",
+    "EncryptedPositiveSmallIntegerField",
+    "EncryptedSmallIntegerField",
+    "EncryptedTextField",
+    "EncryptedTimeField",
+    "EncryptedURLField",
+    "EncryptedUUIDField",
     "ObjectIdAutoField",
     "ObjectIdField",
     "PolymorphicEmbeddedModelArrayField",

django_mongodb_backend/fields/encryption.py

@@ -0,0 +1,130 @@
+from django.db import models
+
+from django_mongodb_backend.fields import ArrayField, EmbeddedModelArrayField, EmbeddedModelField
+from django_mongodb_backend.fields.objectid import ObjectIdField
+
+
+class EncryptedFieldMixin:
+    encrypted = True
+
+    def __init__(self, *args, queries=None, db_index=False, null=False, unique=False, **kwargs):
+        if db_index:
+            raise ValueError("'db_index=True' is not supported on encrypted fields.")
+        if null:
+            raise ValueError("'null=True' is not supported on encrypted fields.")
+        if unique:
+            raise ValueError("'unique=True' is not supported on encrypted fields.")
+        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.encryption"):
+            path = path.replace(
+                "django_mongodb_backend.fields.encryption",
+                "django_mongodb_backend.fields",
+            )
+
+        return name, path, args, kwargs
+
+
+# Django fields
+class EncryptedBinaryField(EncryptedFieldMixin, models.BinaryField):
+    pass
+
+
+class EncryptedBigIntegerField(EncryptedFieldMixin, models.BigIntegerField):
+    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 EncryptedDurationField(EncryptedFieldMixin, models.DurationField):
+    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 EncryptedTextField(EncryptedFieldMixin, models.TextField):
+    pass
+
+
+class EncryptedTimeField(EncryptedFieldMixin, models.TimeField):
+    pass
+
+
+class EncryptedURLField(EncryptedFieldMixin, models.URLField):
+    pass
+
+
+class EncryptedUUIDField(EncryptedFieldMixin, models.UUIDField):
+    pass
+
+
+# MongoDB fields
+class EncryptedArrayField(EncryptedFieldMixin, ArrayField):
+    pass
+
+
+class EncryptedEmbeddedModelArrayField(EncryptedFieldMixin, EmbeddedModelArrayField):
+    pass
+
+
+class EncryptedEmbeddedModelField(EncryptedFieldMixin, EmbeddedModelField):
+    pass
+
+
+class EncryptedObjectIdField(EncryptedFieldMixin, ObjectIdField):
+    pass