mirror of
https://github.com/rowboatlabs/rowboat.git
synced 2026-04-25 00:16:29 +02:00
74 lines
No EOL
2.6 KiB
Python
74 lines
No EOL
2.6 KiB
Python
from pymongo import MongoClient
|
|
from bson import ObjectId
|
|
import os
|
|
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.get_default_database()
|
|
|
|
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()) |