48 lines
1.3 KiB
Python
48 lines
1.3 KiB
Python
from config import logger, states, tasks
|
||
|
||
|
||
def save(doc_id: str, *, task: str | None = None, workspace: str | None = None) -> None:
|
||
"""
|
||
Upsert a document with _id=doc_id, setting any of the provided fields.
|
||
Leave fields you don’t pass unchanged.
|
||
"""
|
||
updates: dict = {}
|
||
if task is not None:
|
||
updates["task"] = task
|
||
if workspace is not None:
|
||
updates["workspace"] = workspace
|
||
if updates:
|
||
states.update_one(
|
||
{"_id": doc_id},
|
||
{"$set": updates},
|
||
upsert=True,
|
||
)
|
||
|
||
|
||
def retrieve(doc_id: str) -> dict[str, str | None]:
|
||
"""
|
||
Fetches the document with _id=doc_id and returns its 'task' and 'workspace'.
|
||
If the document doesn’t exist, both will be None.
|
||
"""
|
||
doc = states.find_one({"_id": doc_id})
|
||
return {
|
||
"task": doc.get("task") if doc else None,
|
||
"workspace": doc.get("workspace") if doc else None,
|
||
}
|
||
|
||
|
||
def init_work_state(wd: dict):
|
||
"""
|
||
init work states with default values
|
||
"""
|
||
if not states.find_one({"_id": "work"}):
|
||
states.insert_one(
|
||
{
|
||
"_id": "work",
|
||
**{
|
||
wd[k]: tasks.find_one({"path": f"work/{wd[k]}"})["task_id"]
|
||
for k in wd
|
||
},
|
||
}
|
||
)
|