2 files changed, 32 insertions, 3 deletions

diff --git a/lib/private/TaskProcessing/Db/Task.php b/lib/private/TaskProcessing/Db/Task.php
index 4d919deaf94..05c0ae9ac74 100644
--- a/lib/private/TaskProcessing/Db/Task.php
+++ b/lib/private/TaskProcessing/Db/Task.php
@@ -45,6 +45,8 @@ use OCP\TaskProcessing\Task as OCPTask;
  * @method int getStartedAt()
  * @method setEndedAt(int $endedAt)
  * @method int getEndedAt()
+ * @method setAllowCleanup(int $allowCleanup)
+ * @method int getAllowCleanup()
  */
 class Task extends Entity {
 	protected $lastUpdated;
@@ -63,16 +65,17 @@ class Task extends Entity {
 	protected $scheduledAt;
 	protected $startedAt;
 	protected $endedAt;
+	protected $allowCleanup;
 
 	/**
 	 * @var string[]
 	 */
-	public static array $columns = ['id', 'last_updated', 'type', 'input', 'output', 'status', 'user_id', 'app_id', 'custom_id', 'completion_expected_at', 'error_message', 'progress', 'webhook_uri', 'webhook_method', 'scheduled_at', 'started_at', 'ended_at'];
+	public static array $columns = ['id', 'last_updated', 'type', 'input', 'output', 'status', 'user_id', 'app_id', 'custom_id', 'completion_expected_at', 'error_message', 'progress', 'webhook_uri', 'webhook_method', 'scheduled_at', 'started_at', 'ended_at', 'allow_cleanup'];
 
 	/**
 	 * @var string[]
 	 */
-	public static array $fields = ['id', 'lastUpdated', 'type', 'input', 'output', 'status', 'userId', 'appId', 'customId', 'completionExpectedAt', 'errorMessage', 'progress', 'webhookUri', 'webhookMethod', 'scheduledAt', 'startedAt', 'endedAt'];
+	public static array $fields = ['id', 'lastUpdated', 'type', 'input', 'output', 'status', 'userId', 'appId', 'customId', 'completionExpectedAt', 'errorMessage', 'progress', 'webhookUri', 'webhookMethod', 'scheduledAt', 'startedAt', 'endedAt', 'allowCleanup'];
 
 
 	public function __construct() {
@@ -94,6 +97,7 @@ class Task extends Entity {
 		$this->addType('scheduledAt', 'integer');
 		$this->addType('startedAt', 'integer');
 		$this->addType('endedAt', 'integer');
+		$this->addType('allowCleanup', 'integer');
 	}
 
 	public function toRow(): array {
@@ -122,6 +126,7 @@ class Task extends Entity {
 			'scheduledAt' => $task->getScheduledAt(),
 			'startedAt' => $task->getStartedAt(),
 			'endedAt' => $task->getEndedAt(),
+			'allowCleanup' => $task->getAllowCleanup() ? 1 : 0,
 		]);
 		return $taskEntity;
 	}
@@ -144,6 +149,7 @@ class Task extends Entity {
 		$task->setScheduledAt($this->getScheduledAt());
 		$task->setStartedAt($this->getStartedAt());
 		$task->setEndedAt($this->getEndedAt());
+		$task->setAllowCleanup($this->getAllowCleanup() !== 0);
 		return $task;
 	}
 }
diff --git a/lib/private/TaskProcessing/Db/TaskMapper.php b/lib/private/TaskProcessing/Db/TaskMapper.php
index 91fd68820ae..fee96534633 100644
--- a/lib/private/TaskProcessing/Db/TaskMapper.php
+++ b/lib/private/TaskProcessing/Db/TaskMapper.php
@@ -183,16 +183,39 @@ class TaskMapper extends QBMapper {
 
 	/**
 	 * @param int $timeout
+	 * @param bool $force If true, ignore the allow_cleanup flag
 	 * @return int the number of deleted tasks
 	 * @throws Exception
 	 */
-	public function deleteOlderThan(int $timeout): int {
+	public function deleteOlderThan(int $timeout, bool $force = false): int {
 		$qb = $this->db->getQueryBuilder();
 		$qb->delete($this->tableName)
 			->where($qb->expr()->lt('last_updated', $qb->createPositionalParameter($this->timeFactory->getDateTime()->getTimestamp() - $timeout)));
+		if (!$force) {
+			$qb->andWhere($qb->expr()->eq('allow_cleanup', $qb->createPositionalParameter(1, IQueryBuilder::PARAM_INT)));
+		}
 		return $qb->executeStatement();
 	}
 
+	/**
+	 * @param int $timeout
+	 * @param bool $force If true, ignore the allow_cleanup flag
+	 * @return \Generator<Task>
+	 * @throws Exception
+	 */
+	public function getTasksToCleanup(int $timeout, bool $force = false): \Generator {
+		$qb = $this->db->getQueryBuilder();
+		$qb->select(Task::$columns)
+			->from($this->tableName)
+			->where($qb->expr()->lt('last_updated', $qb->createPositionalParameter($this->timeFactory->getDateTime()->getTimestamp() - $timeout)));
+		if (!$force) {
+			$qb->andWhere($qb->expr()->eq('allow_cleanup', $qb->createPositionalParameter(1, IQueryBuilder::PARAM_INT)));
+		}
+		foreach ($this->yieldEntities($qb) as $entity) {
+			yield $entity;
+		};
+	}
+
 	public function update(Entity $entity): Entity {
 		$entity->setLastUpdated($this->timeFactory->now()->getTimestamp());
 		return parent::update($entity);