From 51475b73d5e52c765909dcb7b7c0acc4c737cb15 Mon Sep 17 00:00:00 2001
From: Thomas Sileo <t@a4.io>
Date: Sat, 17 Aug 2019 10:27:20 +0200
Subject: [PATCH] Cache attachment one by one

---
 blueprints/tasks.py | 35 +++++++++++++++++++++++++----------
 core/tasks.py       |  6 ++++++
 2 files changed, 31 insertions(+), 10 deletions(-)

diff --git a/blueprints/tasks.py b/blueprints/tasks.py
index f37bdf4..88056fb 100644
--- a/blueprints/tasks.py
+++ b/blueprints/tasks.py
@@ -263,18 +263,15 @@ def task_cache_attachments() -> _Response:
                 # TODO: filter only videogt
                 link = select_video_to_cache(obj.url)
                 if link:
-                    config.MEDIA_CACHE.cache_attachment({"url": link["href"]}, iri)
+                    Tasks.cache_attachment({"url": link["href"]}, iri)
             elif isinstance(obj.url, str):
-                config.MEDIA_CACHE.cache_attachment({"url": obj.url}, iri)
+                Tasks.cache_attachment({"url": obj.url}, iri)
             else:
                 app.logger.warning(f"failed to parse video link {obj!r} for {iri}")
 
         # Iter the attachments
         for attachment in obj._data.get("attachment", []):
-            try:
-                config.MEDIA_CACHE.cache_attachment(attachment, iri)
-            except ValueError:
-                app.logger.exception(f"failed to cache {attachment}")
+            Tasks.cache_attachment(attachment, iri)
 
         app.logger.info(f"attachments cached for {iri}")
 
@@ -287,6 +284,25 @@ def task_cache_attachments() -> _Response:
     return ""
 
 
+@blueprint.route("/task/cache_attachment", methods=["POST"])
+def task_cache_attachment() -> _Response:
+    task = p.parse(flask.request)
+    app.logger.info(f"task={task!r}")
+    iri = task.payload["iri"]
+    attachment = task.payload["attachment"]
+    try:
+        app.logger.info(f"caching attachment {attachment!r} for {iri}")
+
+        config.MEDIA_CACHE.cache_attachment(attachment, iri)
+
+        app.logger.info(f"attachment {attachment!r} cached for {iri}")
+    except Exception as err:
+        app.logger.exception(f"failed to cache attachment {attachment!r} for {iri}")
+        raise TaskError() from err
+
+    return ""
+
+
 @blueprint.route("/task/cache_actor", methods=["POST"])
 def task_cache_actor() -> _Response:
     task = p.parse(flask.request)
@@ -320,10 +336,9 @@ def task_cache_actor() -> _Response:
         if not activity.has_type([ap.ActivityType.CREATE, ap.ActivityType.ANNOUNCE]):
             return ""
 
-        if activity.has_type(ap.ActivityType.CREATE) and (
-            activity.get_object()._data.get("attachment", [])
-            or activity.get_object().has_type(ap.ActivityType.VIDEO)
-        ):
+        if activity.get_object()._data.get(
+            "attachment", []
+        ) or activity.get_object().has_type(ap.ActivityType.VIDEO):
             Tasks.cache_attachments(iri)
 
     except (ActivityGoneError, ActivityNotFoundError):
diff --git a/core/tasks.py b/core/tasks.py
index bb20554..5c0e304 100644
--- a/core/tasks.py
+++ b/core/tasks.py
@@ -1,6 +1,8 @@
 import os
 from datetime import datetime
 from datetime import timezone
+from typing import Any
+from typing import Dict
 
 from poussetaches import PousseTaches
 
@@ -56,6 +58,10 @@ class Tasks:
     def cache_attachments(iri: str) -> None:
         p.push(iri, "/task/cache_attachments")
 
+    @staticmethod
+    def cache_attachment(attachment: Dict[str, Any], iri: str) -> None:
+        p.push({"iri": iri, "attachment": attachment}, "/task/cache_attachment")
+
     @staticmethod
     def finish_post_to_inbox(iri: str) -> None:
         p.push(iri, "/task/finish_post_to_inbox")