Add new gsheets archive capability

Author: PythonCoderAS

Pipfile

@@ -9,6 +9,7 @@ aiohttp = "<4,>=3.9.0"
 asyncpg = ">=0.29.0"
 sqlalchemy = {extras = ["asyncio"], version = "*"}
 sentry-sdk = {extras = ["fastapi"], version = "*"}
+python-multipart = "*"
 
 [dev-packages]
 uvicorn = {extras = ["standard"] }

Pipfile.lock

@@ -38,6 +38,14 @@
     profiles_sample_rate=1.0,
 )
 
+archive_url_regex = re.compile(r"/web/(\d{14})")
+
+
+def get_archive_save_url_timestamp(timestamp: str) -> datetime.datetime:
+    return datetime.datetime.strptime(timestamp, "%Y%m%d%H%M%S").replace(
+        tzinfo=datetime.timezone.utc
+    )
+
 
 # Constants
 min_wait_time_between_archives = datetime.timedelta(hours=1)
@@ -125,12 +133,10 @@ async def url_worker():
                         allow_redirects=False,
                     ) as resp:
                         resp.raise_for_status()
-                        if match := re.search(
-                            r"/web/(\d{14})", resp.headers.get("Location", "")
+                        if match := archive_url_regex.search(
+                            resp.headers.get("Location", "")
                         ):
-                            saved_dt = datetime.datetime.strptime(
-                                match.group(1), "%Y%m%d%H%M%S"
-                            ).replace(tzinfo=datetime.timezone.utc)
+                            saved_dt = get_archive_save_url_timestamp(match.group(1))
                             async with session.begin():
                                 await session.execute(
                                     update(URL)

src/main.py

@@ -0,0 +1,80 @@
+import datetime
+from typing import Iterable
+from pydantic import BaseModel, Field, model_validator
+from sqlalchemy import select
+
+from src.routes.queue.batch import QueueBatchReturn
+from ....models import Job, URL, Batch
+from ....main import app, async_session
+
+
+class BatchItem(BaseModel):
+    url: str
+    completed: datetime.datetime | None = None
+    failed: datetime.datetime | None = None
+    created: datetime.datetime | None = None
+
+    @model_validator(mode="after")
+    def validate(self):
+        if self.completed and self.failed:
+            raise ValueError("Batch item cannot be both completed and failed")
+        if not self.completed and not self.failed:
+            raise ValueError("Batch item must be either completed or failed")
+
+    @property
+    def effective_creation_time(self) -> datetime.datetime:
+        return self.created or self.completed or self.failed
+
+
+class CreateBatchBody(BaseModel):
+    items: list[BatchItem]
+    tags: list[str] = Field(default_factory=list)
+
+
+async def add_filled_batch(
+    items: Iterable[BatchItem],
+    *,
+    priority: int = 0,
+    tags: list[str],
+) -> QueueBatchReturn:
+    urls = [item.url for item in items]
+    async with async_session() as session, session.begin():
+        batch = Batch()
+        stmt = select(URL).where(URL.url.in_(urls))
+        result = await session.scalars(stmt)
+        existing_urls_items = result.all()
+        existing_urls = {url.url for url in existing_urls_items}
+        new_urls = set(urls) - set(existing_urls)
+        if new_urls:
+            new_url_models = [URL(url=url) for url in new_urls]
+            session.add_all(new_url_models)
+            del new_urls, new_url_models, existing_urls, existing_urls_items
+            # Needed because bulk create doesn't return the created models with their IDs
+            stmt = select(URL).where(URL.url.in_(urls))
+            result = await session.scalars(stmt)
+            url_models = result.all()
+        else:
+            url_models = existing_urls_items
+            del new_urls, existing_urls, existing_urls_items
+        url_map: dict[str, URL] = {url.url: url for url in url_models}
+        del url_models
+        jobs = []
+        for item in items:
+            job = Job(
+                url=url_map[item.url],
+                batches=[batch],
+                priority=priority,
+                completed=item.completed,
+                failed=item.failed,
+            )
+            job.created_at = item.effective_creation_time
+            jobs.append(job)
+            url_map[item.url].last_seen = item.completed or url_map[item.url].last_seen
+        session.add_all(jobs)
+
+    return QueueBatchReturn(batch_id=batch.id, job_count=len(jobs))
+
+
+@app.post("/batch/create")
+async def create_batch(body: CreateBatchBody, priority: int = 0) -> QueueBatchReturn:
+    return await add_filled_batch(body.items, priority=priority, tags=body.tags)

src/routes/batch/create/__init__.py

@@ -0,0 +1,61 @@
+import datetime
+from typing import Literal
+from fastapi import UploadFile
+from . import add_filled_batch, BatchItem
+from csv import DictReader
+from typing import TypedDict
+
+from ...queue.batch import QueueBatchReturn
+from ....main import app, archive_url_regex, get_archive_save_url_timestamp
+
+heading = (
+    "url",
+    "capture_dupe_status",
+    "new_capture_status",
+    "archive_message",
+    "archive_info",
+    "first_archive_status",
+)
+
+
+class CSVBatchItem(TypedDict):
+    url: str
+    capture_dupe_status: Literal["New capture", "Already captured", ""]
+    new_capture_status: str | Literal[""]
+    archive_message: str
+    archive_info: str | Literal[""]
+    first_archive_status: Literal["First Archive", ""]
+
+
+@app.post("/batch/create/gsheets_archive")
+async def create_batch_gsheets_archive(
+    file: UploadFile,
+    priority: int = 0,
+    tags: list[str] | None = None,
+    created_override: datetime.datetime | None = None,
+) -> QueueBatchReturn:
+    if tags is None:
+        tags = []
+    tags.append("wayback-machine-gsheets-archive")
+    csv_items = (await file.read()).decode().splitlines(False)
+    reader = DictReader(csv_items, fieldnames=heading)
+    items_list: list[CSVBatchItem] = list(reader)
+    items: list[BatchItem] = []
+    for item in items_list:
+        url = item["url"]
+        if match := archive_url_regex.search(item["archive_message"]):
+            archived_timestamp = get_archive_save_url_timestamp(match.group(1))
+            items.append(
+                BatchItem(
+                    url=url, completed=archived_timestamp, created=created_override
+                )
+            )
+        elif item["archive_message"]:  # A message other than an URL = failure
+            items.append(
+                BatchItem(
+                    url=url,
+                    failed=datetime.datetime.now(tz=datetime.timezone.utc),
+                    created=created_override,
+                )
+            )
+    return await add_filled_batch(items, priority=priority, tags=tags)

src/routes/batch/create/gsheets_archive.py

@@ -1,27 +1,28 @@
-from pydantic import BaseModel
+from typing import Iterable
+from pydantic import BaseModel, Field
 from sqlalchemy import select
-from ...models import Job, URL, Batch
-from ...main import app, async_session
+from ....models import Job, URL, Batch
+from ....main import app, async_session
 
 
 class QueueBatchBody(BaseModel):
     urls: list[str]
-    priority: int = 0
-    unique_only: bool = True
+    tags: list[str] = Field(default_factory=list)
 
 
 class QueueBatchReturn(BaseModel):
     batch_id: int
     job_count: int
 
 
-@app.post("/queue/batch")
-async def queue_batch(body: QueueBatchBody) -> QueueBatchReturn:
-    batch = Batch()
-    urls = body.urls
-    if body.unique_only:
-        urls = set(urls)
+async def add_batch(
+    urls: Iterable[str],
+    *,
+    priority: int = 0,
+    tags: list[str],
+) -> QueueBatchReturn:
     async with async_session() as session, session.begin():
+        batch = Batch()
         stmt = select(URL).where(URL.url.in_(urls))
         result = await session.scalars(stmt)
         existing_urls_items = result.all()
@@ -42,7 +43,14 @@ async def queue_batch(body: QueueBatchBody) -> QueueBatchReturn:
         del url_models
         jobs = []
         for url in urls:
-            jobs.append(Job(url=url_map[url], batches=[batch], priority=body.priority))
+            jobs.append(Job(url=url_map[url], batches=[batch], priority=priority))
         session.add_all(jobs)
 
     return QueueBatchReturn(batch_id=batch.id, job_count=len(jobs))
+
+
+@app.post("/queue/batch")
+async def queue_batch(
+    body: QueueBatchBody, priority: int = 0, unique_only: bool = True
+) -> QueueBatchReturn:
+    return await add_batch(set(body.urls) if unique_only else body.urls, priority, unique_only)

src/routes/queue/batch.py renamed to src/routes/queue/batch/__init__.py

@@ -0,0 +1,14 @@
+from fastapi import UploadFile
+from pydantic import BaseModel, Field
+
+from . import QueueBatchReturn, add_batch
+from ....main import app
+
+class QueueBatchFileBody(BaseModel):
+    file: UploadFile
+    tags: list[str] = Field(default_factory=list)
+
+@app.post("/queue/batch/file")
+async def queue_batch_file(body: QueueBatchFileBody, priority: int = 0, unique_only: bool = False, ) -> QueueBatchReturn:
+    urls = (await body.file.read()).decode().splitlines(False)
+    return await add_batch(set(urls) if unique_only else urls, priority, unique_only, tags=body.tags)