From 71dd63d3f06d8147748fc5ead38bc60ab667fad7 Mon Sep 17 00:00:00 2001
From: simon <simobilleter@gmail.com>
Date: Mon, 19 Dec 2022 11:24:37 +0700
Subject: [PATCH] [API] implement refresh GET and POST endpoints

---
 tubearchivist/api/README.md             | 32 +++++++++++
 tubearchivist/api/views.py              | 17 +++++-
 tubearchivist/home/src/index/reindex.py | 74 +++++++++++++++++++++++++
 tubearchivist/home/src/ta/ta_redis.py   |  7 ++-
 4 files changed, 126 insertions(+), 4 deletions(-)

diff --git a/tubearchivist/api/README.md b/tubearchivist/api/README.md
index 1a0a7237..3b69e763 100644
--- a/tubearchivist/api/README.md
+++ b/tubearchivist/api/README.md
@@ -309,6 +309,38 @@ List of valid task names:
 
 ## Refresh View
 GET /api/refresh/  
+parameters:
+- **type**: one of *video*, *channel*, *playlist*, optional
+- **id**: item id, optional
+
+without specifying type: return total for all queued items:
+```json
+{
+    "total_queued": 2,
+    "type": "all",
+    "state": "running"
+}
+```
+
+specify type: return total items queue of this type:
+```json
+{
+    "total_queued": 2,
+    "type": "video",
+    "state": "running"
+}
+```
+
+specify type *and* id to get state of item in queue:
+```json
+{
+    "total_queued": 2,
+    "type": "video",
+    "state": "in_queue",
+    "id": "video-id"
+}
+```
+
 POST /api/refresh/  
 Parameter:
 - extract_videos: to refresh all videos for channels/playlists, default False
diff --git a/tubearchivist/api/views.py b/tubearchivist/api/views.py
index aebbf297..df8554cc 100644
--- a/tubearchivist/api/views.py
+++ b/tubearchivist/api/views.py
@@ -8,6 +8,7 @@ from home.src.es.connect import ElasticWrap
 from home.src.es.snapshot import ElasticSnapshot
 from home.src.frontend.searching import SearchForm
 from home.src.index.generic import Pagination
+from home.src.index.reindex import ReindexProgress
 from home.src.index.video import SponsorBlock
 from home.src.ta.config import AppConfig
 from home.src.ta.helper import UrlListParser
@@ -597,8 +598,20 @@ class RefreshView(ApiBaseView):
 
     def get(self, request):
         """handle get request"""
-        # pylint: disable=unused-argument
-        return Response({"status": False})
+        request_type = request.GET.get("type")
+        request_id = request.GET.get("id")
+
+        if request_id and not request_type:
+            return Response({"status": "Bad Request"}, status=400)
+
+        try:
+            progress = ReindexProgress(
+                request_type=request_type, request_id=request_id
+            ).get_progress()
+        except ValueError:
+            return Response({"status": "Bad Request"}, status=400)
+
+        return Response(progress)
 
     def post(self, request):
         """handle post request"""
diff --git a/tubearchivist/home/src/index/reindex.py b/tubearchivist/home/src/index/reindex.py
index e02d11a3..780d7dea 100644
--- a/tubearchivist/home/src/index/reindex.py
+++ b/tubearchivist/home/src/index/reindex.py
@@ -342,6 +342,80 @@ class Reindex(ReindexBase):
         return valid
 
 
+class ReindexProgress(ReindexBase):
+    """
+    get progress of reindex task
+    request_type: key of self.REINDEX_CONFIG
+    request_id: id of request_type
+    return = {
+        "state": "running" | "queued" | False
+        "total_queued": int
+        "in_queue_name": "queue_name"
+    }
+    """
+
+    def __init__(self, request_type=False, request_id=False):
+        super().__init__()
+        self.request_type = request_type
+        self.request_id = request_id
+
+    def get_progress(self):
+        """get progress from task"""
+        queue_name, request_type = self._get_queue_name()
+        total = self._get_total_in_queue(queue_name)
+
+        progress = {
+            "total_queued": total,
+            "type": request_type,
+        }
+        state = self._get_state(total, queue_name)
+        progress.update(state)
+
+        return progress
+
+    def _get_queue_name(self):
+        """return queue_name, queue_type, raise exception on error"""
+        if not self.request_type:
+            return "all", "all"
+
+        reindex_config = self.REINDEX_CONFIG.get(self.request_type)
+        if not reindex_config:
+            print(f"reindex_config not found: {self.request_type}")
+            raise ValueError
+
+        return reindex_config["queue_name"], self.request_type
+
+    def _get_total_in_queue(self, queue_name):
+        """get all items in queue"""
+        total = 0
+        if queue_name == "all":
+            queues = [i["queue_name"] for i in self.REINDEX_CONFIG.values()]
+            for queue in queues:
+                total += len(RedisQueue(queue).get_all())
+        else:
+            total += len(RedisQueue(queue_name).get_all())
+
+        return total
+
+    def _get_state(self, total, queue_name):
+        """get state based on request_id"""
+        state_dict = {}
+        if self.request_id:
+            state = RedisQueue(queue_name).in_queue(self.request_id)
+            state_dict.update({"id": self.request_id, "state": state})
+
+            return state_dict
+
+        if total:
+            state = "running"
+        else:
+            state = "empty"
+
+        state_dict.update({"state": state})
+
+        return state_dict
+
+
 class ChannelUrlFixer:
     """fix not matching channel names in reindex"""
 
diff --git a/tubearchivist/home/src/ta/ta_redis.py b/tubearchivist/home/src/ta/ta_redis.py
index 40529cd3..14e90a11 100644
--- a/tubearchivist/home/src/ta/ta_redis.py
+++ b/tubearchivist/home/src/ta/ta_redis.py
@@ -114,10 +114,13 @@ class RedisQueue(RedisBase):
         all_elements = [i.decode() for i in result]
         return all_elements
 
-    def contains(self, element):
+    def in_queue(self, element):
         """check if element is in list"""
         result = self.conn.execute_command("LPOS", self.key, element)
-        return result is not None
+        if result is not None:
+            return "in_queue"
+
+        return False
 
     def add_list(self, to_add):
         """add list to queue"""