Code review fixes

- Remove query type helpers
- Remove KMS AWS patches
- Refactor encrypted fields map creation
- Remove unused test code
- Doc fixes and updates
- Move has_encrypted_fields to model_utils
- Renamed showfieldsmap -> createencryptedfieldsmap

Author: aclark4life

django_mongodb_backend/fields/__init__.py

@@ -23,9 +23,6 @@
     EncryptedTextField,
     EncryptedTimeField,
     EncryptedURLField,
-    EqualityQuery,
-    RangeQuery,
-    has_encrypted_fields,
 )
 from .json import register_json_field
 from .objectid import ObjectIdField
@@ -55,13 +52,10 @@
     "EncryptedTextField",
     "EncryptedTimeField",
     "EncryptedURLField",
-    "EqualityQuery",
     "ObjectIdAutoField",
     "ObjectIdField",
     "PolymorphicEmbeddedModelArrayField",
     "PolymorphicEmbeddedModelField",
-    "RangeQuery",
-    "has_encrypted_fields",
     "register_fields",
 ]
 

django_mongodb_backend/fields/encryption.py

@@ -1,10 +1,6 @@
 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
 
@@ -97,26 +93,3 @@ class EncryptedTextField(EncryptedFieldMixin, models.TextField):
 
 class EncryptedURLField(EncryptedFieldMixin, models.URLField):
     pass
-
-
-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/management/commands/createencryptedfieldsmap.py

@@ -0,0 +1,40 @@
+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
+
+from django_mongodb_backend.model_utils import has_encrypted_fields
+
+
+class Command(BaseCommand):
+    help = "Generate an `encrypted_fields_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.",
+        )
+
+    def handle(self, *args, **options):
+        db = options["database"]
+        connection = connections[db]
+        client = connection.connection
+        encrypted_fields_map = {}
+        auto_encryption_opts = getattr(client._options, "auto_encryption_opts", None)
+        # FIXME:
+        # TypeError: ConnectionRouter.get_migratable_models() missing 2 required
+        # positional arguments: 'app_config' and 'db'
+        # for app_config in router.get_migratable_models():
+        for app_config in apps.get_app_configs():
+            for model in app_config.get_models():
+                if has_encrypted_fields(model):
+                    encrypted_fields_map[model._meta.db_table] = (
+                        connection.schema_editor()._get_encrypted_fields_map(
+                            model, client, auto_encryption_opts
+                        )
+                    )
+        self.stdout.write(json_util.dumps(encrypted_fields_map, indent=2))

django_mongodb_backend/management/commands/showfieldsmap.py

@@ -0,0 +1,3 @@
+# TODO: Move to models.utils
+def has_encrypted_fields(model):
+    return any(getattr(field, "encrypted", False) for field in model._meta.fields)

django_mongodb_backend/model_utils.py

@@ -2,7 +2,7 @@
 from django.core.exceptions import ImproperlyConfigured
 from django.db.utils import ConnectionRouter
 
-from .fields import has_encrypted_fields
+from .model_utils import has_encrypted_fields
 
 
 class MongoRouter:

django_mongodb_backend/routers.py

@@ -5,8 +5,9 @@
 from pymongo.encryption import ClientEncryption
 from pymongo.operations import SearchIndexModel
 
-from .fields import EmbeddedModelField, has_encrypted_fields
+from .fields import EmbeddedModelField
 from .indexes import SearchIndex
+from .model_utils import has_encrypted_fields
 from .query import wrap_database_errors
 from .utils import OperationCollector
 
@@ -432,51 +433,50 @@ def _create_collection(self, model):
         db_table = model._meta.db_table
         if has_encrypted_fields(model):
             client = self.connection.connection
-            options = client._options
-            ae = getattr(options, "auto_encryption_opts", None)
-            if not ae:
+            auto_encryption_opts = getattr(client._options, "auto_encryption_opts", None)
+            if not auto_encryption_opts:
                 raise ImproperlyConfigured(
-                    "Encrypted fields found but the connection does not have "
-                    "auto encryption options set. Please set `auto_encryption_opts` "
+                    "Encrypted fields found but "
+                    "DATABASES[[self.connection.alias}]['OPTIONS'] is missing "
+                    "auto_encryption_opts. Please set `auto_encryption_opts` "
                     "in the connection settings."
                 )
-            encrypted_fields_map = getattr(ae, "_encrypted_fields_map", None)
+            encrypted_fields_map = getattr(auto_encryption_opts, "_encrypted_fields_map", None)
             if not encrypted_fields_map:
-                encrypted_fields_map = {}
-                ce = ClientEncryption(
-                    ae._kms_providers,
-                    ae._key_vault_namespace,
-                    client,
-                    client.codec_options,
+                encrypted_fields_map = self._get_encrypted_fields_map(
+                    model, client, auto_encryption_opts
                 )
-                fields = self._get_encrypted_fields_map(model)
-                kms_provider = router.kms_provider(model)
-                master_key = self.connection.settings_dict.get("KMS_CREDENTIALS").get(kms_provider)
-                for field in fields["fields"]:
-                    key_alt_name = f"{db_table}_{field['path']}"
-                    data_key = ce.create_data_key(
-                        kms_provider=kms_provider,
-                        master_key=master_key,
-                        key_alt_names=[key_alt_name],
-                    )
-                    field["keyId"] = data_key
-                    encrypted_fields_map[db_table] = fields
-            db.create_collection(db_table, encryptedFields=encrypted_fields_map[db_table])
+            db.create_collection(db_table, encryptedFields=encrypted_fields_map)
         else:
             db.create_collection(db_table)
 
-    def _get_encrypted_fields_map(self, model):
+    def _get_encrypted_fields_map(self, model, client, auto_encryption_opts):
         connection = self.connection
         fields = model._meta.fields
-
-        return {
-            "fields": [
-                {
+        kms_provider = router.kms_provider(model)
+        master_key = self.connection.settings_dict.get("KMS_CREDENTIALS").get(kms_provider)
+        client_encryption = ClientEncryption(
+            auto_encryption_opts._kms_providers,
+            auto_encryption_opts._key_vault_namespace,
+            client,
+            client.codec_options,
+        )
+        db_table = model._meta.db_table
+        field_list = []
+        for field in fields:
+            if getattr(field, "encrypted", False):
+                key_alt_name = f"{db_table}_{field.column}"
+                data_key = client_encryption.create_data_key(
+                    kms_provider=kms_provider,
+                    master_key=master_key,
+                    key_alt_names=[key_alt_name],
+                )
+                field_dict = {
                     "bsonType": field.db_type(connection),
                     "path": field.column,
-                    **({"queries": field.queries} if getattr(field, "queries", None) else {}),
+                    "keyId": data_key,
                 }
-                for field in fields
-                if getattr(field, "encrypted", False)
-            ]
-        }
+                if getattr(field, "queries", None):
+                    field_dict["queries"] = field.queries
+                field_list.append(field_dict)
+        return {"fields": field_list}

django_mongodb_backend/schema.py

@@ -58,49 +58,48 @@ see :ref:`configuring-database-routers-setting`.
 Queryable Encryption
 ====================
 
-What about client side configuration?
+What about client-side configuration?
 -------------------------------------
 
 In the :doc:`Queryable Encryption how-to guide <howto/queryable-encryption>`,
-server side Queryable Encryption configuration is covered.
+server-side Queryable Encryption configuration is covered.
 
-Client side Queryable Encryption configuration requires that the entire schema
-for encrypted fields is known at the time of client connection.
+Client side Queryable Encryption configuration requires that the entire
+encrypted fields map be known at the time of client connection.
 
-Schema Map
-~~~~~~~~~~
+Encrypted fields map
+~~~~~~~~~~~~~~~~~~~~
 
 In addition to the
 :ref:`settings described in the how-to guide <server-side-queryable-encryption-settings>`,
-you will need to provide a ``schema_map`` to the ``AutoEncryptionOpts``.
+you will need to provide a ``encrypted_fields_map`` to the
+``AutoEncryptionOpts``.
 
 Fortunately, this is easy to do with Django MongoDB Backend. You can use
-the ``showfieldsmap`` management command to generate the schema map
+the ``createencryptedfieldsmap`` management command to generate the schema map
 for your encrypted fields, and then use the results in your settings.
 
-To generate the schema map, run the following command in your Django project:
-::
+To generate the encrypted fields map, run the following command in your Django
+project::
 
-    python manage.py showfieldsmap
+    python manage.py createencryptedfieldsmap
 
-.. note:: The ``showfieldsmap`` command is only available if you have the
-    ``django_mongodb_backend`` app included in the :setting:`INSTALLED_APPS`
-    setting.
+.. note:: The ``createencryptedfieldsmap`` command is only available if you
+   have the ``django_mongodb_backend`` app included in the
+   :setting:`INSTALLED_APPS` setting.
 
 Settings
 ~~~~~~~~
 
-Now include the generated schema map in your Django settings.
-
-::
+Now include the generated schema map in your Django settings::
 
     …
     DATABASES["encrypted"] = {
         …
         "OPTIONS": {
             "auto_encryption_opts": AutoEncryptionOpts(
                 …
-                schema_map= {
+                encrypted_fields_map = {
                     "encryption__patientrecord": {
                         "fields": [
                             {
@@ -119,4 +118,6 @@ Now include the generated schema map in your Django settings.
         …
     }
 
-You are now ready to use client side :doc:`Queryable Encryption </topics/queryable-encryption>` in your Django project.
+You are now ready to use client-side
+:doc:`Queryable Encryption </topics/queryable-encryption>`
+in your Django project.

docs/source/faq.rst

@@ -6,7 +6,9 @@ Configuring Queryable Encryption in Django is similar to
 :doc:`manual:core/queryable-encryption/quick-start` but with some additional
 steps required for Django.
 
-.. note:: This section describes how to configure server side Queryable
+.. admonition:: Server-side Queryable Encryption
+
+   This section describes how to configure server side Queryable
    Encryption in Django. For configuration of client side Queryable Encryption,
    please refer to this :ref:`FAQ question <queryable-encryption>`.
 
@@ -19,7 +21,9 @@ you will need to install PyMongo with Queryable Encryption support::
 
     pip install django-mongodb-backend[encryption]
 
-.. note:: You can use Queryable Encryption on a MongoDB 7.0 or later replica
+.. admonition:: Queryable Encryption Compatibility
+
+   You can use Queryable Encryption on a MongoDB 7.0 or later replica
    set or sharded cluster, but not a standalone instance.
    :ref:`This table <manual:qe-compatibility-reference>` shows which MongoDB
    server products support which Queryable Encryption mechanisms.

docs/source/howto/queryable-encryption.rst

@@ -28,10 +28,10 @@ Available commands
         Defaults to ``default``.
 
 
-``showfieldsmap``
+``createencryptedfieldsmap``
 ----------------------------
 
-.. django-admin:: showfieldsmap
+.. django-admin:: createencryptedfieldsmap
 
     Creates an encrypted fields map that can be used with `encrypted_fields_map`
     in :class:`~pymongo.encryption_options.AutoEncryptionOpts` to configure