INTPYTHON-527 Add Queryable Encryption support

.evergreen/config.yml

@@ -55,6 +55,46 @@ functions:
         args:
           - ${DRIVERS_TOOLS}/.evergreen/teardown.sh
 
+  # Encryption-specific functions
+  "start csfle servers":
+    - command: ec2.assume_role
+      params:
+        role_arn: ${aws_test_secrets_role}
+    - command: subprocess.exec
+      params:
+        binary: bash
+        include_expansions_in_env: [
+            "AWS_SECRET_ACCESS_KEY",
+            "AWS_ACCESS_KEY_ID",
+            "AWS_SESSION_TOKEN",
+        ]
+        args:
+          - ${DRIVERS_TOOLS}/.evergreen/csfle/setup.sh
+
+  "teardown csfle":
+    - command: subprocess.exec
+      params:
+        binary: bash
+        args:
+          - ${DRIVERS_TOOLS}/.evergreen/csfle/teardown.sh
+
+  "run encryption tests":
+    - command: subprocess.exec
+      type: test
+      params:
+        binary: bash
+        working_dir: "src"
+        include_expansions_in_env: [
+            "AWS_KMS_ARN",
+            "DRIVERS_TOOLS",
+            "MONGODB_URI",
+            "DJANGO_SETTINGS_MODULE",
+            "CRYPT_SHARED_LIB_PATH",
+        ]
+        args:
+          - ./.evergreen/run-tests.sh
+          - encryption
+
 pre:
   - func: setup
   - func: bootstrap mongo-orchestration
@@ -67,6 +107,12 @@ tasks:
     commands:
       - func: "run unit tests"
 
+  - name: run-encryption-tests
+    commands:
+      - func: "start csfle servers"
+      - func: "run encryption tests"
+      - func: "teardown csfle"
+
 buildvariants:
   - name: tests-7-noauth-nossl
     display_name: Run Tests 7.0 NoAuth NoSSL
@@ -111,3 +157,23 @@ buildvariants:
       SSL: "ssl"
     tasks:
       - name: run-tests
+
+  - name: tests-8-qe-local
+    display_name: Run Tests 8.2 QE local KMS
+    run_on: rhel87-small
+    expansions:
+      MONGODB_VERSION: "8.2"
+      TOPOLOGY: replica_set
+      DJANGO_SETTINGS_MODULE: "encrypted_settings"
+    tasks:
+      - name: run-encryption-tests
+
+  - name: tests-8-qe-aws
+    display_name: Run Tests 8.2 QE AWS KMS
+    run_on: rhel87-small
+    expansions:
+      MONGODB_VERSION: "8.2"
+      TOPOLOGY: replica_set
+      DJANGO_SETTINGS_MODULE: "encrypted_aws_settings"
+    tasks:
+      - name: run-encryption-tests

.evergreen/run-tests.sh

@@ -2,11 +2,24 @@
 
 set -eux
 
-# Install django-mongodb-backend
+# Export secrets as environment variables
+# https://github.com/mongodb-labs/drivers-evergreen-tools/blob/master/.evergreen/csfle/README.md#usage
+if [[ "${1:-}" == "encryption" ]]; then
+    . ../secrets-export.sh
+fi
+
+# Set up virtual environment
 /opt/python/3.12/bin/python3 -m venv venv
 . venv/bin/activate
 python -m pip install -U pip
-pip install -e .
+
+# Install django-mongodb-backend
+if [[ "${1:-}" == "encryption" ]]; then
+    # Install encryption dependencies for the Queryable Encryption build
+    pip install -e '.[encryption]'
+else
+    pip install -e .
+fi
 
 # Install django and test dependencies
 git clone --branch mongodb-6.0.x https://github.com/mongodb-forks/django django_repo

.github/workflows/encrypted_aws_settings.py

@@ -0,0 +1,29 @@
+# Settings for django_mongodb_backend/tests with AWS Key Management System.
+import os
+
+from encrypted_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={
+                "aws": {
+                    "accessKeyId": os.environ["AWS_ACCESS_KEY_ID"],
+                    "secretAccessKey": os.environ["AWS_SECRET_ACCESS_KEY"],
+                }
+            },
+            crypt_shared_lib_path=os.environ["CRYPT_SHARED_LIB_PATH"],
+            crypt_shared_lib_required=True,
+        ),
+    },
+    "KMS_CREDENTIALS": {
+        "aws": {
+            "key": os.environ["AWS_KMS_ARN"],
+            "region": os.environ["AWS_DEFAULT_REGION"],
+        }
+    },
+}

.github/workflows/encrypted_settings.py

@@ -0,0 +1,41 @@
+# 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)}},
+            crypt_shared_lib_path=os.environ["CRYPT_SHARED_LIB_PATH"],
+            crypt_shared_lib_required=True,
+        ),
+        "directConnection": True,
+    },
+}
+
+
+class EncryptedRouter:
+    def db_for_read(self, model, **hints):
+        # All models in the encryption_ app use the encrypted database.
+        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):
+        # Create the encryption_ app's models only in the encrypted database.
+        if app_label == "encryption_":
+            return db == "encrypted"
+        # Don't create other apps' 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/sbom.yml

@@ -37,7 +37,7 @@ jobs:
         run: |
           python -m venv .venv
           source .venv/bin/activate
-          pip install .
+          pip install .[encryption]
           pip uninstall -y pip setuptools
           deactivate
           python -m venv .venv-sbom

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

@@ -0,0 +1,65 @@
+name: Python Tests on Atlas with Encryption
+
+on:
+  pull_request:
+      paths:
+      - '**.py'
+      - '!setup.py'
+      - '.github/workflows/test-python-encryption.yml'
+  workflow_dispatch:
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+defaults:
+  run:
+    shell: bash -eux {0}
+
+jobs:
+  build:
+    name: Django Test Suite
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout django-mongodb-backend
+        uses: actions/checkout@v6
+        with:
+          persist-credentials: false
+      - name: install django-mongodb-backend
+        run: |
+          pip3 install --upgrade pip
+          pip3 install -e .[encryption]
+      - name: Checkout Django
+        uses: actions/checkout@v6
+        with:
+          repository: 'mongodb-forks/django'
+          ref: 'mongodb-6.0.x'
+          path: 'django_repo'
+          persist-credentials: false
+      - name: Install system packages for Django's Python test dependencies
+        run: |
+          sudo apt-get update
+          sudo apt-get install libmemcached-dev
+      - name: Install Django and its Python test dependencies
+        run: |
+          cd django_repo/tests/
+          pip3 install -e ..
+          pip3 install -r requirements/py3.txt
+      - name: Copy the test settings files
+        run: cp .github/workflows/*_settings.py django_repo/tests/
+      - name: Copy the test runner file
+        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:8.0
+      - name: Download mongo_crypt_shared
+        run: |
+          wget https://downloads.mongodb.com/linux/mongo_crypt_shared_v1-linux-x86_64-enterprise-ubuntu2404-8.2.3.tgz
+          tar -xvzf mongo_crypt_shared_v1-linux-x86_64-enterprise-ubuntu2404-8.2.3.tgz lib/mongo_crypt_v1.so
+      - name: Run tests
+        run: python3 django_repo/tests/runtests_.py
+    permissions:
+        contents: read
+    env:
+      CRYPT_SHARED_LIB_PATH: "${{ github.workspace }}/lib/mongo_crypt_v1.so"
+      DJANGO_SETTINGS_MODULE: "encrypted_settings"

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
 
@@ -381,3 +382,55 @@ def on_commit(self, func, robust=False):
                     )
             else:
                 func()
+
+    ## Queryable Encryption properties
+    @cached_property
+    def auto_encryption_opts(self):
+        return self.connection._options.auto_encryption_opts
+
+    @cached_property
+    def client_encryption(self):
+        if auto_encryption_opts := self.auto_encryption_opts:
+            return ClientEncryption(
+                auto_encryption_opts._kms_providers,
+                auto_encryption_opts._key_vault_namespace,
+                self.connection,
+                self.connection.codec_options,
+            )
+        return None
+
+    @cached_property
+    def key_vault(self):
+        # The key vault collection.
+        _, key_vault_collection = self.auto_encryption_opts._key_vault_namespace.split(".", 1)
+        return self.get_collection(key_vault_collection)
+
+    @cached_property
+    def kms_provider(self):
+        # The KMS provider.
+        kms_providers = self.auto_encryption_opts._kms_providers
+        if len(kms_providers) == 1:
+            return next(iter(kms_providers.keys()))
+        # Multiple providers are configured since AutoEncryptionOpts requires
+        # at least one.
+        raise ImproperlyConfigured(
+            "Multiple KMS providers per database aren't supported. Please "
+            "create a feature request with details about your use case."
+        )
+
+    @cached_property
+    def kms_credentials(self):
+        # The KMS credentials (also called "master key"). N/A for local KMS.
+        kms_provider = self.kms_provider
+        if kms_provider == "local":
+            return None
+        if "KMS_CREDENTIALS" not in self.settings_dict:
+            raise ImproperlyConfigured(
+                f"DATABASES['{self.alias}'] is missing 'KMS_CREDENTIALS' "
+                f"required for KMS '{kms_provider}'."
+            )
+        if kms_provider not in self.settings_dict["KMS_CREDENTIALS"]:
+            raise ImproperlyConfigured(
+                f"DATABASES['{self.alias}']['KMS_CREDENTIALS'] is missing '{kms_provider}' key."
+            )
+        return self.settings_dict["KMS_CREDENTIALS"][kms_provider]

django_mongodb_backend/creation.py

@@ -1,12 +1,19 @@
 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.
+        if opts := self.connection.settings_dict["OPTIONS"].get("auto_encryption_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 +31,8 @@ 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.
+        if opts := self.connection.settings_dict["OPTIONS"].get("auto_encryption_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,6 +588,14 @@ 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_8_0(self):
+        return self.mongodb_version >= (8, 0)
+
     @cached_property
     def supports_atlas_search(self):
         """Does the server support Atlas search queries and search indexes?"""
@@ -614,3 +622,22 @@ 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):
+        """
+        For testing purposes, Queryable Encryption requires a MongoDB 8.0 or
+        later replica set or sharded cluster, as well as MongoDB Atlas or
+        Enterprise. This flag must not guard any non-test functionality since
+        it would prevent MongoDB 7.0 from being used, which also supports
+        Queryable Encryption. The models in tests/encryption_ aren't compatible
+        with MongoDB 7.0 because {"queryType": "range"} is "rangePreview".
+        """
+        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
+        )