INTPYTHON-527 Add Queryable Encryption support

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,112 @@
+# 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
+
+KEY_VAULT_COLLECTION_NAME = "__keyVault"
+KEY_VAULT_DATABASE_NAME = "keyvault"
+KEY_VAULT_NAMESPACE = f"{KEY_VAULT_DATABASE_NAME}.{KEY_VAULT_COLLECTION_NAME}"
+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 .encryption 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/encryption.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[model._meta.db_table] = fields
+        return schema_map

django_mongodb_backend/models.py

@@ -14,3 +14,11 @@ 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
+        required_db_features = {"supports_queryable_encryption"}

django_mongodb_backend/routers.py

@@ -1,6 +1,5 @@
 from django.apps import apps
-
-from django_mongodb_backend.models import EmbeddedModel
+from django.db.utils import ConnectionRouter
 
 
 class MongoRouter:
@@ -9,10 +8,23 @@ 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 = ConnectionRouter._router_func("kms_provider")

django_mongodb_backend/schema.py

@@ -1,10 +1,12 @@
+from django.conf import settings
+from django.db import router
 from django.db.backends.base.schema import BaseDatabaseSchemaEditor
 from django.db.models import Index, UniqueConstraint
+from pymongo.encryption import ClientEncryption, CodecOptions
 from pymongo.operations import SearchIndexModel
 
-from django_mongodb_backend.indexes import SearchIndex
-
 from .fields import EmbeddedModelField
+from .indexes import SearchIndex
 from .query import wrap_database_errors
 from .utils import OperationCollector
 
@@ -41,7 +43,7 @@ def get_database(self):
     @wrap_database_errors
     @ignore_embedded_models
     def create_model(self, model):
-        self.get_database().create_collection(model._meta.db_table)
+        self._create_collection(model)
         self._create_model_indexes(model)
         # Make implicit M2M tables.
         for field in model._meta.local_many_to_many:
@@ -418,3 +420,61 @@ def _field_should_have_unique(self, field):
         db_type = field.db_type(self.connection)
         # The _id column is automatically unique.
         return db_type and field.unique and field.column != "_id"
+
+    def _create_collection(self, model):
+        """
+        If the model is encrypted create an encrypted collection with the
+        encrypted fields map else create a normal collection.
+        """
+
+    def _create_collection(self, model):
+        """
+        If the model is encrypted, create an encrypted collection with the
+        encrypted fields map; else, create a normal collection.
+        """
+        db = self.get_database()
+        if getattr(model, "encrypted", False):
+            client = self.connection.connection
+            options = client._options.auto_encryption_opts
+            key_vault_namespace = options._key_vault_namespace
+            kms_providers = options._kms_providers
+            codec_options = CodecOptions()
+
+            ce = ClientEncryption(kms_providers, key_vault_namespace, client, codec_options)
+
+            # TODO: Validate schema! `create_encrypted_collection` appears to
+            # succeed no matter what you give it, as long as it's valid JSON.
+            # E.g. encrypted_fields_map = []
+            encrypted_fields_map = self._get_encrypted_fields_map(model)
+            provider = router.kms_provider(model)
+
+            # TODO: Remove ternary condition when `master_key` option is not
+            # inadvertently set to "default" somewhere, which then causes the
+            # `master_key.copy` in libmongocrypt to fail.
+            credentials = settings.DATABASES[db].KMS_CREDENTIALS if provider != "local" else None
+
+            ce.create_encrypted_collection(
+                db,
+                model._meta.db_table,
+                encrypted_fields_map,
+                provider,
+                credentials,
+            )
+        else:
+            db.create_collection(model._meta.db_table)
+
+    def _get_encrypted_fields_map(self, model):
+        connection = self.connection
+        fields = model._meta.fields
+
+        return {
+            "fields": [
+                {
+                    "bsonType": field.db_type(connection),
+                    "path": field.column,
+                    **({"queries": field.queries} if getattr(field, "queries", None) else {}),
+                }
+                for field in fields
+                if getattr(field, "encrypted", False)
+            ]
+        }

docs/source/conf.py

@@ -48,8 +48,6 @@
     "manual": ("https://www.mongodb.com/docs/manual/", None),
 }
 
-root_doc = "contents"
-
 # -- Options for HTML output -------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output