From b3ec7cb9b873b41819bd8e00c7718d8b6194b5a2 Mon Sep 17 00:00:00 2001
From: Lance Edgar <lance@wuttaproject.org>
Date: Thu, 2 Jan 2025 19:48:26 -0600
Subject: [PATCH] fix: add batch handler logic to remove row

also execute() can return whatever it wants, e.g. when creating some
new record(s) based on batch data
---
 src/wuttjamaican/batch.py | 58 ++++++++++++++++++++++++++++++++++++++-
 tests/test_batch.py       | 15 ++++++++++
 2 files changed, 72 insertions(+), 1 deletion(-)

diff --git a/src/wuttjamaican/batch.py b/src/wuttjamaican/batch.py
index 98325ff..c9ca664 100644
--- a/src/wuttjamaican/batch.py
+++ b/src/wuttjamaican/batch.py
@@ -296,6 +296,54 @@ class BatchHandler(GenericHandler):
         that depending on the workflow.
         """
 
+    def do_remove_row(self, row):
+        """
+        Remove a row from its batch.  This will:
+
+        * call :meth:`remove_row()`
+        * decrement the batch
+          :attr:`~wuttjamaican.db.model.batch.BatchMixin.row_count`
+        * call :meth:`refresh_batch_status()`
+
+        So, callers should use ``do_remove_row()``, but subclass
+        should (usually) override :meth:`remove_row()` etc.
+        """
+        batch = row.batch
+
+        self.remove_row(row)
+
+        if batch.row_count is not None:
+            batch.row_count -= 1
+
+        self.refresh_batch_status(batch)
+
+    def remove_row(self, row):
+        """
+        Remove a row from its batch.
+
+        Callers should use :meth:`do_remove_row()` instead, which
+        calls this method automatically.
+
+        Subclass can override this method; the default logic just
+        deletes the row.
+        """
+        session = self.app.get_session(row)
+        session.delete(row)
+
+    def refresh_batch_status(self, batch):
+        """
+        Update the batch status as needed.
+
+        This method is called when some row data has changed for the
+        batch, e.g. from :meth:`do_remove_row()`.
+
+        It does nothing by default; subclass may override to set these
+        attributes on the batch:
+
+        * :attr:`~wuttjamaican.db.model.batch.BatchMixin.status_code`
+        * :attr:`~wuttjamaican.db.model.batch.BatchMixin.status_text`
+        """
+
     def why_not_execute(self, batch, user=None, **kwargs):
         """
         Returns text indicating the reason (if any) that a given batch
@@ -395,6 +443,9 @@ class BatchHandler(GenericHandler):
         :param \**kwargs: Additional kwargs as needed.  These are
            passed as-is to :meth:`why_not_execute()` and
            :meth:`execute()`.
+
+        :returns: Whatever was returned from :meth:`execute()` - often
+           ``None``.
         """
         if batch.executed:
             raise ValueError(f"batch has already been executed: {batch}")
@@ -403,9 +454,10 @@ class BatchHandler(GenericHandler):
         if reason:
             raise RuntimeError(f"batch execution not allowed: {reason}")
 
-        self.execute(batch, user=user, progress=progress, **kwargs)
+        result = self.execute(batch, user=user, progress=progress, **kwargs)
         batch.executed = datetime.datetime.now()
         batch.executed_by = user
+        return result
 
     def execute(self, batch, user=None, progress=None, **kwargs):
         """
@@ -428,6 +480,10 @@ class BatchHandler(GenericHandler):
         :param \**kwargs: Additional kwargs which may affect the batch
            execution behavior.  There are none by default, but some
            handlers may declare/use them.
+
+        :returns: ``None`` by default, but subclass can return
+           whatever it likes, in which case that will be also returned
+           to the caller from :meth:`do_execute()`.
         """
 
     def do_delete(self, batch, user, dry_run=False, progress=None, **kwargs):
diff --git a/tests/test_batch.py b/tests/test_batch.py
index 6075fd0..7d22244 100644
--- a/tests/test_batch.py
+++ b/tests/test_batch.py
@@ -115,6 +115,21 @@ else:
             handler.add_row(batch, row)
             self.assertEqual(batch.row_count, 1)
 
+        def test_remove_row(self):
+            model = self.app.model
+            handler = self.make_handler()
+            user = model.User(username='barney')
+            self.session.add(user)
+            batch = handler.make_batch(self.session, created_by=user)
+            self.session.add(batch)
+            row = handler.make_row()
+            handler.add_row(batch, row)
+            self.session.flush()
+            self.assertEqual(batch.row_count, 1)
+            handler.do_remove_row(row)
+            self.session.flush()
+            self.assertEqual(batch.row_count, 0)
+
         def test_do_execute(self):
             model = self.app.model
             user = model.User(username='barney')