#570 added possibility to auto cleanup scripts after scheduled execution

src/model/script_config.py

@@ -1,4 +1,4 @@
-import json
+import logging
 import logging
 import os
 import re
@@ -56,6 +56,7 @@ def __init__(self,
         self._username = username
         self._audit_name = audit_name
         self.schedulable = False
+        self.scheduling_auto_cleanup = True
 
         self.parameters = ObservableList()
         self.parameter_values = ObservableDict()
@@ -179,8 +180,11 @@ def _reload_config(self):
 
         self.output_files = config.get('output_files', [])
 
-        if config.get('scheduling'):
-            self.schedulable = read_bool_from_config('enabled', config.get('scheduling'), default=False)
+        scheduling_config = config.get('scheduling')
+        if scheduling_config:
+            self.schedulable = read_bool_from_config('enabled', scheduling_config, default=False)
+            self.scheduling_auto_cleanup = read_bool_from_config(
+                'auto_cleanup', scheduling_config, default=not bool(self.output_files))
 
         if not self.script_command:
             raise Exception('No script_path is specified for ' + self.name)

src/scheduling/schedule_service.py

@@ -128,6 +128,12 @@ def _execute_job(self, job: SchedulingJob):
 
             execution_id = self._execution_service.start_script(config, parameter_values, user)
             LOGGER.info('Started script #' + str(execution_id) + ' for ' + job.get_log_name())
+
+            if config.scheduling_auto_cleanup:
+                def cleanup():
+                    self._execution_service.cleanup_execution(execution_id, user)
+
+                self._execution_service.add_finish_listener(cleanup, execution_id)
         except:
             LOGGER.exception('Failed to execute ' + job.get_log_name())
 

src/tests/scheduling/schedule_service_test.py

@@ -64,12 +64,13 @@ def setUp(self) -> None:
         self.create_config('unschedulable-script', scheduling_enabled=False)
 
         self.execution_service = MagicMock()
+        self.execution_service.start_script.side_effect = lambda config, values, user: time.time_ns()
 
         self.schedule_service = ScheduleService(self.config_service, self.execution_service, test_utils.temp_folder)
 
         date_utils._mocked_now = mocked_now
 
-    def create_config(self, name, scheduling_enabled=True, parameters=None):
+    def create_config(self, name, scheduling_enabled=True, parameters=None, auto_cleanup=False):
         if parameters is None:
             parameters = [
                 {'name': 'p1'},
@@ -81,7 +82,7 @@ def create_config(self, name, scheduling_enabled=True, parameters=None):
                 'name': name,
                 'script_path': 'echo 1',
                 'parameters': parameters,
-                'scheduling': {'enabled': scheduling_enabled}
+                'scheduling': {'enabled': scheduling_enabled, 'auto_cleanup': auto_cleanup}
             },
             name)
 
@@ -295,6 +296,7 @@ def test_execute_simple_job(self):
 
         self.execution_service.start_script.assert_called_once_with(
             ANY, job.parameter_values, job.user)
+        self.execution_service.add_finish_listener.assert_not_called()
         self.assert_schedule_calls([])
 
     def test_execute_repeatable_job(self):
@@ -308,6 +310,7 @@ def test_execute_repeatable_job(self):
 
         self.execution_service.start_script.assert_called_once_with(
             ANY, job.parameter_values, job.user)
+        self.execution_service.add_finish_listener.assert_not_called()
         self.assert_schedule_calls([(job, mocked_now_epoch + 86399)])
 
     def test_execute_when_fails(self):
@@ -337,6 +340,7 @@ def test_execute_when_not_schedulable(self):
 
     def test_execute_when_has_secure_parameters(self):
         job = create_job(id=1,
+                         script_name='secure-config',
                          repeatable=True,
                          start_datetime=mocked_now - timedelta(seconds=1),
                          repeat_unit='days',
@@ -349,6 +353,34 @@ def test_execute_when_has_secure_parameters(self):
         self.execution_service.start_script.assert_not_called()
         self.assert_schedule_calls([(job, mocked_now_epoch + 86399)])
 
+    def test_cleanup_execution(self):
+        self.create_config('script_with_cleanup', auto_cleanup=True)
+        job = create_job(id=1,
+                         script_name='script_with_cleanup',
+                         repeatable=False,
+                         start_datetime=mocked_now - timedelta(seconds=1))
+
+        finish_callback = None
+
+        def add_finish_listener(callback_param, execution_id):
+            self.assertIsNotNone(execution_id)
+            nonlocal finish_callback
+            finish_callback = callback_param
+
+        self.execution_service.add_finish_listener.side_effect = add_finish_listener
+
+        self.schedule_service._execute_job(job)
+
+        self.execution_service.start_script.assert_called_once_with(
+            ANY, job.parameter_values, job.user)
+        self.execution_service.cleanup_execution.assert_not_called()
+        self.assertIsNotNone(finish_callback)
+
+        # noinspection PyCallingNonCallable
+        finish_callback()
+
+        self.execution_service.cleanup_execution.assert_called_once_with(ANY, job.user)
+
 
 def create_job(id=None,
                user_id='UserX',

src/tests/script_config_test.py

@@ -913,6 +913,29 @@ def test_create_with_schedulable_true_and_included_secure_parameter(self):
 
         self.assertFalse(config_model.schedulable)
 
+    @parameterized.expand([
+        (True, True, [], True),
+        (False, True, [], True),
+        (True, False, [], False),
+        (False, False, [], False),
+        (True, True, ['test'], True),
+        (True, False, ['test'], False),
+        (True, None, [], True),
+        (False, None, [], True),
+        (True, None, ['test'], False),
+    ])
+    def test_auto_cleanup(self, scheduling_enabled, auto_cleanup, output_files, expected_result):
+        config = {
+            'scheduling': {'enabled': scheduling_enabled},
+            'output_files': output_files
+        }
+        if auto_cleanup is not None:
+            config['scheduling']['auto_cleanup'] = auto_cleanup
+
+        config_model = _create_config_model('some-name', config=config)
+
+        self.assertEqual(config_model.scheduling_auto_cleanup, expected_result)
+
     def tearDown(self) -> None:
         test_utils.cleanup()