browsertrix/backend/btrixcloud/migrations/migration_0048_first_seed_seed_count.py
Tessa Walsh f7ba712646 Add seed file support to Browsertrix backend (#2710)
Fixes #2673 

Changes in this PR:

- Adds a new `file_uploads.py` module and corresponding `/files` API
prefix with methods/endpoints for uploading, GETing, and deleting seed
files (can be extended to other types of files moving forward)
- Seed files are supported via `CrawlConfig.config.seedFileId` on POST
and PATCH endpoints. This seedFileId is replaced by a presigned url when
passed to the crawler by the operator
- Seed files are read when first uploaded to calculate `firstSeed` and
`seedCount` and store them in the database, and this is copied into the
workflow and crawl documents when they are created.
- Logic is added to store `firstSeed` and `seedCount` for other
workflows as well, and a migration added to backfill data, to maintain
consistency and fix some of the pymongo aggregations that previously
assumed all workflows would have at least one `Seed` object in
`CrawlConfig.seeds`
- Seed file and thumbnail storage stats are added to org stats
- Seed file and thumbnail uploads first check that the org's storage
quota has not been exceeded and return a 400 if so
- A cron background job (run weekly each Sunday at midnight by default,
but configurable) is added to look for seed files at least x minutes old
(1440 minutes, or 1 day, by default, but configurable) that are not in
use in any workflows, and to delete them when they are found. The
backend pods will ensure this k8s batch job exists when starting up and
create it if it does not already exist. A database entry for each run of
the job is created in the operator on job completion so that it'll
appear in the `/jobs` API endpoints, but retrying of this type of
regularly scheduled background job is not supported as we don't want to
accidentally create multiple competing scheduled jobs.
- Adds a `min_seed_file_crawler_image` value to the Helm chart that is
checked before creating a crawl from a workflow if set. If a workflow
cannot be run, return the detail of the exception in
`CrawlConfigAddedResponse.errorDetail` so that we can display the reason
in the frontend
- Add SeedFile model from base UserFile (former ImageFIle), ensure all APIs
returning uploaded files return an absolute pre-signed URL (either with external origin or internal service origin)

---------
Co-authored-by: Ilya Kreymer <ikreymer@gmail.com>
2025-07-22 19:11:02 -07:00

96 lines
3.1 KiB
Python

"""
Migration 0048 - Calculate firstSeed/seedCount and store directly in database
"""
from typing import cast, List, Dict, Any
from btrixcloud.migrations import BaseMigration
from btrixcloud.models import CrawlConfig, Crawl, Seed
MIGRATION_VERSION = "0048"
# pylint: disable=duplicate-code
class Migration(BaseMigration):
"""Migration class."""
# pylint: disable=unused-argument
def __init__(self, mdb, **kwargs):
super().__init__(mdb, migration_version=MIGRATION_VERSION)
async def migrate_up(self) -> None:
"""Perform migration up.
Calculate firstSeed and seedCount for workflows and store in db
"""
crawls_mdb = self.mdb["crawls"]
crawl_configs_mdb = self.mdb["crawl_configs"]
match_query = {"$or": [{"firstSeed": None}, {"seedCount": None}]}
# Workflows
async for config_raw in crawl_configs_mdb.find(match_query):
config = CrawlConfig.from_dict(config_raw)
try:
if not config.config.seeds:
print(
f"Unable to find seeds for config {config.id}, skipping",
flush=True,
)
continue
seeds = cast(List[Seed], config.config.seeds)
seed_count = len(seeds)
first_seed = seeds[0].url
await crawl_configs_mdb.find_one_and_update(
{"_id": config.id},
{
"$set": {
"firstSeed": first_seed,
"seedCount": seed_count,
}
},
)
# pylint: disable=broad-exception-caught
except Exception as err:
print(
f"Unable to update seed info for workflow {config.id}: {err}",
flush=True,
)
# Crawls
crawl_query: Dict[str, Any] = {
"type": "crawl",
"$or": [
{"firstSeed": {"$in": [None, ""]}},
{"seedCount": {"$in": [None, 0]}},
],
}
async for crawl_raw in crawls_mdb.find(crawl_query):
crawl_id = crawl_raw["_id"]
try:
crawl = Crawl.from_dict(crawl_raw)
config_raw = await crawl_configs_mdb.find_one({"_id": crawl.cid})
seed_count = config_raw.get("seedCount", 0)
first_seed = config_raw.get("firstSeed", "")
await crawls_mdb.find_one_and_update(
{"_id": crawl_id},
{
"$set": {
"firstSeed": first_seed,
"seedCount": seed_count,
}
},
)
# pylint: disable=broad-exception-caught
except Exception as err:
print(
f"Unable to update seed info for crawl {crawl_id}: {err}",
flush=True,
)