Spaces:
Sleeping
Sleeping
File size: 3,015 Bytes
a4cc2fd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 | 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}")
|