feat(api): add S3-backed screenshot storage

bases/ecoindex/backend/routers/ecoindex.py

@@ -1,4 +1,3 @@
-from os import getcwd
 from typing import Annotated
 
 from ecoindex.backend.models.dependencies_parameters.dates import DateRangeParameters
@@ -22,6 +21,12 @@
 )
 from ecoindex.models import example_ecoindex_not_found, example_file_not_found
 from ecoindex.models.enums import Version
+from ecoindex.screenshot_storage import (
+    get_screenshot_local_path,
+    is_s3_screenshot_storage,
+    read_screenshot,
+    screenshot_exists,
+)
 from fastapi import APIRouter, Depends, HTTPException, Response, status
 from fastapi.params import Query
 from fastapi.responses import FileResponse
@@ -128,8 +133,21 @@ async def get_screenshot_endpoint(
     id: IdParameter,
     version: VersionParameter = Version.v1,
 ):
+    if not screenshot_exists(version=version.value, screenshot_id=str(id)):
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND,
+            detail=f"Screenshot {version.value}/{id}.webp does not exist.",
+        )
+
+    if is_s3_screenshot_storage():
+        return Response(
+            content=read_screenshot(version=version.value, screenshot_id=str(id)),
+            headers={"Content-Disposition": f'inline; filename="{id}.webp"'},
+            media_type="image/webp",
+        )
+
     return FileResponse(
-        path=f"{getcwd()}/screenshots/{version.value}/{id}.webp",
+        path=get_screenshot_local_path(version=version.value, screenshot_id=str(id)),
         filename=f"{id}.webp",
         content_disposition_type="inline",
         media_type="image/webp",

bases/ecoindex/worker/tasks.py

@@ -1,5 +1,4 @@
 from asyncio import run
-from os import getcwd
 from urllib.parse import urlparse
 from uuid import UUID
 
@@ -18,10 +17,14 @@
     EcoindexTimeout,
 )
 from ecoindex.models import ScreenShot, WindowSize
-from ecoindex.models.enums import TaskStatus
+from ecoindex.models.enums import TaskStatus, Version
 from ecoindex.models.tasks import QueueTaskError, QueueTaskResult
 from ecoindex.monitoring import capture_task_failure, init_sentry
 from ecoindex.scraper.scrap import EcoindexScraper
+from ecoindex.screenshot_storage import (
+    get_screenshot_local_folder,
+    persist_screenshot,
+)
 from playwright._impl._errors import Error as WebDriverException
 from rq import get_current_job
 
@@ -59,26 +62,34 @@ async def async_ecoindex_task(
     custom_headers: dict[str, str],
 ) -> QueueTaskResult:
     try:
+        settings = Settings()
         session_generator = get_session()
         session = await session_generator.__anext__()
+        screenshot = (
+            ScreenShot(
+                id=str(task_id),
+                folder=get_screenshot_local_folder(version=Version.v1.value),
+            )
+            if settings.ENABLE_SCREENSHOT
+            else None
+        )
 
         await check_quota(session=session, host=urlparse(url=url).netloc)
 
         ecoindex = await EcoindexScraper(
             url=url,
             window_size=WindowSize(height=height, width=width),
-            wait_after_scroll=Settings().WAIT_AFTER_SCROLL,
-            wait_before_scroll=Settings().WAIT_BEFORE_SCROLL,
-            screenshot=ScreenShot(
-                id=str(task_id), folder=f"{getcwd()}/screenshots/v1"
-            )
-            if Settings().ENABLE_SCREENSHOT
-            else None,
-            screenshot_gid=Settings().SCREENSHOTS_GID,
-            screenshot_uid=Settings().SCREENSHOTS_UID,
+            wait_after_scroll=settings.WAIT_AFTER_SCROLL,
+            wait_before_scroll=settings.WAIT_BEFORE_SCROLL,
+            screenshot=screenshot,
+            screenshot_gid=settings.SCREENSHOTS_GID,
+            screenshot_uid=settings.SCREENSHOTS_UID,
             custom_headers=custom_headers,
         ).get_page_analysis()
 
+        if screenshot:
+            persist_screenshot(screenshot=screenshot, version=Version.v1.value)
+
         db_result = await save_ecoindex_result_db(
             session=session,
             id=task_id,

components/ecoindex/config/settings.py

@@ -26,6 +26,15 @@ class Settings(BaseSettings):
     RQ_JOB_TIMEOUT: int = 600
     RQ_RESULT_TTL: int = 86400
     RQ_WORKERS: int = 3
+    SCREENSHOT_FILESYSTEM_PATH: str = "./screenshots"
+    SCREENSHOT_STORAGE_TYPE: str = "filesystem"
+    SCREENSHOT_S3_ACCESS_KEY_ID: str = ""
+    SCREENSHOT_S3_BUCKET: str = ""
+    SCREENSHOT_S3_ENDPOINT_URL: str = ""
+    SCREENSHOT_S3_FORCE_PATH_STYLE: bool = True
+    SCREENSHOT_S3_PREFIX: str = "screenshots"
+    SCREENSHOT_S3_REGION: str = "us-east-1"
+    SCREENSHOT_S3_SECRET_ACCESS_KEY: str = ""
     SCREENSHOTS_GID: int | None = None
     SCREENSHOTS_UID: int | None = None
     TZ: str = "Europe/Paris"

components/ecoindex/models/response_examples.py

@@ -18,7 +18,7 @@
         "application/json": {
             "example": {
                 "detail": (
-                    "File at path screenshots/v0/"
+                    "Screenshot v1/"
                     "550cdf8c-9c4c-4f8a-819d-cb69d0866fe1.webp does not exist."
                 )
             }

components/ecoindex/screenshot_storage/__init__.py

@@ -0,0 +1,19 @@
+from ecoindex.screenshot_storage.storage import (
+    get_screenshot_local_folder,
+    get_screenshot_local_path,
+    get_screenshot_object_key,
+    is_s3_screenshot_storage,
+    persist_screenshot,
+    read_screenshot,
+    screenshot_exists,
+)
+
+__all__ = [
+    "get_screenshot_local_folder",
+    "get_screenshot_local_path",
+    "get_screenshot_object_key",
+    "is_s3_screenshot_storage",
+    "persist_screenshot",
+    "read_screenshot",
+    "screenshot_exists",
+]

components/ecoindex/screenshot_storage/storage.py

@@ -0,0 +1,135 @@
+from os import getcwd
+from pathlib import Path
+
+import boto3
+from botocore.client import Config
+from botocore.exceptions import ClientError
+from ecoindex.config.settings import Settings
+from ecoindex.models import ScreenShot
+
+TEMPORARY_SCREENSHOT_ROOT = "/tmp/ecoindex-screenshots"
+
+
+def get_screenshot_storage_type() -> str:
+    return Settings().SCREENSHOT_STORAGE_TYPE.lower()
+
+
+def is_s3_screenshot_storage() -> bool:
+    return get_screenshot_storage_type() == "s3"
+
+
+def _get_root_path(path: str) -> Path:
+    root_path = Path(path).expanduser()
+    if root_path.is_absolute():
+        return root_path
+    return Path(getcwd()) / root_path
+
+
+def _get_s3_client():
+    settings = Settings()
+    if not settings.SCREENSHOT_S3_BUCKET:
+        raise RuntimeError("SCREENSHOT_S3_BUCKET must be configured when using S3 storage.")
+
+    if not settings.SCREENSHOT_S3_ENDPOINT_URL:
+        raise RuntimeError(
+            "SCREENSHOT_S3_ENDPOINT_URL must be configured when using S3 storage."
+        )
+
+    if not settings.SCREENSHOT_S3_ACCESS_KEY_ID:
+        raise RuntimeError(
+            "SCREENSHOT_S3_ACCESS_KEY_ID must be configured when using S3 storage."
+        )
+
+    if not settings.SCREENSHOT_S3_SECRET_ACCESS_KEY:
+        raise RuntimeError(
+            "SCREENSHOT_S3_SECRET_ACCESS_KEY must be configured when using S3 storage."
+        )
+
+    return boto3.client(
+        "s3",
+        endpoint_url=settings.SCREENSHOT_S3_ENDPOINT_URL,
+        region_name=settings.SCREENSHOT_S3_REGION,
+        aws_access_key_id=settings.SCREENSHOT_S3_ACCESS_KEY_ID,
+        aws_secret_access_key=settings.SCREENSHOT_S3_SECRET_ACCESS_KEY,
+        config=Config(
+            s3={
+                "addressing_style": (
+                    "path" if settings.SCREENSHOT_S3_FORCE_PATH_STYLE else "auto"
+                )
+            }
+        ),
+    )
+
+
+def get_screenshot_local_folder(version: str) -> str:
+    if is_s3_screenshot_storage():
+        return str(Path(TEMPORARY_SCREENSHOT_ROOT) / version)
+
+    settings = Settings()
+    return str(_get_root_path(settings.SCREENSHOT_FILESYSTEM_PATH) / version)
+
+
+def get_screenshot_local_path(version: str, screenshot_id: str) -> str:
+    return str(Path(get_screenshot_local_folder(version)) / f"{screenshot_id}.webp")
+
+
+def get_screenshot_object_key(version: str, screenshot_id: str) -> str:
+    settings = Settings()
+    prefix = settings.SCREENSHOT_S3_PREFIX.strip("/")
+    filename = f"{version}/{screenshot_id}.webp"
+    return f"{prefix}/{filename}" if prefix else filename
+
+
+def persist_screenshot(screenshot: ScreenShot, version: str) -> None:
+    if not is_s3_screenshot_storage():
+        return
+
+    settings = Settings()
+    screenshot_path = Path(screenshot.get_webp())
+    client = _get_s3_client()
+    bucket = settings.SCREENSHOT_S3_BUCKET
+    try:
+        client.head_bucket(Bucket=bucket)
+    except ClientError as exc:
+        error_code = exc.response.get("Error", {}).get("Code")
+        if error_code not in {"404", "NoSuchBucket", "NotFound"}:
+            raise
+        client.create_bucket(Bucket=bucket)
+
+    client.upload_file(
+        str(screenshot_path),
+        bucket,
+        get_screenshot_object_key(version=version, screenshot_id=screenshot.id),
+        ExtraArgs={"ContentType": "image/webp"},
+    )
+    screenshot_path.unlink(missing_ok=True)
+
+
+def screenshot_exists(version: str, screenshot_id: str) -> bool:
+    if not is_s3_screenshot_storage():
+        return Path(get_screenshot_local_path(version, screenshot_id)).is_file()
+
+    settings = Settings()
+    try:
+        _get_s3_client().head_object(
+            Bucket=settings.SCREENSHOT_S3_BUCKET,
+            Key=get_screenshot_object_key(version=version, screenshot_id=screenshot_id),
+        )
+        return True
+    except ClientError as exc:
+        error_code = exc.response.get("Error", {}).get("Code")
+        if error_code in {"404", "NoSuchKey", "NotFound"}:
+            return False
+        raise
+
+
+def read_screenshot(version: str, screenshot_id: str) -> bytes:
+    if not is_s3_screenshot_storage():
+        return Path(get_screenshot_local_path(version, screenshot_id)).read_bytes()
+
+    settings = Settings()
+    response = _get_s3_client().get_object(
+        Bucket=settings.SCREENSHOT_S3_BUCKET,
+        Key=get_screenshot_object_key(version=version, screenshot_id=screenshot_id),
+    )
+    return response["Body"].read()

projects/ecoindex_api/.env.template

@@ -9,6 +9,15 @@
 # DEBUG=1
 # ENABLE_SCREENSHOT=1
 # EXCLUDED_HOSTS='["localhost","127.0.0.1"]'
+# SCREENSHOT_FILESYSTEM_PATH=/code/screenshots
+# SCREENSHOT_STORAGE_TYPE=s3
+# SCREENSHOT_S3_ACCESS_KEY_ID=ecoindex-access-key
+# SCREENSHOT_S3_BUCKET=ecoindex-screenshots
+# SCREENSHOT_S3_ENDPOINT_URL=http://localhost:9000
+# SCREENSHOT_S3_FORCE_PATH_STYLE=1
+# SCREENSHOT_S3_PREFIX=screenshots
+# SCREENSHOT_S3_REGION=us-east-1
+# SCREENSHOT_S3_SECRET_ACCESS_KEY=ecoindex-secret-key-change-me
 # SENTRY_DSN=
 # SENTRY_ENVIRONMENT=production
 # SENTRY_TRACES_SAMPLE_RATE=0.1