Spaces:
Running
Running
| from __future__ import annotations | |
| import logging | |
| from datetime import datetime | |
| from typing import Any, Dict, List, Optional | |
| from bson import ObjectId | |
| from pymongo import DESCENDING | |
| from connections import get_runs_collection | |
| log = logging.getLogger(__name__) | |
| # ------------------- Index Setup ------------------- | |
| def _ensure_runs_indexes(col): | |
| try: | |
| col.create_index([("updated_at", DESCENDING)]) | |
| col.create_index([("flow", 1), ("updated_at", DESCENDING)]) | |
| col.create_index([("created_by", 1), ("updated_at", DESCENDING)]) | |
| except Exception: | |
| pass | |
| _rc = get_runs_collection() | |
| if _rc is not None: | |
| _ensure_runs_indexes(_rc) | |
| # ------------------- Runs ------------------- | |
| def upsert_run( | |
| run_id: str, | |
| *, | |
| flow: str, | |
| state_json: str, | |
| stages: List[Dict[str, Any]], | |
| urls: Dict[str, Any], | |
| created_by: Optional[str] = None, | |
| ) -> bool: | |
| """Insert or update a run. Returns False when there is no database.""" | |
| col = get_runs_collection() | |
| if col is None: | |
| log.error("run %s not saved: no database", run_id) | |
| return False | |
| now = datetime.utcnow() | |
| set_fields: Dict[str, Any] = { | |
| "flow": flow, | |
| "state": state_json, | |
| "stages": stages, | |
| "urls": urls, | |
| "updated_at": now, | |
| } | |
| try: | |
| col.update_one( | |
| {"_id": ObjectId(run_id)}, | |
| {"$set": set_fields, "$setOnInsert": {"created_by": created_by, "created_at": now}}, | |
| upsert=True, | |
| ) | |
| return True | |
| except Exception as e: | |
| log.error(f"upsert_run failed: {e}") | |
| return False | |
| def find_run(run_id: str) -> Optional[Dict[str, Any]]: | |
| col = get_runs_collection() | |
| if col is None: | |
| return None | |
| return col.find_one({"_id": ObjectId(run_id)}) | |
| def find_runs( | |
| *, | |
| flow: Optional[str] = None, | |
| start_date: Optional[datetime] = None, | |
| end_date: Optional[datetime] = None, | |
| limit: int = 200, | |
| created_by: Optional[str] = None, | |
| ) -> List[Dict[str, Any]]: | |
| """Runs newest first, optionally filtered by flow ("Research"/"Upload"), user and date.""" | |
| col = get_runs_collection() | |
| if col is None: | |
| return [] | |
| q: Dict[str, Any] = {"updated_at": {"$exists": True}} # skip docs from before this shape | |
| if flow: | |
| q["flow"] = flow | |
| if created_by: | |
| q["created_by"] = created_by | |
| if start_date or end_date: | |
| rng: Dict[str, Any] = {} | |
| if start_date: rng["$gte"] = start_date | |
| if end_date: rng["$lt"] = end_date | |
| q["updated_at"] = rng | |
| cur = col.find(q).sort("updated_at", DESCENDING).limit(max(1, int(limit))) | |
| out: List[Dict[str, Any]] = [] | |
| for d in cur: | |
| d["_id"] = str(d.get("_id")) | |
| out.append(d) | |
| return out | |
| def delete_run(run_id: str) -> None: | |
| col = get_runs_collection() | |
| if col is None: | |
| return | |
| try: | |
| col.delete_one({"_id": ObjectId(run_id)}) | |
| except Exception as e: | |
| log.error(f"delete_run failed: {e}") | |