From 2fa462387641f32293c5bdb58b94fe899cc8e977 Mon Sep 17 00:00:00 2001
From: Jon Herron <jon.herron@yahoo.com>
Date: Wed, 22 Feb 2023 10:39:35 -0500
Subject: [PATCH] Add cumulative task data size to script

---
 .../scripts/get_data_sizes.py                 | 37 +++++++++++++++++++
 .../services/process_instance_processor.py    | 22 +++++++----
 2 files changed, 52 insertions(+), 7 deletions(-)
 create mode 100644 spiffworkflow-backend/src/spiffworkflow_backend/scripts/get_data_sizes.py

diff --git a/spiffworkflow-backend/src/spiffworkflow_backend/scripts/get_data_sizes.py b/spiffworkflow-backend/src/spiffworkflow_backend/scripts/get_data_sizes.py
new file mode 100644
index 00000000..1721bcdd
--- /dev/null
+++ b/spiffworkflow-backend/src/spiffworkflow_backend/scripts/get_data_sizes.py
@@ -0,0 +1,37 @@
+"""Get_data_sizes."""
+from typing import Any
+
+from spiffworkflow_backend.models.script_attributes_context import (
+    ScriptAttributesContext,
+)
+from spiffworkflow_backend.scripts.script import Script
+#from spiffworkflow_backend.servces.process_instance_processor import ProcessInstanceProcessor
+from spiffworkflow_backend.services.process_instance_processor import (
+    ProcessInstanceProcessor,
+)
+
+
+class GetDataSizes(Script):
+    """GetDataSizes."""
+
+    @staticmethod
+    def requires_privileged_permissions() -> bool:
+        """We have deemed this function safe to run without elevated permissions."""
+        return False
+
+    def get_description(self) -> str:
+        """Get_description."""
+        return """Returns a dictionary of information about the size of task data and the python environment for the currently running process."""
+
+    def run(
+        self,
+        script_attributes_context: ScriptAttributesContext,
+        *_args: Any,
+        **kwargs: Any
+    ) -> Any:
+        """Run."""
+        task = script_attributes_context.task
+        return {
+            "cumulative_task_data_size": ProcessInstanceProcessor.get_task_data_size(task.workflow),
+            "python_env_size": 0,
+        }
diff --git a/spiffworkflow-backend/src/spiffworkflow_backend/services/process_instance_processor.py b/spiffworkflow-backend/src/spiffworkflow_backend/services/process_instance_processor.py
index 5aabe5ac..62bf61ad 100644
--- a/spiffworkflow-backend/src/spiffworkflow_backend/services/process_instance_processor.py
+++ b/spiffworkflow-backend/src/spiffworkflow_backend/services/process_instance_processor.py
@@ -1602,16 +1602,24 @@ class ProcessInstanceProcessor:
         except WorkflowTaskException as we:
             raise ApiError.from_workflow_exception("task_error", str(we), we) from we
 
-    def check_task_data_size(self) -> None:
-        """CheckTaskDataSize."""
-        tasks_to_check = self.bpmn_process_instance.get_tasks(TaskState.FINISHED_MASK)
-        task_data = [task.data for task in tasks_to_check]
-        task_data_to_check = list(filter(len, task_data))
+    @classmethod
+    def _get_data_size(cls, data: Dict[Any, Any]) -> int:
+        data_to_check = list(filter(len, data))
 
         try:
-            task_data_len = len(json.dumps(task_data_to_check))
+            return len(json.dumps(data_to_check))
         except Exception:
-            task_data_len = 0
+            return 0
+
+    @classmethod
+    def get_task_data_size(cls, bpmn_process_instance: BpmnWorkflow) -> int:
+        tasks_to_check = bpmn_process_instance.get_tasks(TaskState.FINISHED_MASK)
+        task_data = [task.data for task in tasks_to_check]
+        return cls._get_data_size(task_data)
+
+    def check_task_data_size(self) -> None:
+        """CheckTaskDataSize."""
+        task_data_len = self.get_task_data_size(self.bpmn_process_instance)
 
         # Not sure what the number here should be but this now matches the mysql
         # max_allowed_packet variable on dev - 1073741824