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,126 @@
+# Queryable Encryption helpers
+#
+# TODO: Decide if these helpers should even exist, and if so, find a permanent
+# place for them.
+
+from bson.binary import STANDARD
+from bson.codec_options import CodecOptions
+from pymongo.encryption import AutoEncryptionOpts, ClientEncryption
+
+KEY_VAULT_DATABASE_NAME = "keyvault"
+KEY_VAULT_COLLECTION_NAME = "__keyVault"
+KMS_PROVIDER = "local" # e.g., "aws", "azure", "gcp", "kmip", or "local"
+
+
+class EqualityQuery:
+ """
+ Represents an encrypted equality query for encrypted fields in MongoDB's
+ Queryable Encryption.
+ """
+
+ def __init__(self, contention=None):
+ self.queryType = "equality"
+ self.contention = contention
+
+ def to_dict(self):
+ query_type = {"queryType": self.queryType}
+ if self.contention is not None:
+ query_type["contention"] = self.contention
+ return [query_type]
+
+
+class RangeQuery:
+ """Represents an encrypted range query configuration for encrypted fields in
+ MongoDB's Queryable Encryption.
+ """
+
+ def __init__(self, sparsity=None, precision=None, trimFactor=None):
+ self.queryType = "range"
+ self.sparsity = sparsity
+ self.precision = precision
+ self.trimFactor = trimFactor
+
+ def to_dict(self):
+ query_type = {"queryType": self.queryType}
+ if self.sparsity is not None:
+ query_type["sparsity"] = self.sparsity
+ if self.precision is not None:
+ query_type["precision"] = self.precision
+ if self.trimFactor is not None:
+ query_type["trimFactor"] = self.trimFactor
+ return query_type
+
+
+class QueryTypes:
+ """
+ Factory class for creating query type configurations for
+ MongoDB Queryable Encryption.
+ """
+
+ def equality(self, *, contention=None):
+ return EqualityQuery(contention=contention)
+
+ def range(self, *, sparsity=None, precision=None, trimFactor=None):
+ return RangeQuery(sparsity=sparsity, precision=precision, trimFactor=trimFactor)
+
+
+def get_auto_encryption_opts(
+ key_vault_namespace=None, crypt_shared_lib_path=None, kms_providers=None
+):
+ """
+ Returns an `AutoEncryptionOpts` instance for MongoDB Client-Side Field
+ Level Encryption (CSFLE) that can be used to create an encrypted connection.
+ """
+ return AutoEncryptionOpts(
+ key_vault_namespace=key_vault_namespace,
+ kms_providers=kms_providers,
+ crypt_shared_lib_path=crypt_shared_lib_path,
+ )
+
+
+def get_client_encryption(client, key_vault_namespace=None, kms_providers=None):
+ """
+ Returns a `ClientEncryption` instance for MongoDB Client-Side Field Level
+ Encryption (CSFLE) that can be used to create an encrypted collection.
+ """
+
+ codec_options = CodecOptions(uuid_representation=STANDARD)
+ return ClientEncryption(kms_providers, key_vault_namespace, client, codec_options)
+
+
+def get_customer_master_key():
+ """
+ Returns a 96-byte local master key for use with MongoDB Client-Side Field Level
+ Encryption (CSFLE). For local testing purposes only. In production, use a secure KMS
+ like AWS, Azure, GCP, or KMIP.
+ Returns:
+ bytes: A 96-byte key.
+ """
+ # WARNING: This is a static key for testing only.
+ # Generate with: os.urandom(96)
+ return bytes.fromhex(
+ "000102030405060708090a0b0c0d0e0f"
+ "101112131415161718191a1b1c1d1e1f"
+ "202122232425262728292a2b2c2d2e2f"
+ "303132333435363738393a3b3c3d3e3f"
+ "404142434445464748494a4b4c4d4e4f"
+ "505152535455565758595a5b5c5d5e5f"
+ )
+
+
+def get_key_vault_namespace(
+ key_vault_database_name=KEY_VAULT_DATABASE_NAME,
+ key_vault_collection_name=KEY_VAULT_COLLECTION_NAME,
+):
+ return f"{key_vault_database_name}.{key_vault_collection_name}"
+
+
+def get_kms_providers():
+ """
+ Return supported KMS providers for MongoDB Client-Side Field Level Encryption (CSFLE).
+ """
+ return {
+ "local": {
+ "key": get_customer_master_key(),
+ },
+ }

django_mongodb_backend/features.py

@@ -624,3 +624,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

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,26 @@
+from django.db import models
+
+
+class EncryptedCharField(models.CharField):
+ encrypted = True
+ queries = []
+
+ def __init__(self, *args, queries=None, **kwargs):
+ self.queries = queries
+ super().__init__(*args, **kwargs)
+
+ def deconstruct(self):
+ name, path, args, kwargs = super().deconstruct()
+
+ # Add 'queries' to kwargs if it was set
+ if self.queries is not None:
+ kwargs["queries"] = self.queries
+
+ # Normalize path if needed
+ 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/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
+
+
+class Command(BaseCommand):
+ help = "Generate an encryptedFieldsMap for MongoDB automatic encryption"
+
+ def handle(self, *args, **options):
+ connection = connections[DEFAULT_DB_ALIAS]
+
+ 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 model in apps.get_models():
+ 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):
+ fields = model._meta.fields
+ encrypted_fields = []
+
+ for field in fields:
+ if getattr(field, "encrypted", False):
+ field_map = {
+ "path": field.column,
+ "bsonType": field.db_type(conn),
+ }
+
+ if getattr(field, "queries", None):
+ field_map["queries"] = field.queries[0].to_dict()
+
+ encrypted_fields.append(field_map)
+
+ return encrypted_fields

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,11 @@
+from django.conf import settings
 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 +42,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 +419,46 @@ 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 not encrypted, create a normal collection otherwise
+ create an encrypted collection with the encrypted fields map.
+ """
+
+ db = self.get_database()
+ if getattr(model, "encrypted", False):
+ client = self.connection.connection
+ ce = get_client_encryption(
+ client,
+ key_vault_namespace=settings.KEY_VAULT_NAMESPACE,
+ kms_providers=settings.KMS_PROVIDERS,
+ )
+ ce.create_encrypted_collection(
+ db,
+ model._meta.db_table,
+ self._get_encrypted_fields_map(model),
+ settings.KMS_PROVIDER,
+ )
+ 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[0].to_dict()}
+ if getattr(field, "queries", None)
+ else {}
+ ),
+ }
+ for field in fields
+ if getattr(field, "encrypted", False)
+ ]
+ }

docs/source/topics/encrypted-models.rst

@@ -0,0 +1,22 @@
+Encrypted models
+================
+
+``EncryptedCharField``
+----------------------
+
+The basics
+~~~~~~~~~~
+
+Let's consider this example::
+
+ from django.db import models
+
+ from django_mongodb_backend.fields import EncryptedCharField
+ from django_mongodb_backend.models import EncryptedModel
+
+
+ class Person(EncryptedModel):
+ ssn = EncryptedCharField("ssn", max_length=11)
+
+ def __str__(self):
+ return self.ssn

docs/source/topics/index.rst

@@ -10,4 +10,5 @@ know:
 
  cache
  embedded-models
+ encrypted-models
  known-issues

tests/backend_/test_features.py

@@ -44,3 +44,21 @@ def mocked_command(command):
 
  with patch("pymongo.synchronous.database.Database.command", wraps=mocked_command):
  self.assertIs(connection.features.supports_transactions, False)
+
+
+class SupportsQueryableEncryptionTests(TestCase):
+ def setUp(self):
+ # Clear the cached property.
+ connection.features.__dict__.pop("supports_queryable_encryption", None)
+
+ def tearDown(self):
+ connection.features.__dict__.pop("supports_queryable_encryption", None)
+
+ def test_supports_queryable_encryption(self):
+ def mocked_command(command):
+ if command == "buildInfo":
+ return {"modules": ["enterprise"]}
+ raise Exception("Unexpected command")
+
+ with patch("pymongo.synchronous.database.Database.command", wraps=mocked_command):
+ self.assertIs(connection.features.supports_queryable_encryption, True)

tests/encryption_/models.py

@@ -0,0 +1,20 @@
+from django.db import models
+
+from django_mongodb_backend.encryption import QueryTypes
+from django_mongodb_backend.fields import EncryptedCharField
+from django_mongodb_backend.models import EncryptedModel
+
+# Query types for encrypted fields with optional parameters
+query_types = QueryTypes()
+queries = [query_types.equality(contention=1), query_types.range(sparsity=2, precision=3)]
+
+
+class Person(EncryptedModel):
+ name = models.CharField("name", max_length=100)
+ ssn = EncryptedCharField("ssn", max_length=11, queries=queries)
+
+ class Meta:
+ required_db_features = {"supports_queryable_encryption"}
+
+ def __str__(self):
+ return self.name