INTPYTHON-527 Add Queryable Encryption support

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,129 @@
+# Queryable Encryption helpers
+import os
+
+from bson.binary import STANDARD
+from bson.codec_options import CodecOptions
+from pymongo.encryption import AutoEncryptionOpts, ClientEncryption
+
+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"),
+    },
+    # TODO: Provide a valid test key
+    #
+    # "Failed to parse KMS provider gcp: unable to parse base64 from UTF-8 field privateKey"
+    #
+    # "gcp": {
+    #     "email": os.getenv("GCP_EMAIL", "not an email"),
+    #     "privateKey": os.getenv("GCP_PRIVATE_KEY", "not a private key"),
+    # },
+    "kmip": {
+        "endpoint": os.getenv("KMIP_KMS_ENDPOINT", "not a valid endpoint"),
+    },
+    "local": {
+        "key": bytes.fromhex(
+            "000102030405060708090a0b0c0d0e0f"
+            "101112131415161718191a1b1c1d1e1f"
+            "202122232425262728292a2b2c2d2e2f"
+            "303132333435363738393a3b3c3d3e3f"
+            "404142434445464748494a4b4c4d4e4f"
+            "505152535455565758595a5b5c5d5e5f"
+        )
+    },
+}
+
+
+class EncryptedRouter:
+    def _get_db_for_model(self, model):
+        if getattr(model, "encrypted", False):
+            return "encrypted"
+        return "default"
+
+    def db_for_read(self, model, **hints):
+        return self._get_db_for_model(model)
+
+    def db_for_write(self, model, **hints):
+        return self._get_db_for_model(model)
+
+    def allow_migrate(self, db, app_label, model_name=None, model=None, **hints):
+        if model:
+            return db == self._get_db_for_model(model)
+        return db == "default"
+
+
+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, *, sparsity=None, precision=None, trimFactor=None):
+        query = {"queryType": "range"}
+        if sparsity is not None:
+            query["sparsity"] = sparsity
+        if precision is not None:
+            query["precision"] = precision
+        if trimFactor is not None:
+            query["trimFactor"] = trimFactor
+        return query
+
+
+def get_auto_encryption_opts(
+    *, key_vault_namespace, crypt_shared_lib_path=None, kms_providers=None, schema_map=None
+):
+    """
+    Returns an `AutoEncryptionOpts` instance for use with Queryable Encryption.
+    """
+    # WARNING: Provide a schema map for production use. You can generate a schema map
+    # with the management command `get_encrypted_fields_map` after adding
+    # django_mongodb_backend to INSTALLED_APPS.
+    return AutoEncryptionOpts(
+        key_vault_namespace=key_vault_namespace,
+        kms_providers=kms_providers,
+        crypt_shared_lib_path=crypt_shared_lib_path,
+        schema_map=schema_map,
+    )
+
+
+def get_client_encryption(client, key_vault_namespace=None, kms_providers=None):
+    """
+    Returns a `ClientEncryption` instance for use with Queryable Encryption.
+    """
+
+    codec_options = CodecOptions(uuid_representation=STANDARD)
+    return ClientEncryption(kms_providers, key_vault_namespace, client, codec_options)

django_mongodb_backend/features.py

@@ -592,6 +592,10 @@ def django_test_expected_failures(self):
     def is_mongodb_6_3(self):
         return self.connection.get_database_version() >= (6, 3)
 
+    @cached_property
+    def is_mongodb_7_0(self):
+        return self.connection.get_database_version() >= (7, 0)
+
     @cached_property
     def supports_atlas_search(self):
         """Does the server support Atlas search queries and search indexes?"""
@@ -624,3 +628,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 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
+        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
 from .json import register_json_field
 from .objectid import ObjectIdField
 
@@ -11,6 +12,7 @@
     "ArrayField",
     "EmbeddedModelArrayField",
     "EmbeddedModelField",
+    "EncryptedCharField",
     "ObjectIdAutoField",
     "ObjectIdField",
 ]

django_mongodb_backend/fields/encryption.py

@@ -0,0 +1,23 @@
+from django.db import models
+
+
+class EncryptedCharField(models.CharField):
+    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

django_mongodb_backend/functions.py

@@ -38,7 +38,9 @@
     Trim,
     Upper,
 )
+from django.db.utils import ConnectionRouter
 
+from .encryption import KMS_CREDENTIALS
 from .query_utils import process_lhs
 
 MONGO_OPERATORS = {
@@ -268,10 +270,20 @@ def trunc_time(self, compiler, connection):
     }
 
 
+def kms_credentials(self, provider):  # noqa: ARG001
+    return KMS_CREDENTIALS.get(provider, None)
+
+
+def kms_provider(self):  # noqa: ARG001
+    return getattr(settings, "KMS_PROVIDER", None)
+
+
 def register_functions():
     Cast.as_mql = cast
     Concat.as_mql = concat
     ConcatPair.as_mql = concat_pair
+    ConnectionRouter.kms_credentials = kms_credentials
+    ConnectionRouter.kms_provider = kms_provider
     Cot.as_mql = cot
     Extract.as_mql = extract
     Func.as_mql = func

django_mongodb_backend/management/commands/get_encrypted_fields_map.py

@@ -0,0 +1,45 @@
+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 `get_autoencryption_opts` 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.generate_encrypted_fields_schema_map(connection)
+
+        self.stdout.write(json.dumps(schema_map, indent=2))
+
+    def generate_encrypted_fields_schema_map(self, conn):
+        schema_map = {}
+
+        for app_config in apps.get_app_configs():
+            for model in router.get_migratable_models(
+                app_config, conn.alias, include_auto_created=False
+            ):
+                if getattr(model, "encrypted", False):
+                    encrypted_fields = self.get_encrypted_fields(model, conn)
+                    if encrypted_fields:
+                        collection = model._meta.db_table
+                        schema_map[collection] = {"fields": encrypted_fields}
+
+        return schema_map
+
+    def get_encrypted_fields(self, model, conn):
+        return conn.schema_editor()._get_encrypted_fields_map(model)

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/schema.py

@@ -1,10 +1,12 @@
+from django.core.exceptions import ImproperlyConfigured
+from django.db import router
 from django.db.backends.base.schema import BaseDatabaseSchemaEditor
 from django.db.models import Index, UniqueConstraint
 from pymongo.operations import SearchIndexModel
 
-from django_mongodb_backend.indexes import SearchIndex
-
+from .encryption import get_client_encryption
 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,54 @@ 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.
+        """
+
+        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
+
+            ce = get_client_encryption(
+                client,
+                key_vault_namespace=key_vault_namespace,
+                kms_providers=kms_providers,
+            )
+            if not router.kms_provider():
+                raise ImproperlyConfigured(
+                    "No KMS_PROVIDER found. Please configure KMS_PROVIDER in settings."
+                )
+            provider = router.kms_provider()
+            credentials = router.kms_credentials(provider)
+            ce.create_encrypted_collection(
+                db,
+                model._meta.db_table,
+                self._get_encrypted_fields_map(model),
+                provider,
+                credentials,
+            )
+        else:
+            db.create_collection(model._meta.db_table)
+
+    def _get_encrypted_fields_map(self, model):
+        conn = self.connection
+        fields = model._meta.fields
+
+        return {
+            "fields": [
+                {
+                    "path": field.column,
+                    "bsonType": field.db_type(conn),
+                    **({"queries": field.queries} if getattr(field, "queries", None) else {}),
+                }
+                for field in fields
+                if getattr(field, "encrypted", False)
+            ]
+        }