INTPYTHON-527 Add Queryable Encryption support

.evergreen/config.yml

@@ -48,13 +48,32 @@ functions:
         args:
           - ./.evergreen/run-tests.sh
 
+  "run encryption tests":
+    - command: ec2.assume_role
+      params:
+        role_arn: ${aws_role_arn}
+      type: test
+    - command: subprocess.exec
+      params:
+        binary: bash
+        args:
+          - .evergreen/run-encryption-tests.sh
+        working_dir: src
+        include_expansions_in_env:
+          - AWS_ACCESS_KEY_ID
+          - AWS_SECRET_ACCESS_KEY
+          - AWS_SESSION_TOKEN
+
   "teardown":
     - command: subprocess.exec
       params:
         binary: bash
         args:
           - ${DRIVERS_TOOLS}/.evergreen/teardown.sh
 
+
+
+
 pre:
   - func: setup
   - func: bootstrap mongo-orchestration

.evergreen/run-encryption-tests.sh

@@ -0,0 +1,36 @@
+#!/usr/bin/bash
+
+set -eux
+
+# Clone drivers-tools and set var
+git clone https://github.com/mongodb-labs/drivers-evergreen-tools.git drivers_tools
+export DRIVERS_TOOLS=$(pwd)/drivers_tools
+
+git clone https://github.com/mongodb/pymongo/ pymongo_repo
+pushd pymongo_repo
+
+# Setup encryption
+just run-server --topology=replica_set
+just setup-tests encryption
+
+# Install django-mongodb-backend
+/opt/python/3.10/bin/python3 -m venv venv
+. venv/bin/activate
+python -m pip install -U pip
+pip install -e .\[encryption\]
+
+# Install django and test dependencies
+git clone --branch mongodb-5.2.x https://github.com/mongodb-forks/django django_repo
+pushd django_repo/tests/
+pip install -e ..
+pip install -r requirements/py3.txt
+popd
+
+# Copy the test settings file
+cp ./.github/workflows/mongodb_settings.py django_repo/tests/
+
+# Copy the test runner file
+cp ./.github/workflows/runtests.py django_repo/tests/runtests_.py
+
+# Run tests
+python django_repo/tests/runtests_.py

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

@@ -286,4 +286,7 @@ def validate_no_broken_transaction(self):
 
     def get_database_version(self):
         """Return a tuple of the database's version."""
-        return tuple(self.connection.server_info()["versionArray"])
+        # Avoid PyMongo or require PyMongo>=4.14.0 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,109 @@
+# Queryable Encryption helper functions and constants for MongoDB
+#
+# These helper functions and constants are optional and Queryable
+# Encryption can be used in Django without them. They are provided
+# to make it easier configure Queryable Encryption in Django.
+
+import base64
+import os
+
+KMS_CREDENTIALS = {
+    "aws": {
+        "key": os.getenv("AWS_KEY_ARN", ""),
+        "region": os.getenv("AWS_KEY_REGION", ""),
+    },
+    "azure": {
+        "keyName": os.getenv("AZURE_KEY_NAME", ""),
+        "keyVaultEndpoint": os.getenv("AZURE_KEY_VAULT_ENDPOINT", ""),
+    },
+    "gcp": {
+        "projectId": os.getenv("GCP_PROJECT_ID", ""),
+        "location": os.getenv("GCP_LOCATION", ""),
+        "keyRing": os.getenv("GCP_KEY_RING", ""),
+        "keyName": os.getenv("GCP_KEY_NAME", ""),
+    },
+    "kmip": {},
+    "local": {},
+}
+KMS_PROVIDERS = {
+    "aws": {
+        "accessKeyId": os.getenv("AWS_ACCESS_KEY_ID", "not an access key"),
+        "secretAccessKey": os.getenv("AWS_SECRET_ACCESS_KEY", "not a secret key"),
+    },
+    "azure": {
+        "tenantId": os.getenv("AZURE_TENANT_ID", "not a tenant ID"),
+        "clientId": os.getenv("AZURE_CLIENT_ID", "not a client ID"),
+        "clientSecret": os.getenv("AZURE_CLIENT_SECRET", "not a client secret"),
+    },
+    "gcp": {
+        "email": os.getenv("GCP_EMAIL", "not an email"),
+        "privateKey": os.getenv(
+            "GCP_PRIVATE_KEY",
+            base64.b64encode(b"not a private key").decode("ascii"),
+        ),
+    },
+    "kmip": {
+        "endpoint": os.getenv("KMIP_KMS_ENDPOINT", "not a valid endpoint"),
+    },
+    "local": {
+        "key": bytes.fromhex(
+            "000102030405060708090a0b0c0d0e0f"
+            "101112131415161718191a1b1c1d1e1f"
+            "202122232425262728292a2b2c2d2e2f"
+            "303132333435363738393a3b3c3d3e3f"
+            "404142434445464748494a4b4c4d4e4f"
+            "505152535455565758595a5b5c5d5e5f"
+        )
+    },
+}
+
+
+class EncryptedRouter:
+    """A sample database router for Django that routes encrypted
+    models to an encrypted database with a local KMS provider.
+    """
+
+    def allow_migrate(self, db, app_label, model_name=None, model=None, **hints):
+        if model:
+            return db == ("encrypted" if getattr(model, "encrypted", False) else "default")
+        return db == "default"
+
+    def db_for_read(self, model, **hints):
+        if getattr(model, "encrypted", False):
+            return "encrypted"
+        return "default"
+
+    db_for_write = db_for_read
+
+    def kms_provider(self, model):
+        return "local"
+
+
+class QueryType:
+    """
+    Class that supports building encrypted equality and range queries
+    for MongoDB's Queryable Encryption.
+    """
+
+    @classmethod
+    def equality(cls, *, contention=None):
+        query = {"queryType": "equality"}
+        if contention is not None:
+            query["contention"] = contention
+        return query
+
+    @classmethod
+    def range(
+        cls, *, contention=None, max=None, min=None, precision=None, sparsity=None, trimFactor=None
+    ):
+        query = {"queryType": "range"}
+        options = {
+            "contention": contention,
+            "max": max,
+            "min": min,
+            "precision": precision,
+            "sparsity": sparsity,
+            "trimFactor": trimFactor,
+        }
+        query.update({k: v for k, v in options.items() if v is not None})
+        return query

django_mongodb_backend/features.py

@@ -588,9 +588,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):
@@ -624,3 +632,19 @@ 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 is supported if the server is Atlas or Enterprise
+        and is configured as a replica set or sharded cluster.
+        """
+        self.connection.ensure_connection()
+        client = self.connection.connection.admin
+        build_info = client.command("buildInfo")
+        is_enterprise = "enterprise" in build_info.get("modules")
+        # `supports_transactions` already checks if the server is a
+        # replica set or sharded cluster.
+        is_not_single = self.supports_transactions
+        # TODO: check if the server is Atlas
+        return is_enterprise and is_not_single and self.is_mongodb_7_0

django_mongodb_backend/fields/__init__.py

@@ -3,6 +3,7 @@
 from .duration import register_duration_field
 from .embedded_model import EmbeddedModelField
 from .embedded_model_array import EmbeddedModelArrayField
+from .encrypted_model import EncryptedCharField, EncryptedIntegerField
 from .json import register_json_field
 from .objectid import ObjectIdField
 
@@ -11,6 +12,8 @@
     "ArrayField",
     "EmbeddedModelArrayField",
     "EmbeddedModelField",
+    "EncryptedCharField",
+    "EncryptedIntegerField",
     "ObjectIdAutoField",
     "ObjectIdField",
 ]

django_mongodb_backend/fields/encrypted_model.py

@@ -0,0 +1,31 @@
+from django.db import models
+
+
+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.encryption"):
+            path = path.replace(
+                "django_mongodb_backend.fields.encryption",
+                "django_mongodb_backend.fields",
+            )
+
+        return name, path, args, kwargs
+
+
+class EncryptedCharField(EncryptedFieldMixin, models.CharField):
+    pass
+
+
+class EncryptedIntegerField(EncryptedFieldMixin, models.IntegerField):
+    pass

django_mongodb_backend/management/commands/get_encrypted_fields_map.py

@@ -0,0 +1,36 @@
+import json
+
+from django.apps import apps
+from django.core.management.base import BaseCommand
+from django.db import DEFAULT_DB_ALIAS, connections, router
+
+
+class Command(BaseCommand):
+    help = "Generate a `schema_map` of encrypted fields for all encrypted"
+    " models in the database for use with `AutoEncryptionOpts` in"
+    " production environments."
+
+    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.",
+        )
+
+    def handle(self, *args, **options):
+        db = options["database"]
+        connection = connections[db]
+        schema_map = self.get_encrypted_fields_map(connection)
+        self.stdout.write(json.dumps(schema_map, indent=2))
+
+    def get_encrypted_fields_map(self, connection):
+        schema_map = {}
+        for app_config in apps.get_app_configs():
+            for model in router.get_migratable_models(
+                app_config, connection.alias, include_auto_created=False
+            ):
+                if getattr(model, "encrypted", False):
+                    fields = connection.schema_editor()._get_encrypted_fields_map(model)
+                    schema_map[f"{connection.alias}.{model._meta.db_table}"] = fields
+        return schema_map

django_mongodb_backend/models.py

@@ -14,3 +14,10 @@ def delete(self, *args, **kwargs):
 
     def save(self, *args, **kwargs):
         raise NotSupportedError("EmbeddedModels cannot be saved.")
+
+
+class EncryptedModel(models.Model):
+    encrypted = True
+
+    class Meta:
+        abstract = True

django_mongodb_backend/routers.py

@@ -1,6 +1,7 @@
 from django.apps import apps
+from django.db.utils import ConnectionRouter
 
-from django_mongodb_backend.models import EmbeddedModel
+from .utils import _router_func
 
 
 class MongoRouter:
@@ -9,10 +10,24 @@ def allow_migrate(self, db, app_label, model_name=None, **hints):
         EmbeddedModels don't have their own collection and must be ignored by
         dumpdata.
         """
+
         if not model_name:
             return None
         try:
             model = apps.get_model(app_label, model_name)
         except LookupError:
             return None
+
+        # Delay import for `register_routers` patching.
+        from django_mongodb_backend.models import EmbeddedModel
+
         return False if issubclass(model, EmbeddedModel) else None
+
+
+def register_routers():
+    """
+    Patch the ConnectionRouter with methods to get KMS credentials and provider
+    from the SchemaEditor.
+    """
+
+    ConnectionRouter.kms_provider = _router_func("kms_provider")