Add --create to showencryptedfieldsmap

Author: aclark4life

django_mongodb_backend/management/commands/createencryptedfieldsmap.py renamed to django_mongodb_backend/management/commands/showencryptedfieldsmap.py

@@ -18,9 +18,15 @@ def add_arguments(self, parser):
             help="Specify the database to use for generating the encrypted"
             "fields map. Defaults to the 'default' database.",
         )
+        parser.add_argument(
+            "--create",
+            action="store_true",
+            help="Create the encrypted fields map.",
+        )
 
     def handle(self, *args, **options):
         db = options["database"]
+        create = options.get("create", False)
         connection = connections[db]
         client = connection.connection
         encrypted_fields_map = {}
@@ -32,9 +38,9 @@ def handle(self, *args, **options):
         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
-                        )
+                    from_db = not create
+                    fields = connection.schema_editor()._get_encrypted_fields_map(
+                        model, client, auto_encryption_opts, from_db=from_db
                     )
+                    encrypted_fields_map[model._meta.db_table] = fields
         self.stdout.write(json_util.dumps(encrypted_fields_map, indent=2))

django_mongodb_backend/schema.py

@@ -450,7 +450,7 @@ def _create_collection(self, model):
         else:
             db.create_collection(db_table)
 
-    def _get_encrypted_fields_map(self, model, client, auto_encryption_opts):
+    def _get_encrypted_fields_map(self, model, client, auto_encryption_opts, from_db=False):
         connection = self.connection
         fields = model._meta.fields
         kms_provider = router.kms_provider(model)
@@ -461,16 +461,25 @@ def _get_encrypted_fields_map(self, model, client, auto_encryption_opts):
             client,
             client.codec_options,
         )
+        # Parse key vault namespace
+        key_vault_db, key_vault_coll = auto_encryption_opts._key_vault_namespace.split(".", 1)
+        key_vault_collection = client[key_vault_db][key_vault_coll]
         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],
-                )
+                if from_db:
+                    key_doc = key_vault_collection.find_one({"keyAltNames": key_alt_name})
+                    if not key_doc:
+                        raise ValueError(f"No key found in keyvault for keyAltName={key_alt_name}")
+                    data_key = key_doc["_id"]
+                else:
+                    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,

docs/source/howto/queryable-encryption.rst

@@ -109,15 +109,15 @@ 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 ``createencryptedfieldsmap`` management command to generate the schema map
+the ``showencryptedfieldsmap`` management command to generate the schema map
 for your encrypted fields, and then use the results in your settings.
 
 To generate the encrypted fields map, run the following command in your Django
 project::
 
-    python manage.py createencryptedfieldsmap
+    python manage.py showencryptedfieldsmap
 
-.. note:: The ``createencryptedfieldsmap`` command is only available if you
+.. note:: The ``showencryptedfieldsmap`` command is only available if you
    have the ``django_mongodb_backend`` app included in the
    :setting:`INSTALLED_APPS` setting.
 

docs/source/ref/django-admin.rst

@@ -28,16 +28,19 @@ Available commands
         Defaults to ``default``.
 
 
-``createencryptedfieldsmap``
+``showencryptedfieldsmap``
 ----------------------------
 
-.. django-admin:: createencryptedfieldsmap
+.. django-admin:: showencryptedfieldsmap
 
-    Creates an encrypted fields map that can be used with `encrypted_fields_map`
-    in :class:`~pymongo.encryption_options.AutoEncryptionOpts` to configure
-    client-side Queryable Encryption.
+   Show mapping of models and their encrypted fields.
 
     .. django-admin-option:: --database DATABASE
 
         Specifies the database to use.
         Defaults to ``default``.
+
+    .. django-admin-option:: --create
+
+        If specified, creates the data keys instead of getting them from the
+        database.

tests/encryption_/test_management.py

@@ -107,25 +107,42 @@
     INSTALLED_APPS={"prepend": "django_mongodb_backend"},
 )
 @override_settings(DATABASE_ROUTERS=[TestEncryptedRouter()])
-class EncryptedFieldTests(TransactionTestCase):
+class EncryptedFieldsManagementCommandTests(TransactionTestCase):
     databases = {"default", "encrypted"}
     available_apps = ["django_mongodb_backend", "encryption_"]
 
+    def _compare_json(self, json1, json2):
+        # Remove keyIds since they are different for each run.
+        for table in json2:
+            for field in json2[table]["fields"]:
+                del field["keyId"]
+        # TODO: probably we don't need to test the entire mapping, otherwise it
+        # requires updates every time a new model or field is added!
+        self.assertEqual(json1, json2)
+
+    def test_show_encrypted_fields_map(self):
+        self.maxDiff = None
+        out = StringIO()
+        call_command(
+            "showencryptedfieldsmap",
+            "--database",
+            "encrypted",
+            verbosity=0,
+            stdout=out,
+        )
+        output_json = json_util.loads(out.getvalue())
+        self._compare_json(EXPECTED_ENCRYPTED_FIELDS_MAP, output_json)
+
     def test_create_encrypted_fields_map(self):
         self.maxDiff = None
         out = StringIO()
         call_command(
-            "createencryptedfieldsmap",
+            "showencryptedfieldsmap",
             "--database",
             "encrypted",
+            "--create",
             verbosity=0,
             stdout=out,
         )
-        # Remove keyIds since they are different for each run.
         output_json = json_util.loads(out.getvalue())
-        for table in output_json:
-            for field in output_json[table]["fields"]:
-                del field["keyId"]
-        # TODO: probably we don't need to test the entire mapping, otherwise it
-        # requires updates every time a new model or field is added!
-        self.assertEqual(EXPECTED_ENCRYPTED_FIELDS_MAP, output_json)
+        self._compare_json(EXPECTED_ENCRYPTED_FIELDS_MAP, output_json)