From db0e362b7d7e143a898f398b28b91328ac1f955e Mon Sep 17 00:00:00 2001
From: simon <simobilleter@gmail.com>
Date: Wed, 22 Mar 2023 17:01:34 +0700
Subject: [PATCH] make update_subscribed stoppable

---
 .../home/src/download/subscriptions.py        | 24 ++++++++++++-------
 tubearchivist/home/src/ta/task_manager.py     |  7 ++++++
 tubearchivist/home/tasks.py                   |  5 ++++
 3 files changed, 28 insertions(+), 8 deletions(-)

diff --git a/tubearchivist/home/src/download/subscriptions.py b/tubearchivist/home/src/download/subscriptions.py
index 4bb6f689..dd8a6bf8 100644
--- a/tubearchivist/home/src/download/subscriptions.py
+++ b/tubearchivist/home/src/download/subscriptions.py
@@ -126,10 +126,14 @@ class ChannelSubscription:
             if not self.task:
                 continue
 
-            self.task.send_progress(
-                message_lines=[f"Scanning Channel {idx + 1}/{total}"],
-                progress=(idx + 1) / total,
-            )
+            if self.task:
+                self.task.send_progress(
+                    message_lines=[f"Scanning Channel {idx + 1}/{total}"],
+                    progress=(idx + 1) / total,
+                )
+                if self.task.is_stopped():
+                    self.task.send_progress(["Received Stop signal."])
+                    break
 
         return missing_videos
 
@@ -261,10 +265,14 @@ class PlaylistSubscription:
             if not self.task:
                 continue
 
-            self.task.send_progress(
-                message_lines=[f"Scanning Playlists {idx + 1}/{total}"],
-                progress=(idx + 1) / total,
-            )
+            if self.task:
+                self.task.send_progress(
+                    message_lines=[f"Scanning Playlists {idx + 1}/{total}"],
+                    progress=(idx + 1) / total,
+                )
+                if self.task.is_stopped():
+                    self.task.send_progress(["Received Stop signal."])
+                    break
 
         return missing_videos
 
diff --git a/tubearchivist/home/src/ta/task_manager.py b/tubearchivist/home/src/ta/task_manager.py
index 517318a6..62fff923 100644
--- a/tubearchivist/home/src/ta/task_manager.py
+++ b/tubearchivist/home/src/ta/task_manager.py
@@ -40,6 +40,13 @@ class TaskManager:
 
         return bool([i for i in tasks if i.get("status") == "PENDING"])
 
+    def is_stopped(self, task_id):
+        """check if task_id has received STOP command"""
+        task = self.get_task(task_id)
+        print(task)
+
+        return task.get("command") == "STOP"
+
     def get_pending(self, task_name):
         """get all pending tasks of task_name"""
         tasks = self.get_tasks_by_name(task_name)
diff --git a/tubearchivist/home/tasks.py b/tubearchivist/home/tasks.py
index f044b1f6..5c307205 100644
--- a/tubearchivist/home/tasks.py
+++ b/tubearchivist/home/tasks.py
@@ -53,6 +53,7 @@ class BaseTask(Task):
             "title": "Rescan your Subscriptions",
             "group": "download:scan",
             "api-start": True,
+            "api-stop": True,
         },
         "download_pending": {
             "title": "Downloading",
@@ -146,6 +147,10 @@ class BaseTask(Task):
         key = f"message:{message.get('group')}:{task_id.split('-')[0]}"
         return message, key
 
+    def is_stopped(self):
+        """check if task is stopped"""
+        return TaskManager().is_stopped(self.request.id)
+
 
 @shared_task(name="update_subscribed", bind=True, base=BaseTask)
 def update_subscribed(self):