Add tags support

PythonCoderAS · PythonCoderAS · commit cb524f87a0c5 · 2024-01-17T20:53:56.000-05:00
diff --git a/alembic/versions/9a098a265cbd_add_batch_tags.py b/alembic/versions/9a098a265cbd_add_batch_tags.py
@@ -0,0 +1,92 @@
+"""Add batch tags
+
+Revision ID: 9a098a265cbd
+Revises: f79ac07ca876
+Create Date: 2024-01-17 15:29:50.515286
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = "9a098a265cbd"
+down_revision: Union[str, None] = "f79ac07ca876"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        "batch_tags",
+        sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
+        sa.Column("name", sa.String(length=256), nullable=False),
+        sa.PrimaryKeyConstraint("id"),
+    )
+    op.create_index(op.f("ix_batch_tags_name"), "batch_tags", ["name"], unique=True)
+    op.create_table(
+        "batch_tag_batches",
+        sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
+        sa.Column("batch_id", sa.Integer(), nullable=False),
+        sa.Column("batch_tag_id", sa.Integer(), nullable=False),
+        sa.ForeignKeyConstraint(
+            ["batch_id"],
+            ["batches.id"],
+        ),
+        sa.ForeignKeyConstraint(
+            ["batch_tag_id"],
+            ["batch_tags.id"],
+        ),
+        sa.PrimaryKeyConstraint("id"),
+        sa.UniqueConstraint("batch_id", "batch_tag_id", name="_batch_batch_tag_uc"),
+    )
+    op.create_index(
+        op.f("ix_batch_tag_batches_batch_id"),
+        "batch_tag_batches",
+        ["batch_id"],
+        unique=False,
+    )
+    op.create_index(
+        op.f("ix_batch_tag_batches_batch_tag_id"),
+        "batch_tag_batches",
+        ["batch_tag_id"],
+        unique=False,
+    )
+    op.create_index(
+        op.f("ix_batches_created_at"), "batches", ["created_at"], unique=False
+    )
+    op.create_index(op.f("ix_jobs_completed"), "jobs", ["completed"], unique=False)
+    op.create_index(op.f("ix_jobs_created_at"), "jobs", ["created_at"], unique=False)
+    op.create_index(
+        op.f("ix_jobs_delayed_until"), "jobs", ["delayed_until"], unique=False
+    )
+    op.create_index(op.f("ix_jobs_failed"), "jobs", ["failed"], unique=False)
+    op.create_index(
+        op.f("ix_repeat_urls_created_at"), "repeat_urls", ["created_at"], unique=False
+    )
+    op.create_index(op.f("ix_urls_first_seen"), "urls", ["first_seen"], unique=False)
+    op.create_index(op.f("ix_urls_last_seen"), "urls", ["last_seen"], unique=False)
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_index(op.f("ix_urls_last_seen"), table_name="urls")
+    op.drop_index(op.f("ix_urls_first_seen"), table_name="urls")
+    op.drop_index(op.f("ix_repeat_urls_created_at"), table_name="repeat_urls")
+    op.drop_index(op.f("ix_jobs_failed"), table_name="jobs")
+    op.drop_index(op.f("ix_jobs_delayed_until"), table_name="jobs")
+    op.drop_index(op.f("ix_jobs_created_at"), table_name="jobs")
+    op.drop_index(op.f("ix_jobs_completed"), table_name="jobs")
+    op.drop_index(op.f("ix_batches_created_at"), table_name="batches")
+    op.drop_index(
+        op.f("ix_batch_tag_batches_batch_tag_id"), table_name="batch_tag_batches"
+    )
+    op.drop_index(op.f("ix_batch_tag_batches_batch_id"), table_name="batch_tag_batches")
+    op.drop_table("batch_tag_batches")
+    op.drop_index(op.f("ix_batch_tags_name"), table_name="batch_tags")
+    op.drop_table("batch_tags")
+    # ### end Alembic commands ###
diff --git a/src/main.py b/src/main.py
@@ -17,6 +17,7 @@
 )
 
 import sqlalchemy
+import sqlalchemy.ext.asyncio
 from aiohttp import ClientSession
 from fastapi import Depends, FastAPI, HTTPException, Query, Request, Response
 from fastapi.middleware.cors import CORSMiddleware
diff --git a/src/models.py b/src/models.py
@@ -1,4 +1,6 @@
 import datetime
+from typing import Self
+from sqlalchemy import select
 from sqlalchemy.orm import Mapped, mapped_column
 import sqlalchemy.ext.asyncio
 import sqlalchemy.orm
@@ -12,6 +14,71 @@ class Base(
     pass
 
 
+class BatchTag(Base):
+    __tablename__ = "batch_tags"
+    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True, init=False)
+    name: Mapped[str] = mapped_column(
+        sqlalchemy.String(length=256), unique=True, index=True
+    )
+
+    batches: Mapped[list["Batch"]] = sqlalchemy.orm.relationship(
+        "Batch",
+        secondary="batch_tag_batches",
+        back_populates="tags",
+        init=False,
+        repr=False,
+    )
+
+    @classmethod
+    async def resolve_list(cls, names: set[str]) -> list[Self]:
+        from .main import async_session
+
+        if not names:
+            return []
+
+        async with async_session() as session, session.begin():
+            stmt = select(cls).where(cls.name.in_(names))
+            result = (await session.scalars(stmt)).all()
+            seen = {r.name: r for r in result}
+            seen_set = set(seen.keys())
+            missing = names - seen_set
+            if missing:
+                newly_created: list[BatchTag] = [
+                    BatchTag(name=name) for name in missing
+                ]
+                session.add_all(newly_created)
+                return [*seen.values(), *newly_created]
+            return list(seen.values())
+
+
+class BatchTagBatch(Base):
+    __tablename__ = "batch_tag_batches"
+    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True, init=False)
+    batch_id: Mapped[int] = mapped_column(
+        sqlalchemy.ForeignKey("batches.id"), index=True
+    )
+    batch_tag_id: Mapped[int] = mapped_column(
+        sqlalchemy.ForeignKey(BatchTag.id), index=True
+    )
+
+    batch: Mapped["Batch"] = sqlalchemy.orm.relationship(
+        "Batch", lazy="joined", innerjoin=True, foreign_keys=[batch_id], viewonly=True
+    )
+    batch_tag: Mapped["BatchTag"] = sqlalchemy.orm.relationship(
+        "BatchTag",
+        lazy="joined",
+        innerjoin=True,
+        foreign_keys=[batch_tag_id],
+        viewonly=True,
+    )
+
+    __table_args__ = (
+        sqlalchemy.UniqueConstraint(
+            "batch_id", "batch_tag_id", name="_batch_batch_tag_uc"
+        ),
+    )
+
+
 class Batch(Base):
     __tablename__ = "batches"
 
@@ -23,10 +90,19 @@ class Batch(Base):
         init=False,
         index=True,
     )
+    locked: Mapped[bool] = mapped_column(
+        default=False
+    )  # Indicates that a batch is locked (no more jobs can be added to it)
 
     jobs: Mapped[list["Job"]] = sqlalchemy.orm.relationship(
         "Job", secondary="batch_jobs", back_populates="batches", init=False, repr=False
     )
+    tags: Mapped[list[BatchTag]] = sqlalchemy.orm.relationship(
+        BatchTag,
+        secondary="batch_tag_batches",
+        back_populates="batches",
+        default_factory=list,
+    )
 
 
 class URL(Base):
@@ -119,9 +195,6 @@ class Job(Base):
     id: Mapped[int] = mapped_column(
         sqlalchemy.BigInteger, primary_key=True, autoincrement=True, init=False
     )
-    batches: Mapped[list[Batch]] = sqlalchemy.orm.relationship(
-        Batch, secondary="batch_jobs", back_populates="jobs"
-    )
     url_id: Mapped[int] = mapped_column(
         sqlalchemy.ForeignKey(URL.id),
         nullable=False,
@@ -132,6 +205,9 @@ class Job(Base):
     url: Mapped[URL] = sqlalchemy.orm.relationship(
         URL, lazy="joined", innerjoin=True, foreign_keys=[url_id], back_populates="jobs"
     )
+    batches: Mapped[list[Batch]] = sqlalchemy.orm.relationship(
+        Batch, secondary="batch_jobs", back_populates="jobs", default_factory=list
+    )
     created_at: Mapped[datetime.datetime] = mapped_column(
         sqlalchemy.DateTime(timezone=True),
         server_default=sqlalchemy.sql.func.now(),
diff --git a/src/routes/batch/__init__.py b/src/routes/batch/__init__.py
@@ -1,6 +1,6 @@
 import datetime
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 import sqlalchemy
 from sqlalchemy import select
 
@@ -20,6 +20,7 @@ class BatchReturn(BaseModel):
     created_at: datetime.datetime
     repeat_url: int | None = None
     jobs: int
+    tags: list[str] = Field(default_factory=list)
 
 
 @app.get("/batch")
diff --git a/src/routes/batch/create/__init__.py b/src/routes/batch/create/__init__.py
@@ -1,10 +1,11 @@
 import datetime
+from itertools import batched
 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 ....models import Job, URL, Batch, BatchTag
 from ....main import app, async_session
 
 
@@ -35,44 +36,50 @@ async def add_filled_batch(
     items: Iterable[BatchItem],
     *,
     priority: int = 0,
-    tags: list[str],
+    tags: Iterable[str],
 ) -> QueueBatchReturn:
-    urls = [item.url for item in items]
+    items_batched = batched(items, 30000)
+    job_count = 0
     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
+        batch = Batch(tags=await BatchTag.resolve_list(set(tags)))
+        for item_set in items_batched:
+            urls = [item.url for item in item_set]
             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)
+            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 item_set:
+                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)
+            job_count += len(jobs)
 
-    return QueueBatchReturn(batch_id=batch.id, job_count=len(jobs))
+    return QueueBatchReturn(batch_id=batch.id, job_count=job_count)
 
 
 @app.post("/batch/create")
diff --git a/src/routes/batch/create/gsheets_archive.py b/src/routes/batch/create/gsheets_archive.py
@@ -32,6 +32,7 @@ class CSVBatchItem(TypedDict):
 @app.post("/batch/create/gsheets_archive")
 async def create_batch_gsheets_archive(
     file: UploadFile,
+    exclude_ratelimited_daily: bool = True,
     priority: int = 0,
     tags: list[str] | None = None,
     created_override: datetime.datetime | None = None,
@@ -53,6 +54,11 @@ async def create_batch_gsheets_archive(
                 )
             )
         elif item["archive_message"]:  # A message other than an URL = failure
+            if (
+                "You cannot make more than" in item["archive_message"]
+            ):  # We got ratelimited for too many captures in a day
+                if exclude_ratelimited_daily:
+                    continue
             items.append(
                 BatchItem(
                     url=url,
diff --git a/src/routes/queue/batch/__init__.py b/src/routes/queue/batch/__init__.py

Original file line number	Diff line number	Diff line change
`@@ -17,6 +17,7 @@`
`17`	`17`	`)`
`18`	`18`
`19`	`19`	`import sqlalchemy`
	`20`	`+import sqlalchemy.ext.asyncio`
`20`	`21`	`from aiohttp import ClientSession`
`21`	`22`	`from fastapi import Depends, FastAPI, HTTPException, Query, Request, Response`
`22`	`23`	`from fastapi.middleware.cors import CORSMiddleware`