Remove EncryptedModel, check for encrypted fields

aclark4life · aclark4life · commit 1af38d858c45 · 2025-07-31T20:31:32.000-04:00
diff --git a/django_mongodb_backend/encryption.py b/django_mongodb_backend/encryption.py
@@ -2,6 +2,8 @@
 
 import os
 
+from .fields import has_encrypted_fields
+
 KMS_CREDENTIALS = {
     "aws": {
         "key": os.getenv("AWS_KEY_ARN", ""),
@@ -44,11 +46,11 @@
 class EncryptedRouter:
     def allow_migrate(self, db, app_label, model_name=None, model=None, **hints):
         if model:
-            return db == ("other" if getattr(model, "encrypted", False) else "default")
+            return db == ("other" if has_encrypted_fields(model) else "default")
         return db == "default"
 
     def db_for_read(self, model, **hints):
-        if getattr(model, "encrypted", False):
+        if has_encrypted_fields(model):
             return "other"
         return "default"
 
diff --git a/django_mongodb_backend/fields/__init__.py b/django_mongodb_backend/fields/__init__.py
@@ -3,7 +3,7 @@
 from .duration import register_duration_field
 from .embedded_model import EmbeddedModelField
 from .embedded_model_array import EmbeddedModelArrayField
-from .encrypted_model import (
+from .encryption import (
     EncryptedBigIntegerField,
     EncryptedBinaryField,
     EncryptedBooleanField,
@@ -19,11 +19,13 @@
     EncryptedTextField,
     EncryptedTimeField,
     EncryptedURLField,
+    has_encrypted_fields,
 )
 from .json import register_json_field
 from .objectid import ObjectIdField
 
 __all__ = [
+    "has_encrypted_fields",
     "register_fields",
     "ArrayField",
     "EmbeddedModelArrayField",
diff --git a/django_mongodb_backend/fields/encryption.py b/django_mongodb_backend/fields/encryption.py
@@ -1,6 +1,10 @@
 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
 
diff --git a/django_mongodb_backend/models.py b/django_mongodb_backend/models.py
@@ -14,10 +14,3 @@ 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
diff --git a/django_mongodb_backend/schema.py b/django_mongodb_backend/schema.py
@@ -5,7 +5,7 @@
 from pymongo.encryption import ClientEncryption
 from pymongo.operations import SearchIndexModel
 
-from .fields import EmbeddedModelField
+from .fields import EmbeddedModelField, has_encrypted_fields
 from .indexes import SearchIndex
 from .query import wrap_database_errors
 from .utils import OperationCollector
@@ -432,7 +432,7 @@ def _create_collection(self, model):
         """
         db = self.get_database()
         db_table = model._meta.db_table
-        if getattr(model, "encrypted", False):
+        if has_encrypted_fields(model):
             client = self.connection.connection
             options = getattr(client._options, "auto_encryption_opts", None)
             if options is not None:
diff --git a/docs/source/howto/encryption.rst b/docs/source/howto/encryption.rst
@@ -19,7 +19,7 @@ In addition to :doc:`installing </intro/install>` and
 :doc:`configuring </intro/configure>` Django MongoDB Backend,
 you will need to install PyMongo with Queryable Encryption support::
 
-    pip install pymongo[aws,encryption]
+    pip install django-mongodb-backend[encryption]
 
 Settings
 --------
@@ -28,7 +28,7 @@ Add an encrypted database, encrypted database router and KMS credentials to
 your Django settings.
 
 .. note:: Use of the helpers provided in ``django_mongodb_backend.encryption``
-    requires an encrypted database named "my_encrypted_database".
+    requires an encrypted database named "other".
 
 ::
 
@@ -40,21 +40,21 @@ your Django settings.
             MONGODB_URI,
             db_name="my_database",
         ),
-        "my_encrypted_database": parse_uri(
+        "other": parse_uri(
             MONGODB_URI,
-            db_name="my_encrypted_database",
+            db_name="other",
             options={
                 "auto_encryption_opts": AutoEncryptionOpts(
                     kms_providers=encryption.KMS_PROVIDERS,
-                    key_vault_namespace="my_encrypted_database.keyvault",
+                    key_vault_namespace="other.keyvault",
                 )
             },
         ),
 
-    DATABASES["my_encrypted_database"]["KMS_CREDENTIALS"] = encryption.KMS_CREDENTIALS
+    DATABASES["other"]["KMS_CREDENTIALS"] = encryption.KMS_CREDENTIALS
     DATABASE_ROUTERS = [encryption.EncryptedRouter()]
 
-You are now ready to use :doc:`encrypted models </topics/encrypted-models>` in your Django project.
+You are now ready to use :doc:`encrypted fields </topics/encrypted-fields>` in your Django project.
 
 
 Helper classes and settings
diff --git a/docs/source/index.rst b/docs/source/index.rst
@@ -45,7 +45,7 @@ Models
 **Topic guides:**
 
 - :doc:`topics/embedded-models`
-- :doc:`topics/encrypted-models`
+- :doc:`topics/encrypted-fields`
 
 Forms
 =====
diff --git a/docs/source/intro/configure.rst b/docs/source/intro/configure.rst
@@ -175,15 +175,15 @@ present.)
 Queryable Encryption
 --------------------
 
-If you intend to use :doc:`encrypted models </topics/encrypted-models>`, you may
+If you intend to use :doc:`encrypted fields </topics/encrypted-fields>`, you may
 optionally configure the :setting:`DATABASE_ROUTERS` setting so that collections
-for encrypted models are created in an encrypted database.
+for encrypted fields are created in an encrypted database.
 
 `Router configuration <https://docs.djangoproject.com/en/stable/topics/db/multi-db/#database-routers>`__
 is unique to each project and beyond the scope of Django MongoDB Backend, but an
 example is included in the :doc:`encryption helpers </howto/encryption>`
-that routes collection operations for encrypted models to a database named
-"my_encrypted_database"::
+that routes collection operations for encrypted fields to a database named
+"other"::
 
     DATABASE_ROUTERS = ["django_mongodb_backend.encryption.EncryptedRouter"]
 
diff --git a/docs/source/ref/django-admin.rst b/docs/source/ref/django-admin.rst
@@ -33,12 +33,12 @@ Available commands
 
 .. django-admin:: get_encrypted_fields_map
 
-    Creates a schema map for the encrypted fields in your encrypted models. This
+    Creates a schema map for the encrypted fields in your models. This
     map can be provided to
     :class:`~pymongo.encryption_options.AutoEncryptionOpts` for use with
-    production deployments of :class:`~pymongo.encryption.ClientEncryption`.
+    :class:`~pymongo.encryption.ClientEncryption`.
 
     .. django-admin-option:: --database DATABASE
 
-        Specifies the database to use to generate an encrypted fields map
-        for all encrypted models. Defaults to ``default``.
+        Specifies the database to use to generate an encrypted fields map.
+        Defaults to ``default``.
diff --git a/docs/source/ref/models/fields.rst b/docs/source/ref/models/fields.rst
@@ -349,7 +349,7 @@ they encrypt the data before storing it in the database.
 | ``EncryptedURLField``                 | :class:`~django.db.models.URLField`              |
 +---------------------------------------+--------------------------------------------------+
 
-.. _encrypted-model-unsupported-fields:
+.. _encrypted-fields-unsupported-fields:
 
 Unsupported fields
 ~~~~~~~~~~~~~~~~~~
diff --git a/docs/source/ref/models/models.rst b/docs/source/ref/models/models.rst
@@ -17,8 +17,3 @@ Two MongoDB-specific models are available in ``django_mongodb_backend.models``.
 
     Embedded model instances won't have a value for their primary key unless
     one is explicitly set.
-
-.. class:: EncryptedModel
-
-    An abstract model which all :doc:`encrypted models </topics/encrypted-models>`
-    must subclass.
diff --git a/docs/source/topics/encrypted-fields.rst b/docs/source/topics/encrypted-fields.rst
@@ -1,28 +1,24 @@
-Encrypted models
+Encrypted fields
 ================
 
 Use :class:`~django_mongodb_backend.models.EncryptedModel` and
 :mod:`~django_mongodb_backend.fields` to structure
 your data using `Queryable Encryption <https://www.mongodb.com/docs/manual/core/queryable-encryption/>`_.
 
-.. _encrypted-model-field-example:
-
-``EncryptedModelField``
------------------------
+.. _encrypted-field-example:
 
 The basics
-~~~~~~~~~~
+----------
 
 Let's consider this example::
 
     from django.db import models
 
-    from django_mongodb_backend.models import EncryptedModel
     from django_mongodb_backend.fields import EncryptedCharField
     from django_mongodb_backend.encryption import EqualityQuery
 
 
-    class Patient(EncryptedModel):
+    class Patient(models.Model):
         ssn = EncryptedCharField(max_length=11, queries=EqualityQuery())
 
         def __str__(self):
@@ -55,10 +51,10 @@ the patient data looks like this:
       ssn: Binary.createFromBase64('DkrbD67ejkt2u…', 6),
     }
 
-Querying Encrypted Models
-~~~~~~~~~~~~~~~~~~~~~~~~~
+Querying Encrypted Fields
+-------------------------
 
-You can query encrypted fields in encrypted models using a `limited set of
+You can query encrypted fields using a `limited set of
 query operators <https://www.mongodb.com/docs/manual/core/queryable-encryption/reference/supported-operations/#std-label-qe-supported-query-operators>`_
 which must be specified in the field definition. For example, to query the ``ssn`` field for equality, you can use the
 ``EqualityQuery`` operator as shown in the example above.
diff --git a/docs/source/topics/index.rst b/docs/source/topics/index.rst
@@ -10,5 +10,5 @@ know:
 
    cache
    embedded-models
-   encrypted-models
+   encrypted-fields
    known-issues
diff --git a/docs/source/topics/known-issues.rst b/docs/source/topics/known-issues.rst
@@ -106,4 +106,4 @@ Consider these
 `limitations and restrictions <https://www.mongodb.com/docs/manual/core/queryable-encryption/reference/limitations/>`_
 before enabling Queryable Encryption. Some operations are unsupported, and others behave differently.
 
-Also see :ref:`unsupported fields <encrypted-model-unsupported-fields>`.
+Also see :ref:`unsupported fields <encrypted-fields-unsupported-fields>`.
diff --git a/tests/encryption_/models.py b/tests/encryption_/models.py
@@ -1,3 +1,5 @@
+from django.db import models
+
 from django_mongodb_backend.encryption import EqualityQuery, RangeQuery
 from django_mongodb_backend.fields import (
     EncryptedBigIntegerField,
@@ -15,17 +17,16 @@
     EncryptedTimeField,
     EncryptedURLField,
 )
-from django_mongodb_backend.models import EncryptedModel
 
 
-class Appointment(EncryptedModel):
+class Appointment(models.Model):
     time = EncryptedTimeField(queries=EqualityQuery())
 
     class Meta:
         required_db_features = {"supports_queryable_encryption"}
 
 
-class Billing(EncryptedModel):
+class Billing(models.Model):
     cc_type = EncryptedCharField(max_length=20, queries=EqualityQuery())
     cc_number = EncryptedBigIntegerField(queries=EqualityQuery())
     account_balance = EncryptedDecimalField(max_digits=10, decimal_places=2, queries=RangeQuery())
@@ -34,15 +35,15 @@ class Meta:
         required_db_features = {"supports_queryable_encryption"}
 
 
-class PatientPortalUser(EncryptedModel):
+class PatientPortalUser(models.Model):
     ip_address = EncryptedGenericIPAddressField(queries=EqualityQuery())
     url = EncryptedURLField(queries=EqualityQuery())
 
     class Meta:
         required_db_features = {"supports_queryable_encryption"}
 
 
-class PatientRecord(EncryptedModel):
+class PatientRecord(models.Model):
     ssn = EncryptedCharField(max_length=11, queries=EqualityQuery())
     birth_date = EncryptedDateField(queries=RangeQuery())
     profile_picture = EncryptedBinaryField(queries=EqualityQuery())
@@ -56,7 +57,7 @@ class Meta:
         required_db_features = {"supports_queryable_encryption"}
 
 
-class Patient(EncryptedModel):
+class Patient(models.Model):
     patient_id = EncryptedIntegerField("patient_id", queries=EqualityQuery())
     patient_name = EncryptedCharField(max_length=100)
     patient_notes = EncryptedTextField(queries=EqualityQuery())
diff --git a/tests/encryption_/routers.py b/tests/encryption_/routers.py
@@ -1,16 +1,31 @@
+from django_mongodb_backend.fields import has_encrypted_fields
+
+
 class TestEncryptedRouter:
     """Router for testing encrypted models in Django. `kms_provider`
     must be set on the global test router since table creation happens
     at the start of the test suite, before @override_settings(
     DATABASE_ROUTERS=[TestEncryptedRouter()]) takes effect.
     """
 
+    # def allow_migrate(self, db, app_label, model_name=None, model=None, **hints):
+    #     return getattr(model, "encrypted", False)
+
+    # def db_for_read(self, model, **hints):
+    #     if getattr(model, "encrypted", False):
+    #         return "other"
+    #     return None
+
+    # db_for_write = db_for_read
+
     def allow_migrate(self, db, app_label, model_name=None, model=None, **hints):
-        return getattr(model, "encrypted", False)
+        if model:
+            return db == ("other" if has_encrypted_fields(model) else "default")
+        return db == "default"
 
     def db_for_read(self, model, **hints):
-        if getattr(model, "encrypted", False):
+        if has_encrypted_fields(model):
             return "other"
-        return None
+        return "default"
 
     db_for_write = db_for_read
diff --git a/tests/encryption_/tests.py b/tests/encryption_/tests.py
@@ -13,7 +13,6 @@
 
 from django_mongodb_backend.encryption import EqualityQuery
 from django_mongodb_backend.fields import EncryptedFieldMixin
-from django_mongodb_backend.models import EncryptedModel
 
 from .models import (
     Appointment,
@@ -264,7 +263,7 @@ def test_appointment(self):
         with self.assertRaises(AssertionError):  # noqa: SIM117
             with self.assertRaises(pymongo.errors.OperationFailure):
 
-                class DurationFieldTest(EncryptedModel):
+                class DurationFieldTest(models.Model):
                     duration = EncryptedDurationField(EqualityQuery())
 
                 appointment = DurationFieldTest(duration=timedelta(hours=2, minutes=30))
@@ -301,7 +300,7 @@ def test_patientportaluser(self):
         with self.assertRaises(AssertionError):  # noqa: SIM117
             with self.assertRaises(pymongo.errors.OperationFailure):
 
-                class SlugFieldTest(EncryptedModel):
+                class SlugFieldTest(models.Model):
                     slug = EncryptedSlugField(EqualityQuery())
 
     def test_patientrecord(self):
