INTPYTHON-527 Add Queryable Encryption support

Author: aclark4life

.evergreen/config.yml

@@ -90,6 +90,28 @@ buildvariants:
     tasks:
       - name: run-tests
 
+  - name: tests-7-noauth-nossl
+    display_name: Run Tests 7.0 NoAuth NoSSL
+    run_on: rhel87-small
+    expansions:
+      MONGODB_VERSION: "7.0"
+      TOPOLOGY: server
+      AUTH: "noauth"
+      SSL: "nossl"
+    tasks:
+      - name: run-tests
+
+  - name: tests-7-auth-ssl
+    display_name: Run Tests 7.0 Auth SSL
+    run_on: rhel87-small
+    expansions:
+      MONGODB_VERSION: "7.0"
+      TOPOLOGY: server
+      AUTH: "auth"
+      SSL: "ssl"
+    tasks:
+      - name: run-tests
+
   - name: tests-8-noauth-nossl
     display_name: Run Tests 8.0 NoAuth NoSSL
     run_on: rhel87-small

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

@@ -229,4 +229,7 @@ def cursor(self):
 
     def get_database_version(self):
         """Return a tuple of the database's version."""
-        return tuple(self.connection.server_info()["versionArray"])
+        # Avoid using PyMongo to check the database version or require
+        # pymongocrypt>=1.14.2 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,83 @@
+# Queryable Encryption helper classes and settings
+
+import os
+
+from .fields import has_encrypted_fields
+
+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": {},
+    "azure": {},
+    "gcp": {},
+    "kmip": {
+        "endpoint": os.getenv("KMIP_KMS_ENDPOINT", "not a valid endpoint"),
+    },
+    "local": {
+        "key": bytes.fromhex(
+            "000102030405060708090a0b0c0d0e0f"
+            "101112131415161718191a1b1c1d1e1f"
+            "202122232425262728292a2b2c2d2e2f"
+            "303132333435363738393a3b3c3d3e3f"
+            "404142434445464748494a4b4c4d4e4f"
+            "505152535455565758595a5b5c5d5e5f"
+        )
+    },
+}
+
+
+class EncryptedRouter:
+    def allow_migrate(self, db, app_label, model_name=None, model=None, **hints):
+        if model:
+            return db == ("other" if has_encrypted_fields(model) else "default")
+        return db == "default"
+
+    def db_for_read(self, model, **hints):
+        if has_encrypted_fields(model):
+            return "other"
+        return "default"
+
+    db_for_write = db_for_read
+
+    def kms_provider(self, model):
+        return "local"
+
+
+class EqualityQuery(dict):
+    def __init__(self, *, contention=None):
+        super().__init__(queryType="equality")
+        if contention is not None:
+            self["contention"] = contention
+
+
+class RangeQuery(dict):
+    def __init__(
+        self, *, contention=None, max=None, min=None, precision=None, sparsity=None, trimFactor=None
+    ):
+        super().__init__(queryType="range")
+        options = {
+            "contention": contention,
+            "max": max,
+            "min": min,
+            "precision": precision,
+            "sparsity": sparsity,
+            "trimFactor": trimFactor,
+        }
+        self.update({k: v for k, v in options.items() if v is not None})

django_mongodb_backend/features.py

@@ -97,6 +97,18 @@ class DatabaseFeatures(BaseDatabaseFeatures):
         "expressions.tests.ExpressionOperatorTests.test_lefthand_transformed_field_bitwise_or",
     }
 
+    _django_test_expected_failures_queryable_encryption = {
+        "encryption.tests.EncryptedFieldTests.test_get_encrypted_fields_map_method",
+        "encryption.tests.EncryptedFieldTests.test_get_encrypted_fields_map_command",
+        "encryption.tests.EncryptedFieldTests.test_set_encrypted_fields_map_in_client",
+        "encryption.tests.EncryptedFieldTests.test_appointment",
+        "encryption.tests.EncryptedFieldTests.test_billing",
+        "encryption.tests.EncryptedFieldTests.test_patientportaluser",
+        "encryption.tests.EncryptedFieldTests.test_patientrecord",
+        "encryption.tests.EncryptedFieldTests.test_patient",
+        "encryption.tests.EncryptedFieldTests.test_env",
+    }
+
     @cached_property
     def django_test_expected_failures(self):
         expected_failures = super().django_test_expected_failures
@@ -569,9 +581,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):
@@ -601,3 +621,20 @@ 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 a sharded cluster.
+        """
+        self.connection.ensure_connection()
+        client = self.connection.connection.admin
+        build_info = client.command("buildInfo")
+        is_enterprise = "enterprise" in build_info.get("modules")
+        # Queryable Encryption requires transaction support which
+        # is only available on replica sets or sharded clusters
+        # which we already check in `supports_transactions`.
+        supports_transactions = self.supports_transactions
+        # TODO: check if the server is Atlas
+        return is_enterprise and supports_transactions and self.is_mongodb_7_0

django_mongodb_backend/fields/__init__.py

@@ -3,15 +3,50 @@
 from .duration import register_duration_field
 from .embedded_model import EmbeddedModelField
 from .embedded_model_array import EmbeddedModelArrayField
+from .encryption import (
+    EncryptedBigIntegerField,
+    EncryptedBinaryField,
+    EncryptedBooleanField,
+    EncryptedCharField,
+    EncryptedDateField,
+    EncryptedDateTimeField,
+    EncryptedDecimalField,
+    EncryptedEmailField,
+    EncryptedFieldMixin,
+    EncryptedFloatField,
+    EncryptedGenericIPAddressField,
+    EncryptedIntegerField,
+    EncryptedTextField,
+    EncryptedTimeField,
+    EncryptedURLField,
+    has_encrypted_fields,
+)
 from .json import register_json_field
 from .objectid import ObjectIdField
 from .polymorphic_embedded_model import PolymorphicEmbeddedModelField
 from .polymorphic_embedded_model_array import PolymorphicEmbeddedModelArrayField
 
 __all__ = [
+    "has_encrypted_fields",
+    "register_fields",
     "ArrayField",
     "EmbeddedModelArrayField",
     "EmbeddedModelField",
+    "EncryptedBigIntegerField",
+    "EncryptedBinaryField",
+    "EncryptedBooleanField",
+    "EncryptedCharField",
+    "EncryptedDateTimeField",
+    "EncryptedDateField",
+    "EncryptedDecimalField",
+    "EncryptedEmailField",
+    "EncryptedFieldMixin",
+    "EncryptedFloatField",
+    "EncryptedGenericIPAddressField",
+    "EncryptedIntegerField",
+    "EncryptedTextField",
+    "EncryptedTimeField",
+    "EncryptedURLField",
     "ObjectIdAutoField",
     "ObjectIdField",
     "PolymorphicEmbeddedModelArrayField",

django_mongodb_backend/fields/encryption.py

@@ -0,0 +1,94 @@
+from django.db import models
+
+
+def has_encrypted_fields(model):
+    return any(getattr(field, "encrypted", False) for field in model._meta.fields)
+
+
+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.encrypted_model"):
+            path = path.replace(
+                "django_mongodb_backend.fields.encrypted_model",
+                "django_mongodb_backend.fields",
+            )
+
+        return name, path, args, kwargs
+
+
+class EncryptedBigIntegerField(EncryptedFieldMixin, models.BigIntegerField):
+    pass
+
+
+class EncryptedBinaryField(EncryptedFieldMixin, models.BinaryField):
+    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 EncryptedEmailField(EncryptedFieldMixin, models.EmailField):
+    pass
+
+
+class EncryptedFloatField(EncryptedFieldMixin, models.FloatField):
+    pass
+
+
+class EncryptedGenericIPAddressField(EncryptedFieldMixin, models.GenericIPAddressField):
+    pass
+
+
+class EncryptedIntegerField(EncryptedFieldMixin, models.IntegerField):
+    pass
+
+
+class EncryptedSlugField(EncryptedFieldMixin, models.SlugField):
+    pass
+
+
+class EncryptedTimeField(EncryptedFieldMixin, models.TimeField):
+    pass
+
+
+class EncryptedTextField(EncryptedFieldMixin, models.TextField):
+    pass
+
+
+class EncryptedURLField(EncryptedFieldMixin, models.URLField):
+    pass
+
+
+# TODO: Add more encrypted fields
+# - PositiveBigIntegerField
+# - PositiveIntegerField
+# - PositiveSmallIntegerField
+# - SmallIntegerField

django_mongodb_backend/management/commands/get_encrypted_fields_map.py

@@ -0,0 +1,59 @@
+from bson import json_util
+from django.apps import apps
+from django.core.management.base import BaseCommand
+from django.db import DEFAULT_DB_ALIAS, connections, router
+from pymongo.encryption import ClientEncryption
+
+
+class Command(BaseCommand):
+    help = "Generate a `schema_map` of encrypted fields for all encrypted"
+    " models in the database for use with `AutoEncryptionOpts` in"
+    " client configuration."
+
+    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.",
+        )
+        parser.add_argument(
+            "--kms-provider",
+            default="local",
+            help="Specify the KMS provider to use for encryption. Defaults to 'local'.",
+        )
+
+    def handle(self, *args, **options):
+        db = options["database"]
+        kms_provider = options["kms_provider"]
+        connection = connections[db]
+        schema_map = json_util.dumps(
+            self.get_encrypted_fields_map(connection, kms_provider), indent=2
+        )
+        self.stdout.write(schema_map)
+
+    def get_client_encryption(self, connection):
+        client = connection.connection
+        options = client._options.auto_encryption_opts
+        key_vault_namespace = options._key_vault_namespace
+        kms_providers = options._kms_providers
+        return ClientEncryption(kms_providers, key_vault_namespace, client, client.codec_options)
+
+    def get_encrypted_fields_map(self, connection, kms_provider):
+        schema_map = {}
+        for app_config in apps.get_app_configs():
+            for model in router.get_migratable_models(
+                app_config, connection.settings_dict["NAME"], include_auto_created=False
+            ):
+                if getattr(model, "encrypted", False):
+                    fields = connection.schema_editor()._get_encrypted_fields_map(model)
+                    ce = self.get_client_encryption(connection)
+                    master_key = connection.settings_dict.get("KMS_CREDENTIALS").get(kms_provider)
+                    for field in fields["fields"]:
+                        data_key = ce.create_data_key(
+                            kms_provider=kms_provider,
+                            master_key=master_key,
+                        )
+                        field["keyId"] = data_key
+                    schema_map[model._meta.db_table] = fields
+        return schema_map