From a400443dc0bdb8434de0ce6e7be07d5d207d50b3 Mon Sep 17 00:00:00 2001
From: msramalho <19508417+msramalho@users.noreply.github.com>
Date: Fri, 10 Mar 2023 11:00:02 +0000
Subject: [PATCH] task retries

---
 src/main.py   | 7 +++++--
 src/worker.py | 3 +--
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/src/main.py b/src/main.py
index a4443f5..3b5f933 100644
--- a/src/main.py
+++ b/src/main.py
@@ -1,5 +1,5 @@
 from celery.result import AsyncResult
-from fastapi import Body, FastAPI, Depends, Request
+from fastapi import Body, FastAPI, Depends, Request, HTTPException
 from fastapi.encoders import jsonable_encoder
 from fastapi.responses import JSONResponse, FileResponse
 from fastapi.staticfiles import StaticFiles
@@ -68,7 +68,10 @@ def search(skip: int = 0, limit: int = 100, db: Session = Depends(get_db), email
 
 @app.post("/tasks", status_code=201)
 def run_task(payload = Body(...), email = Depends(get_bearer_auth)):
-    logger.info(f"new task for user {email}: {payload.get('url')}")
+    url = payload.get('url')
+    logger.info(f"new task for user {email}: {url}")
+    if type(url)!=str or len(url)<=5:
+        raise HTTPException(status_code=422, detail=f"Invalid URL received: {url}")
     task = create_archive_task.delay(url=payload.get('url'), email=email)
     return JSONResponse({"id": task.id})
 
diff --git a/src/worker.py b/src/worker.py
index 1627d9a..08fbae4 100644
--- a/src/worker.py
+++ b/src/worker.py
@@ -31,9 +31,8 @@ if (config_bcat_file := os.environ.get("ORCHESTRATION_CONFIG_BELLINGCAT")):
 
 orchestrators = {"bellingcat": None, "default": None}
 
-@celery.task(name="create_archive_task", bind=True)
+@celery.task(name="create_archive_task", bind=True, autoretry_for=(Exception,), retry_backoff=True, retry_kwargs={'max_retries': 5})
 def create_archive_task(self, url: str, email:str=""):
-    assert type(url)==str and len(url)>5, f"Invalid URL received: {url}"
     orchestrator = choose_orchestrator(email)
     result = orchestrator.feed_item(Metadata().set_url(url)).to_json()
     with get_db() as session: