feat: add MongoBuildCleanupManager and related commands for managing MongoDB builds

Author: DylanWelzel

biothings/hub/__init__.py

@@ -701,6 +701,7 @@ def configure_index_manager(self):
 
     def configure_snapshot_manager(self):
         assert "index" in self.features, "'snapshot' feature requires 'index'"
+        from biothings.hub.dataindex.mongo_build_cleanup import MongoBuildCleanupManager
         from biothings.hub.dataindex.snapshooter import SnapshotManager
 
         args = self.mixargs("snapshot")
@@ -713,6 +714,7 @@ def configure_snapshot_manager(self):
         snapshot_manager.configure(config.SNAPSHOT_CONFIG)
         snapshot_manager.poll("snapshot", snapshot_manager.snapshot_a_build)
         self.managers["snapshot_manager"] = snapshot_manager
+        self.managers["mongo_build_cleanup_manager"] = MongoBuildCleanupManager(job_manager=self.managers["job_manager"])
 
     def configure_auto_snapshot_cleaner_manager(self):
         assert "snapshot" in self.features, "'auto_snapshot_cleaner' feature requires 'snapshot'"
@@ -1148,6 +1150,9 @@ def configure_commands(self):
             self.commands["list_snapshots"] = self.managers["snapshot_manager"].list_snapshots
             self.commands["delete_snapshots"] = self.managers["snapshot_manager"].delete_snapshots
             self.commands["validate_snapshots"] = self.managers["snapshot_manager"].validate_snapshots
+        if self.managers.get("mongo_build_cleanup_manager"):
+            self.commands["list_mongo_builds"] = self.managers["mongo_build_cleanup_manager"].list_mongo_builds
+            self.commands["delete_mongo_builds"] = self.managers["mongo_build_cleanup_manager"].delete_mongo_builds
         # data release commands
         if self.managers.get("release_manager"):
             self.commands["create_release_note"] = self.managers["release_manager"].create_release_note
@@ -1514,6 +1519,12 @@ def configure_api_endpoints(self):
             )
         if "validate_snapshots" in cmdnames:
             self.api_endpoints["validate_snapshots"] = EndpointDefinition(name="validate_snapshots", method="post")
+        if "list_mongo_builds" in cmdnames:
+            self.api_endpoints["list_mongo_builds"] = EndpointDefinition(name="list_mongo_builds", method="get")
+        if "delete_mongo_builds" in cmdnames:
+            self.api_endpoints["delete_mongo_builds"] = EndpointDefinition(
+                name="delete_mongo_builds", method="put", force_bodyargs=True
+            )
         if "sync" in cmdnames:
             self.api_endpoints["sync"] = EndpointDefinition(name="sync", method="post", force_bodyargs=True)
         if "whatsnew" in cmdnames:

biothings/hub/dataindex/mongo_build_cleanup.py

@@ -0,0 +1,74 @@
+from functools import partial
+
+from biothings.hub.manager import BaseManager
+from biothings.utils.hub_db import get_src_build
+from config import logger as logging
+
+
+class MongoBuildCleaner:
+    def __init__(self, job_manager):
+        self.job_manager = job_manager
+
+    def list_builds(self, build_config=None, build_name=None):
+        collection = get_src_build()
+
+        filters = {}
+        if build_config:
+            filters["build_config._id"] = build_config
+        if build_name:
+            filters["_id"] = build_name
+
+        projection = {
+            "_id": 1,
+            "build_config": 1,
+            "started_at": 1,
+            "archived": 1,
+            "target_name": 1,
+        }
+        builds = list(collection.find(filters, projection).sort("started_at", -1))
+
+        grouped = {}
+        for build in builds:
+            group_name = build.get("build_config", {}).get("_id") or "N/A"
+            grouped.setdefault(group_name, []).append(build)
+
+        return [{"_id": key, "items": items} for key, items in grouped.items()]
+
+    async def delete_builds(self, build_ids):
+        if not build_ids:
+            return {"deleted_count": 0}
+
+        from biothings.utils import mongo
+
+        conn = mongo.get_hub_db_async_conn()
+        try:
+            src_build = mongo.get_src_build_async(conn)
+            result = await src_build.delete_many({"_id": {"$in": build_ids}})
+            return {"deleted_count": result.deleted_count}
+        finally:
+            conn.close()
+
+    def done(self, future):
+        try:
+            result = future.result()
+            logging.info("Deleted %d MongoDB builds", result.get("deleted_count", 0), extra={"notify": True})
+        except Exception as exc:
+            logging.exception("Failed to delete MongoDB builds: %s", exc, extra={"notify": True})
+
+
+class MongoBuildCleanupManager(BaseManager):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.cleaner = MongoBuildCleaner(self.job_manager)
+
+    def list_mongo_builds(self, build_config=None, build_name=None):
+        return self.cleaner.list_builds(build_config=build_config, build_name=build_name)
+
+    def delete_mongo_builds(self, build_ids):
+        try:
+            job = self.job_manager.submit(partial(self.cleaner.delete_builds, build_ids))
+            job.add_done_callback(self.cleaner.done)
+        except Exception as ex:
+            logging.exception("Error while submitting MongoDB build deletion job: %s", ex, extra={"notify": True})
+            raise
+        return job

biothings/hub/dataindex/snapshooter.py

@@ -8,9 +8,8 @@
 from functools import partial
 
 import boto3
-from config import logger as logging
 from elasticsearch import Elasticsearch
-from elasticsearch.exceptions import TransportError, NotFoundError
+from elasticsearch.exceptions import NotFoundError, TransportError
 
 from biothings import config as btconfig
 from biothings.hub import SNAPSHOOTER_CATEGORY
@@ -22,6 +21,7 @@
 from biothings.utils.hub import template_out
 from biothings.utils.hub_db import get_src_build
 from biothings.utils.loggers import get_logger
+from config import logger as logging
 
 from . import snapshot_cleanup as cleaner, snapshot_registrar as registrar
 from .snapshot_repo import Repository

biothings/utils/mongo.py

@@ -10,7 +10,7 @@
 
 import bson
 import dateutil.parser as date_parser
-from pymongo import DESCENDING, MongoClient
+from pymongo import DESCENDING, AsyncMongoClient, MongoClient
 from pymongo.client_session import ClientSession
 from pymongo.collection import Collection as PymongoCollection
 from pymongo.database import Database as PymongoDatabase
@@ -155,6 +155,10 @@ def __getitem__(self, name):
         return Database(self, name)
 
 
+class AsyncDatabaseClient(AsyncMongoClient):
+    pass
+
+
 def requires_config(func):
     @wraps(func)
     def func_wrapper(*args, **kwargs):
@@ -192,6 +196,12 @@ def get_hub_db_conn():
     return conn
 
 
+@requires_config
+def get_hub_db_async_conn():
+    conn = AsyncDatabaseClient(config.HUB_DB_BACKEND["uri"])
+    return conn
+
+
 @requires_config
 def get_src_conn():
     return get_conn(config.DATA_SRC_SERVER, getattr(config, "DATA_SRC_PORT", 27017))
@@ -221,6 +231,12 @@ def get_src_build(conn=None):
     return conn[config.DATA_HUB_DB_DATABASE][config.DATA_SRC_BUILD_COLLECTION]
 
 
+@requires_config
+def get_src_build_async(conn=None):
+    conn = conn or get_hub_db_async_conn()
+    return conn[config.DATA_HUB_DB_DATABASE][config.DATA_SRC_BUILD_COLLECTION]
+
+
 @requires_config
 def get_src_build_config(conn=None):
     conn = conn or get_hub_db_conn()

pyproject.toml

@@ -86,11 +86,11 @@ opensearch = [
 ]
 # minimal requirements for running biothings.hub, e.g. in CLI mode
 hubcore = [
-    "pymongo>=4.1.0,<5.0",  # support MongoDB 5.0 since v3.12.0
+    "pymongo>=4.13.0,<5.0",  # AsyncMongoClient stable since 4.13.0
 ]
 # extra requirements to run a full biothings.hub
 hub = [
-    "pymongo>=4.1.0,<5.0",
+    "pymongo>=4.13.0,<5.0",
     "beautifulsoup4",  # used in dumper.GoogleDriveDumper
     "aiocron==1.8",  # setup scheduled jobs
     # "aiohttp==3.8.4",  # elasticsearch requires aiohttp>=3,<4
@@ -120,7 +120,7 @@ hub = [
 ]
 # minimal requirements for to run biothings CLI
 cli = [
-    "pymongo>=4.1.0,<5.0",  # support MongoDB 5.0 since v3.12.0
+    "pymongo>=4.13.0,<5.0",  # AsyncMongoClient stable since 4.13.0
     "psutil",
     "jsonschema>=2.6.0",
     "typer>=0.17.0",  # required for CLI, also installs rich package