rowboat/apps/simulation_runner/db.py

from pymongo import MongoClient
from bson import ObjectId
import os
from datetime import datetime, timedelta, timezone
from scenario_types import SimulationRun, Scenario, SimulationResult, SimulationAggregateResult

MONGO_URI = os.environ.get("MONGODB_URI", "mongodb://localhost:27017/rowboat").strip()

SCENARIOS_COLLECTION_NAME = "scenarios"
API_KEYS_COLLECTION = "api_keys"
SIMULATIONS_COLLECTION_NAME = "simulation_runs"
SIMULATION_RESULT_COLLECTION_NAME = "simulation_result"
SIMULATION_AGGREGATE_RESULT_COLLECTION_NAME = "simulation_aggregate_result"

def get_db():
    client = MongoClient(MONGO_URI)
    return client["rowboat"]

def get_collection(collection_name: str):
    db = get_db()
    return db[collection_name]

def get_api_key(project_id: str):
    collection = get_collection(API_KEYS_COLLECTION)
    doc = collection.find_one({"projectId": project_id})
    if doc:
        return doc["key"]
    else:
        return None

def get_pending_simulation_run():
    collection = get_collection(SIMULATIONS_COLLECTION_NAME)
    doc = collection.find_one_and_update(
        {"status": "pending"},
        {"$set": {"status": "running"}},
        return_document=True
    )
    if doc:
        return SimulationRun(
            id=str(doc["_id"]),
            projectId=doc["projectId"],
            status="running",
            scenarioIds=doc["scenarioIds"],
            workflowId=doc["workflowId"],
            startedAt=doc["startedAt"],
            completedAt=doc.get("completedAt")
        )
    return None

def set_simulation_run_to_completed(simulation_run: SimulationRun, aggregate_result: SimulationAggregateResult):
    collection = get_collection(SIMULATIONS_COLLECTION_NAME)
    collection.update_one({"_id": ObjectId(simulation_run.id)}, {"$set": {"status": "completed", "aggregateResults": aggregate_result.model_dump(by_alias=True)}})

def get_scenarios_for_run(simulation_run: SimulationRun):
    if simulation_run is None:
        return []
    collection = get_collection(SCENARIOS_COLLECTION_NAME)
    scenarios = []
    for doc in collection.find():
        if doc["_id"] in [ObjectId(sid) for sid in simulation_run.scenarioIds]:
            scenarios.append(Scenario(
                id=str(doc["_id"]),
                projectId=doc["projectId"],
                name=doc["name"],
                description=doc["description"],
                criteria=doc["criteria"],
                context=doc["context"],
                createdAt=doc["createdAt"],
                lastUpdatedAt=doc["lastUpdatedAt"]
            ))
    return scenarios

def write_simulation_result(result: SimulationResult):
    collection = get_collection(SIMULATION_RESULT_COLLECTION_NAME)
    collection.insert_one(result.model_dump())

def update_simulation_run_heartbeat(simulation_run_id: str):
    """
    Updates the 'last_heartbeat' timestamp for a SimulationRun.
    """

    collection = get_collection(SIMULATIONS_COLLECTION_NAME)
    collection.update_one(
        {"_id": ObjectId(simulation_run_id)},
        {"$set": {"lastHeartbeat": datetime.now(timezone.utc)}}
    )

def mark_stale_jobs_as_failed():
    """
    Finds any job in 'running' status whose last_heartbeat is older than 5 minutes,
    and sets it to 'failed'.
    """

    collection = get_collection(SIMULATIONS_COLLECTION_NAME)
    stale_threshold = datetime.now(timezone.utc) - timedelta(minutes=20)
    result = collection.update_many(
        {
            "status": "running",
            "lastHeartbeat": {"$lt": stale_threshold}
        },
        {
            "$set": {"status": "failed"}
        }
    )
    return result.modified_count  # Number of jobs marked failed