Batch with additional historical model fields (#1248)

* docs * add history attribute to batch * doc updates * uses `custom_historical_attrs` not kwargs * updates docs with PollWithHistoricalSessionAttr and session attr * changes example values to not be same as other examples * adds PollWithHistoricalSessionAttr for testing * updates test for PollWithHistoricalSessionAttr * adds custom_historical_attrs to test_bulk_create_no_ids_return test * corrected change list * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * updates `SessionsHistoricalModel` with explicit default=None * updates CustomHistoricalAttrsTest test to use TestCase setUp * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fixes: id value can not be 0 * Update CHANGES.rst Co-authored-by: Anders <[email protected]> * Update simple_history/tests/tests/test_utils.py Co-authored-by: Anders <[email protected]> * Update docs/common_issues.rst Co-authored-by: Anders <[email protected]> * Update docs/common_issues.rst Co-authored-by: Anders <[email protected]> * Update docs/common_issues.rst Co-authored-by: Anders <[email protected]> * Update docs/common_issues.rst Co-authored-by: Anders <[email protected]> * Update simple_history/tests/tests/test_utils.py Co-authored-by: Anders <[email protected]> * Update simple_history/tests/tests/test_utils.py Co-authored-by: Anders <[email protected]> * Update simple_history/tests/tests/test_utils.py Co-authored-by: Anders <[email protected]> * Update simple_history/tests/tests/test_utils.py Co-authored-by: Anders <[email protected]> * Update simple_history/tests/tests/test_utils.py Co-authored-by: Anders <[email protected]> * Update simple_history/manager.py Co-authored-by: Anders <[email protected]> * updates docstring --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Anders <[email protected]>
jazzband · Sep 25, 2023 · 1032b40 · 1032b40
1 parent bde32e1
commit 1032b40
Show file tree

Hide file tree

Showing 7 changed files with 120 additions and 0 deletions.
diff --git a/AUTHORS.rst b/AUTHORS.rst
@@ -107,6 +107,7 @@ Authors
 - Nathan Villagaray-Carski (`ncvc <https://github.com/ncvc>`_)
 - Nianpeng Li
 - Nick Träger
+- Noel James (`NoelJames <https://github.com/NoelJames>`_)
 - Phillip Marshall
 - Prakash Venkatraman (`dopatraman <https://github.com/dopatraman>`_)
 - Rajesh Pappula

diff --git a/CHANGES.rst b/CHANGES.rst
@@ -7,6 +7,10 @@ Unreleased
 - Allow ``HistoricalRecords.m2m_fields`` as str (gh-1243)
 - Fixed ``HistoryRequestMiddleware`` deleting non-existent
   ``HistoricalRecords.context.request`` in very specific circumstances (gh-1256)
+- Added ``custom_historical_attrs`` to ``bulk_create_with_history()`` and
+  ``bulk_update_with_history()`` for setting additional fields on custom history models
+  (gh-1248)
+
 
 3.4.0 (2023-08-18)
 ------------------

diff --git a/docs/common_issues.rst b/docs/common_issues.rst
@@ -55,6 +55,27 @@ You can also specify a default user or default change reason responsible for the
     >>> Poll.history.get(id=data[0].id).history_user == user
     True
 
+If you're using `additional fields in historical models`_ and have custom fields to
+batch-create into the history, pass the optional dict argument ``custom_historical_attrs``
+containing the field names and values.
+A field ``session`` would be passed as ``custom_historical_attrs={'session': 'training'}``.
+
+.. _additional fields in historical models: historical_model.html#adding-additional-fields-to-historical-models
+
+.. code-block:: pycon
+
+    >>> from simple_history.tests.models import PollWithHistoricalSessionAttr
+    >>> data = [
+        PollWithHistoricalSessionAttr(id=x, question=f'Question {x}')
+        for x in range(10)
+    ]
+    >>> objs = bulk_create_with_history(
+            data, PollWithHistoricalSessionAttr,
+            custom_historical_attrs={'session': 'training'}
+        )
+    >>> data[0].history.get().session
+    'training'
+
 Bulk Updating a Model with History (New)
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
@@ -88,6 +109,22 @@ default manager returns a filtered set), you can specify which manager to use wi
     >>> data = [PollWithAlternativeManager(id=x, question='Question ' + str(x), pub_date=now()) for x in range(1000)]
     >>> objs = bulk_create_with_history(data, PollWithAlternativeManager, batch_size=500, manager=PollWithAlternativeManager.all_polls)
 
+If you're using `additional fields in historical models`_ and have custom fields to
+batch-update into the history, pass the optional dict argument ``custom_historical_attrs``
+containing the field names and values.
+A field ``session`` would be passed as ``custom_historical_attrs={'session': 'jam'}``.
+
+.. _additional fields in historical models: historical_model.html#adding-additional-fields-to-historical-models
+
+.. code-block:: pycon
+
+    >>> bulk_update_with_history(
+            data, PollWithHistoricalSessionAttr,
+            custom_historical_attrs={'session': 'jam'}
+        )
+    >>> data[0].history.latest().session
+    'jam'
+
 QuerySet Updates with History (Updated in Django 2.2)
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 Unlike with ``bulk_create``, `queryset updates`_ perform an SQL update query on

diff --git a/simple_history/manager.py b/simple_history/manager.py
@@ -230,6 +230,7 @@ def bulk_history_create(
         default_user=None,
         default_change_reason="",
         default_date=None,
+        custom_historical_attrs=None,
     ):
         """
         Bulk create the history for the objects specified by objs.
@@ -262,6 +263,7 @@ def bulk_history_create(
                     field.attname: getattr(instance, field.attname)
                     for field in self.model.tracked_fields
                 },
+                **(custom_historical_attrs or {}),
             )
             if hasattr(self.model, "history_relation"):
                 row.history_relation_id = instance.pk

diff --git a/simple_history/tests/models.py b/simple_history/tests/models.py
@@ -125,6 +125,18 @@ def get_absolute_url(self):
         return reverse("poll-detail", kwargs={"pk": self.pk})
 
 
+class SessionsHistoricalModel(models.Model):
+    session = models.CharField(max_length=200, null=True, default=None)
+
+    class Meta:
+        abstract = True
+
+
+class PollWithHistoricalSessionAttr(models.Model):
+    question = models.CharField(max_length=200)
+    history = HistoricalRecords(bases=[SessionsHistoricalModel])
+
+
 class PollWithManyToMany(models.Model):
     question = models.CharField(max_length=200)
     pub_date = models.DateTimeField("date published")

diff --git a/simple_history/tests/tests/test_utils.py b/simple_history/tests/tests/test_utils.py
@@ -16,12 +16,14 @@
     Poll,
     PollWithAlternativeManager,
     PollWithExcludeFields,
+    PollWithHistoricalSessionAttr,
     PollWithUniqueQuestion,
     Street,
 )
 from simple_history.utils import (
     bulk_create_with_history,
     bulk_update_with_history,
+    get_history_manager_for_model,
     update_change_reason,
 )
 
@@ -288,6 +290,7 @@ def test_bulk_create_no_ids_return(self, hist_manager_mock):
             default_user=None,
             default_change_reason=None,
             default_date=None,
+            custom_historical_attrs=None,
         )
 
 
@@ -509,6 +512,58 @@ def test_bulk_update_history_wrong_manager(self):
             )
 
 
+class CustomHistoricalAttrsTest(TestCase):
+    def setUp(self):
+        self.data = [
+            PollWithHistoricalSessionAttr(id=x, question=f"Question {x}")
+            for x in range(1, 6)
+        ]
+
+    def test_bulk_create_history_with_custom_model_attributes(self):
+        bulk_create_with_history(
+            self.data,
+            PollWithHistoricalSessionAttr,
+            custom_historical_attrs={"session": "jam"},
+        )
+
+        self.assertEqual(PollWithHistoricalSessionAttr.objects.count(), 5)
+        self.assertEqual(
+            PollWithHistoricalSessionAttr.history.filter(session="jam").count(),
+            5,
+        )
+
+    def test_bulk_update_history_with_custom_model_attributes(self):
+        bulk_create_with_history(
+            self.data,
+            PollWithHistoricalSessionAttr,
+            custom_historical_attrs={"session": None},
+        )
+        bulk_update_with_history(
+            self.data,
+            PollWithHistoricalSessionAttr,
+            fields=["question"],
+            custom_historical_attrs={"session": "training"},
+        )
+
+        self.assertEqual(PollWithHistoricalSessionAttr.objects.count(), 5)
+        self.assertEqual(
+            PollWithHistoricalSessionAttr.history.filter(session="training").count(),
+            5,
+        )
+
+    def test_bulk_manager_with_custom_model_attributes(self):
+        history_manager = get_history_manager_for_model(PollWithHistoricalSessionAttr)
+        history_manager.bulk_history_create(
+            self.data, custom_historical_attrs={"session": "co-op"}
+        )
+
+        self.assertEqual(PollWithHistoricalSessionAttr.objects.count(), 0)
+        self.assertEqual(
+            PollWithHistoricalSessionAttr.history.filter(session="co-op").count(),
+            5,
+        )
+
+
 class UpdateChangeReasonTestCase(TestCase):
     def test_update_change_reason_with_excluded_fields(self):
         poll = PollWithExcludeFields(

diff --git a/simple_history/utils.py b/simple_history/utils.py
@@ -65,6 +65,7 @@ def bulk_create_with_history(
     default_user=None,
     default_change_reason=None,
     default_date=None,
+    custom_historical_attrs=None,
 ):
     """
     Bulk create the objects specified by objs while also bulk creating
@@ -81,6 +82,8 @@ def bulk_create_with_history(
         in each historical record
     :param default_date: Optional date to specify as the history_date in each historical
         record
+    :param custom_historical_attrs: Optional dict of field `name`:`value` to specify
+        values for custom fields
     :return: List of objs with IDs
     """
     # Exclude ManyToManyFields because they end up as invalid kwargs to
@@ -106,6 +109,7 @@ def bulk_create_with_history(
                 default_user=default_user,
                 default_change_reason=default_change_reason,
                 default_date=default_date,
+                custom_historical_attrs=custom_historical_attrs,
             )
     if second_transaction_required:
         with transaction.atomic(savepoint=False):
@@ -143,6 +147,7 @@ def bulk_create_with_history(
                 default_user=default_user,
                 default_change_reason=default_change_reason,
                 default_date=default_date,
+                custom_historical_attrs=custom_historical_attrs,
             )
         objs_with_id = obj_list
     return objs_with_id
@@ -157,6 +162,7 @@ def bulk_update_with_history(
     default_change_reason=None,
     default_date=None,
     manager=None,
+    custom_historical_attrs=None,
 ):
     """
     Bulk update the objects specified by objs while also bulk creating
@@ -173,6 +179,8 @@ def bulk_update_with_history(
         record
     :param manager: Optional model manager to use for the model instead of the default
         manager
+    :param custom_historical_attrs: Optional dict of field `name`:`value` to specify
+        values for custom fields
     :return: The number of model rows updated, not including any history objects
     """
     history_manager = get_history_manager_for_model(model)
@@ -189,6 +197,7 @@ def bulk_update_with_history(
             default_user=default_user,
             default_change_reason=default_change_reason,
             default_date=default_date,
+            custom_historical_attrs=custom_historical_attrs,
         )
     return rows_updated