Code review fixes (1/x)

Author: aclark4life

django_mongodb_backend/fields/encryption.py

@@ -1,7 +1,7 @@
 from django.db import models
 
 
-class EncryptedFieldMixin(models.Field):
+class EncryptedFieldMixin:
     encrypted = True
 
     def __init__(self, *args, queries=None, **kwargs):

django_mongodb_backend/management/commands/showencryptedfieldsmap.py

@@ -8,10 +8,11 @@
 
 class Command(BaseCommand):
     help = """
-    This command shows the mapping of encrypted fields to attributes
-    including  data type, data keys and query types. It can be used to set the
-    ``encrypted_fields_map`` in ``AutoEncryptionOpts``. Defaults to showing
-    existing keys from the configured key vault.
+    Shows the mapping of encrypted fields to field attributes, including data
+    type, data keys and query types. The output can be used to set
+    ``encrypted_fields_map`` in ``AutoEncryptionOpts``.
+
+    Defaults to showing keys from the ``key_vault_namespace`` collection.
     """
 
     def add_arguments(self, parser):

docs/source/ref/models/encrypted-fields.rst

@@ -63,7 +63,7 @@ supported by Queryable Encryption.
 ``EncryptedFieldMixin``
 =======================
 
-.. class:: EncryptedFieldMixin(models.Field)
+.. class:: EncryptedFieldMixin
 
     .. versionadded:: 5.2.0rc1
 

tests/encryption_/models.py

@@ -11,7 +11,6 @@
     EncryptedFloatField,
     EncryptedGenericIPAddressField,
     EncryptedIntegerField,
-    EncryptedPositiveSmallIntegerField,
     EncryptedSmallIntegerField,
     EncryptedTextField,
     EncryptedTimeField,
@@ -22,45 +21,42 @@
 RANGE_QUERY = {"queryType": "range"}
 
 
-class Appointment(models.Model):
-    time = EncryptedTimeField(queries=EQUALITY_QUERY)
-
+class QueryableEncryptionModelBase(models.Model):
     class Meta:
+        abstract = True
         required_db_features = {"supports_queryable_encryption"}
 
 
-class Billing(models.Model):
-    cc_type = EncryptedCharField(max_length=20, queries=EQUALITY_QUERY)
-    cc_number = EncryptedIntegerField(queries=EQUALITY_QUERY)
+class Appointment(QueryableEncryptionModelBase):
+    time = EncryptedTimeField(queries=EQUALITY_QUERY)
+
+
+class Billing(QueryableEncryptionModelBase):
     account_balance = EncryptedDecimalField(max_digits=10, decimal_places=2, queries=RANGE_QUERY)
 
-    class Meta:
-        required_db_features = {"supports_queryable_encryption"}
+
+class CreditCard(QueryableEncryptionModelBase):
+    cc_type = EncryptedCharField(max_length=20, queries=EQUALITY_QUERY)
+    cc_number = EncryptedIntegerField(queries=EQUALITY_QUERY)
 
 
-class PatientPortalUser(models.Model):
+class PatientPortalUser(QueryableEncryptionModelBase):
     ip_address = EncryptedGenericIPAddressField(queries=EQUALITY_QUERY)
     url = EncryptedURLField(queries=EQUALITY_QUERY)
 
-    class Meta:
-        required_db_features = {"supports_queryable_encryption"}
-
 
-class PatientRecord(models.Model):
+class PatientRecord(QueryableEncryptionModelBase):
     ssn = EncryptedCharField(max_length=11, queries=EQUALITY_QUERY)
     birth_date = EncryptedDateField(queries=RANGE_QUERY)
     profile_picture = EncryptedBinaryField(queries=EQUALITY_QUERY)
     patient_age = EncryptedSmallIntegerField(queries={**RANGE_QUERY, "min": 0, "max": 100})
     weight = EncryptedFloatField(queries=RANGE_QUERY)
 
     # TODO: Embed Billing model
-    # billing =
+    # billing = EncryptedEmbeddedField(Billing)
 
-    class Meta:
-        required_db_features = {"supports_queryable_encryption"}
 
-
-class Patient(models.Model):
+class Patient(QueryableEncryptionModelBase):
     patient_id = EncryptedIntegerField(queries=EQUALITY_QUERY)
     patient_name = EncryptedCharField(max_length=100)
     patient_notes = EncryptedTextField(queries=EQUALITY_QUERY)
@@ -69,27 +65,4 @@ class Patient(models.Model):
     email = EncryptedEmailField(queries=EQUALITY_QUERY)
 
     # TODO: Embed PatientRecord model
-    # patient_record =
-
-    class Meta:
-        required_db_features = {"supports_queryable_encryption"}
-
-
-class EncryptedNumbers(models.Model):
-    # Not tested elsewhere
-    pos_smallint = EncryptedPositiveSmallIntegerField(queries=EQUALITY_QUERY)
-    smallint = EncryptedSmallIntegerField(queries=EQUALITY_QUERY)
-
-    class Meta:
-        required_db_features = {"supports_queryable_encryption"}
-
-
-class SensitiveData(models.Model):
-    # Example from documentation
-    name = EncryptedCharField(max_length=100)
-    email = EncryptedEmailField()
-    phone_number = EncryptedCharField(max_length=15)
-
-    sensitive_text = EncryptedTextField()
-    sensitive_integer = EncryptedIntegerField()
-    sensitive_date = EncryptedDateField()
+    # patient_record = EncryptedEmbeddedField(PatientRecord)

tests/encryption_/test_base.py

@@ -1,52 +1,10 @@
-from datetime import datetime
-
 from django.test import TransactionTestCase, override_settings, skipUnlessDBFeature
 
-from .models import (
-    Appointment,
-    Billing,
-    EncryptedNumbers,
-    Patient,
-    PatientPortalUser,
-    PatientRecord,
-)
 from .routers import TestEncryptedRouter
 
 
 @skipUnlessDBFeature("supports_queryable_encryption")
 @override_settings(DATABASE_ROUTERS=[TestEncryptedRouter()])
-class QueryableEncryptionTests(TransactionTestCase):
+class QueryableEncryptionTestCase(TransactionTestCase):
     databases = {"default", "encrypted"}
     available_apps = ["encryption_"]
-
-    def setUp(self):
-        """
-        Used in schema and field tests.
-        """
-        self.appointment = Appointment.objects.create(time="8:00")
-        self.billing = Billing.objects.create(
-            cc_type="Visa", cc_number=1234567890123456, account_balance=100.50
-        )
-        self.portal_user = PatientPortalUser.objects.create(
-            ip_address="127.0.0.1",
-            url="https://example.com",
-        )
-        self.patientrecord = PatientRecord.objects.create(
-            ssn="123-45-6789",
-            birth_date="1970-01-01",
-            profile_picture=b"image data",
-            weight=175.5,
-            patient_age=47,
-        )
-        self.patient = Patient.objects.create(
-            patient_id=1,
-            patient_name="John Doe",
-            patient_notes="patient notes " * 25,
-            registration_date=datetime(2023, 10, 1, 12, 0, 0),
-            is_active=True,
-            email="john.doe@example.com",
-        )
-        EncryptedNumbers.objects.create(
-            pos_smallint=12345,
-            smallint=-12345,
-        )

tests/encryption_/test_connection.py

@@ -1,7 +1,7 @@
-from .test_base import QueryableEncryptionTests
+from .test_base import QueryableEncryptionTestCase
 
 
-class TestConnection(QueryableEncryptionTests):
+class TestConnection(QueryableEncryptionTestCase):
     def test_connection(self):
         # raise ImproperlyConfigured(
         #       "Encrypted fields found but "

tests/encryption_/test_fields.py

@@ -1,4 +1,4 @@
-from datetime import datetime, time
+from datetime import date, datetime, time
 
 import pymongo
 from bson.binary import Binary
@@ -8,43 +8,97 @@
 from .models import (
     Appointment,
     Billing,
-    EncryptedNumbers,
+    CreditCard,
     Patient,
     PatientPortalUser,
     PatientRecord,
 )
-from .test_base import QueryableEncryptionTests
+from .test_base import QueryableEncryptionTestCase
 
 
-class FieldTests(QueryableEncryptionTests):
-    def test_appointment(self):
-        self.assertEqual(Appointment.objects.get(time="8:00").time, time(8, 0))
+class FieldTests(QueryableEncryptionTestCase):
+    def setUp(self):
+        Patient.objects.create(
+            patient_id=1,
+            patient_name="John Doe",
+            patient_notes="patient notes " * 25,
+            registration_date=datetime(2023, 10, 1, 12, 0, 0),
+            is_active=True,
+            email="john.doe@example.com",
+        )
+        PatientRecord.objects.create(
+            ssn="123-45-6789",
+            birth_date="1969-01-01",
+            profile_picture=b"image data",
+            patient_age=50,
+            weight=180.0,
+        )
+
+    def test_binaryfield(self):
+        self.assertEqual(
+            PatientRecord.objects.get(profile_picture=b"image data").profile_picture, b"image data"
+        )
+
+    def test_booleanfield(self):
+        self.assertTrue(Patient.objects.get(patient_id=1).is_active)
+
+    def test_charfield(self):
+        CreditCard.objects.create(cc_type="Visa", cc_number="1234567890123456")
+        self.assertEqual(CreditCard.objects.get(cc_type="Visa").cc_type, "Visa")
+        self.assertEqual(PatientRecord.objects.get(ssn="123-45-6789").ssn, "123-45-6789")
 
-    def test_billing(self):
+    def test_datefield(self):
         self.assertEqual(
-            Billing.objects.get(cc_number=1234567890123456).cc_number, 1234567890123456
+            PatientRecord.objects.get(birth_date="1969-1-1").birth_date, date(1969, 1, 1)
         )
-        self.assertEqual(Billing.objects.get(cc_type="Visa").cc_type, "Visa")
+
+    def test_datetimefield(self):
+        self.assertEqual(
+            Patient.objects.get(
+                registration_date=datetime(2023, 10, 1, 12, 0, 0)
+            ).registration_date,
+            datetime(2023, 10, 1, 12, 0, 0),
+        )
+
+    def test_decimalfield(self):
+        Billing.objects.create(account_balance=100.50)
         self.assertTrue(Billing.objects.filter(account_balance__gte=100.0).exists())
 
-    def test_patientportaluser(self):
+    def test_emailfield(self):
         self.assertEqual(
-            PatientPortalUser.objects.get(ip_address="127.0.0.1").ip_address, "127.0.0.1"
+            Patient.objects.get(email="john.doe@example.com").email, "john.doe@example.com"
         )
 
-    def test_patientrecord(self):
-        self.assertEqual(PatientRecord.objects.get(ssn="123-45-6789").ssn, "123-45-6789")
-        with self.assertRaises(PatientRecord.DoesNotExist):
-            PatientRecord.objects.get(ssn="000-00-0000")
-        self.assertTrue(PatientRecord.objects.filter(birth_date__gte="1969-01-01").exists())
+    def test_floatfield(self):
+        self.assertTrue(PatientRecord.objects.filter(weight__gte=175.0).exists())
+
+    def test_integerfield(self):
+        CreditCard.objects.create(cc_type="Visa", cc_number="1234567890123456")
+        self.assertEqual(
+            CreditCard.objects.get(cc_number=1234567890123456).cc_number, 1234567890123456
+        )
+
+    def test_ipaddressfield(self):
+        PatientPortalUser.objects.create(ip_address="127.0.0.1", url="https://example.com")
         self.assertEqual(
-            PatientRecord.objects.get(ssn="123-45-6789").profile_picture, b"image data"
+            PatientPortalUser.objects.get(ip_address="127.0.0.1").ip_address, "127.0.0.1"
         )
+
+    def test_smallintegerfield(self):
         self.assertTrue(PatientRecord.objects.filter(patient_age__gte=40).exists())
         self.assertFalse(PatientRecord.objects.filter(patient_age__gte=80).exists())
-        self.assertTrue(PatientRecord.objects.filter(weight__gte=175.0).exists())
 
-        # Test encrypted patient record in unencrypted database.
+    def test_timefield(self):
+        Appointment.objects.create(time="8:00")
+        self.assertEqual(Appointment.objects.get(time="8:00").time, time(8, 0))
+
+    def test_encrypted_patient_record_in_encrypted_database(self):
+        patients = connections["encrypted"].database.encryption__patient.find()
+        self.assertEqual(len(list(patients)), 1)
+        records = connections["encrypted"].database.encryption__patientrecord.find()
+        self.assertTrue("__safeContent__" in records[0])
+
+    def test_encrypted_patient_record_in_unencrypted_database(self):
         conn_params = connections["encrypted"].get_connection_params()
         db_name = settings.DATABASES["encrypted"]["NAME"]
         if conn_params.pop("auto_encryption_opts", False):
@@ -56,32 +110,8 @@ def test_patientrecord(self):
                 ssn = patientrecords[0]["ssn"]
                 self.assertTrue(isinstance(ssn, Binary))
 
-    def test_patient(self):
+    def test_textfield(self):
         self.assertEqual(
             Patient.objects.get(patient_notes="patient notes " * 25).patient_notes,
             "patient notes " * 25,
         )
-        self.assertEqual(
-            Patient.objects.get(
-                registration_date=datetime(2023, 10, 1, 12, 0, 0)
-            ).registration_date,
-            datetime(2023, 10, 1, 12, 0, 0),
-        )
-        self.assertTrue(Patient.objects.get(patient_id=1).is_active)
-        self.assertEqual(
-            Patient.objects.get(email="john.doe@example.com").email, "john.doe@example.com"
-        )
-
-        # Test decrypted patient record in encrypted database.
-        patients = connections["encrypted"].database.encryption__patient.find()
-        self.assertEqual(len(list(patients)), 1)
-        records = connections["encrypted"].database.encryption__patientrecord.find()
-        self.assertTrue("__safeContent__" in records[0])
-
-    def test_pos_small_int(self):
-        obj = EncryptedNumbers.objects.get(pos_smallint=12345)
-        self.assertEqual(obj.pos_smallint, 12345)
-
-    def test_small_int(self):
-        obj = EncryptedNumbers.objects.get(smallint=-12345)
-        self.assertEqual(obj.smallint, -12345)

tests/encryption_/test_schema.py

@@ -1,9 +1,10 @@
 from django.db import connections
 
-from .test_base import QueryableEncryptionTests
+from .models import Patient
+from .test_base import QueryableEncryptionTestCase
 
 
-class SchemaTests(QueryableEncryptionTests):
+class SchemaTests(QueryableEncryptionTestCase):
     maxDiff = None
 
     def test_get_encrypted_fields_map(self):
@@ -53,7 +54,7 @@ def test_get_encrypted_fields_map(self):
         connection = connections["encrypted"]
         with connection.schema_editor() as editor:
             client = connection.connection
-            encrypted_fields_map = editor._get_encrypted_fields_map(self.patient, client)
+            encrypted_fields_map = editor._get_encrypted_fields_map(Patient, client)
             for field in encrypted_fields_map["fields"]:
                 # Remove data keys from the output; they are expected to differ
                 field.pop("keyId", None)