Skip to content

Commit a5f41bc

Browse files
kraftpkraftp
authored
Update for Latest API (#12)
1 parent 4c2b93e commit a5f41bc

2 files changed

Lines changed: 126 additions & 43 deletions

File tree

‎dbos_mcp/client.py‎

Lines changed: 81 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import time
55
from pathlib import Path
66
from typing import Any
7+
from urllib.parse import quote
78

89
import httpx
910

@@ -39,6 +40,11 @@ def _save_credentials(credentials: dict[str, str]) -> None:
3940
CREDENTIALS_PATH.write_text(json.dumps(credentials, indent=2))
4041

4142

43+
def _path(param: str) -> str:
44+
"""Percent-encode a URL path parameter (workflow IDs may contain /, ?, #, etc.)."""
45+
return quote(param, safe="")
46+
47+
4248
def _get_credentials() -> dict[str, str]:
4349
"""Get credentials or raise if not logged in."""
4450
creds = _load_credentials()
@@ -158,6 +164,10 @@ async def list_workflows(
158164
authenticated_user: str | list[str] | None = None,
159165
start_time: str | None = None,
160166
end_time: str | None = None,
167+
completed_after: str | None = None,
168+
completed_before: str | None = None,
169+
dequeued_after: str | None = None,
170+
dequeued_before: str | None = None,
161171
status: str | list[str] | None = None,
162172
application_version: str | list[str] | None = None,
163173
forked_from: str | list[str] | None = None,
@@ -172,6 +182,7 @@ async def list_workflows(
172182
executor_id: str | list[str] | None = None,
173183
queues_only: bool | None = None,
174184
was_forked_from: bool | None = None,
185+
has_parent: bool | None = None,
175186
) -> list[dict[str, Any]]:
176187
"""List workflows with optional filters."""
177188
creds = _get_credentials()
@@ -187,6 +198,14 @@ async def list_workflows(
187198
body["start_time"] = start_time
188199
if end_time is not None:
189200
body["end_time"] = end_time
201+
if completed_after is not None:
202+
body["completed_after"] = completed_after
203+
if completed_before is not None:
204+
body["completed_before"] = completed_before
205+
if dequeued_after is not None:
206+
body["dequeued_after"] = dequeued_after
207+
if dequeued_before is not None:
208+
body["dequeued_before"] = dequeued_before
190209
if status is not None:
191210
body["status"] = status
192211
if application_version is not None:
@@ -215,10 +234,12 @@ async def list_workflows(
215234
body["queues_only"] = queues_only
216235
if was_forked_from is not None:
217236
body["was_forked_from"] = was_forked_from
237+
if has_parent is not None:
238+
body["has_parent"] = has_parent
218239

219240
async with httpx.AsyncClient() as client:
220241
response = await client.post(
221-
f"{CONDUCTOR_URL}/api/{creds['organization']}/applications/{application_name}/workflows/",
242+
f"{CONDUCTOR_URL}/api/{creds['organization']}/applications/{_path(application_name)}/workflows/",
222243
json=body,
223244
headers={
224245
"Content-Type": "application/json",
@@ -239,7 +260,7 @@ async def get_workflow(
239260
creds = _get_credentials()
240261
async with httpx.AsyncClient() as client:
241262
response = await client.get(
242-
f"{CONDUCTOR_URL}/api/{creds['organization']}/applications/{application_name}/workflows/{workflow_id}",
263+
f"{CONDUCTOR_URL}/api/{creds['organization']}/applications/{_path(application_name)}/workflows/{_path(workflow_id)}",
243264
headers={
244265
"Content-Type": "application/json",
245266
"Authorization": f"Bearer {creds['token']}",
@@ -254,12 +275,21 @@ async def get_workflow(
254275
async def list_steps(
255276
application_name: str,
256277
workflow_id: str,
278+
limit: int | None = None,
279+
offset: int | None = None,
257280
) -> list[dict[str, Any]]:
258281
"""Get execution steps for a workflow."""
259282
creds = _get_credentials()
283+
params: dict[str, Any] = {}
284+
if limit is not None:
285+
params["limit"] = limit
286+
if offset is not None:
287+
params["offset"] = offset
288+
260289
async with httpx.AsyncClient() as client:
261290
response = await client.get(
262-
f"{CONDUCTOR_URL}/api/{creds['organization']}/applications/{application_name}/workflows/{workflow_id}/steps",
291+
f"{CONDUCTOR_URL}/api/{creds['organization']}/applications/{_path(application_name)}/workflows/{_path(workflow_id)}/steps",
292+
params=params,
263293
headers={
264294
"Content-Type": "application/json",
265295
"Authorization": f"Bearer {creds['token']}",
@@ -278,7 +308,7 @@ async def list_executors(
278308
creds = _get_credentials()
279309
async with httpx.AsyncClient() as client:
280310
response = await client.get(
281-
f"{CONDUCTOR_URL}/api/{creds['organization']}/applications/{application_name}/executors",
311+
f"{CONDUCTOR_URL}/api/{creds['organization']}/applications/{_path(application_name)}/executors",
282312
headers={
283313
"Content-Type": "application/json",
284314
"Authorization": f"Bearer {creds['token']}",
@@ -293,12 +323,18 @@ async def list_executors(
293323
async def cancel_workflow(
294324
application_name: str,
295325
workflow_id: str,
326+
cancel_children: bool = False,
296327
) -> None:
297328
"""Cancel a workflow."""
298329
creds = _get_credentials()
330+
params: dict[str, Any] = {}
331+
if cancel_children:
332+
params["cancel_children"] = "true"
333+
299334
async with httpx.AsyncClient() as client:
300335
response = await client.post(
301-
f"{CONDUCTOR_URL}/api/{creds['organization']}/applications/{application_name}/workflows/{workflow_id}/cancel",
336+
f"{CONDUCTOR_URL}/api/{creds['organization']}/applications/{_path(application_name)}/workflows/{_path(workflow_id)}/cancel",
337+
params=params,
302338
headers={
303339
"Content-Type": "application/json",
304340
"Authorization": f"Bearer {creds['token']}",
@@ -311,12 +347,18 @@ async def cancel_workflow(
311347
async def resume_workflow(
312348
application_name: str,
313349
workflow_id: str,
350+
queue_name: str | None = None,
314351
) -> None:
315352
"""Resume a workflow."""
316353
creds = _get_credentials()
354+
body: dict[str, Any] = {}
355+
if queue_name is not None:
356+
body["queue_name"] = queue_name
357+
317358
async with httpx.AsyncClient() as client:
318359
response = await client.post(
319-
f"{CONDUCTOR_URL}/api/{creds['organization']}/applications/{application_name}/workflows/{workflow_id}/resume",
360+
f"{CONDUCTOR_URL}/api/{creds['organization']}/applications/{_path(application_name)}/workflows/{_path(workflow_id)}/resume",
361+
json=body,
320362
headers={
321363
"Content-Type": "application/json",
322364
"Authorization": f"Bearer {creds['token']}",
@@ -332,6 +374,8 @@ async def fork_workflow(
332374
start_step: int,
333375
application_version: str | None = None,
334376
new_workflow_id: str | None = None,
377+
queue_name: str | None = None,
378+
queue_partition_key: str | None = None,
335379
) -> dict[str, Any]:
336380
"""Fork a workflow from a specific step."""
337381
creds = _get_credentials()
@@ -340,10 +384,14 @@ async def fork_workflow(
340384
body["application_version"] = application_version
341385
if new_workflow_id is not None:
342386
body["new_workflow_id"] = new_workflow_id
387+
if queue_name is not None:
388+
body["queue_name"] = queue_name
389+
if queue_partition_key is not None:
390+
body["queue_partition_key"] = queue_partition_key
343391

344392
async with httpx.AsyncClient() as client:
345393
response = await client.post(
346-
f"{CONDUCTOR_URL}/api/{creds['organization']}/applications/{application_name}/workflows/{workflow_id}/fork",
394+
f"{CONDUCTOR_URL}/api/{creds['organization']}/applications/{_path(application_name)}/workflows/{_path(workflow_id)}/fork",
347395
json=body,
348396
headers={
349397
"Content-Type": "application/json",
@@ -359,13 +407,18 @@ async def fork_workflow(
359407
async def bulk_cancel_workflows(
360408
application_name: str,
361409
workflow_ids: list[str],
410+
cancel_children: bool = False,
362411
) -> None:
363412
"""Cancel multiple workflows."""
364413
creds = _get_credentials()
414+
body: dict[str, Any] = {"workflow_ids": workflow_ids}
415+
if cancel_children:
416+
body["cancel_children"] = True
417+
365418
async with httpx.AsyncClient() as client:
366419
response = await client.post(
367-
f"{CONDUCTOR_URL}/api/{creds['organization']}/applications/{application_name}/workflows/cancel",
368-
json={"workflow_ids": workflow_ids},
420+
f"{CONDUCTOR_URL}/api/{creds['organization']}/applications/{_path(application_name)}/workflows/cancel",
421+
json=body,
369422
headers={
370423
"Content-Type": "application/json",
371424
"Authorization": f"Bearer {creds['token']}",
@@ -388,7 +441,7 @@ async def bulk_resume_workflows(
388441

389442
async with httpx.AsyncClient() as client:
390443
response = await client.post(
391-
f"{CONDUCTOR_URL}/api/{creds['organization']}/applications/{application_name}/workflows/resume",
444+
f"{CONDUCTOR_URL}/api/{creds['organization']}/applications/{_path(application_name)}/workflows/resume",
392445
json=body,
393446
headers={
394447
"Content-Type": "application/json",
@@ -412,7 +465,7 @@ async def bulk_delete_workflows(
412465

413466
async with httpx.AsyncClient() as client:
414467
response = await client.post(
415-
f"{CONDUCTOR_URL}/api/{creds['organization']}/applications/{application_name}/workflows/delete",
468+
f"{CONDUCTOR_URL}/api/{creds['organization']}/applications/{_path(application_name)}/workflows/delete",
416469
json=body,
417470
headers={
418471
"Content-Type": "application/json",
@@ -454,7 +507,7 @@ async def fork_from_failure(
454507

455508
async with httpx.AsyncClient() as client:
456509
response = await client.post(
457-
f"{CONDUCTOR_URL}/api/{creds['organization']}/applications/{application_name}/workflows/fork-from-failure",
510+
f"{CONDUCTOR_URL}/api/{creds['organization']}/applications/{_path(application_name)}/workflows/fork-from-failure",
458511
json=body,
459512
headers={
460513
"Content-Type": "application/json",
@@ -481,7 +534,7 @@ async def delete_workflow(
481534

482535
async with httpx.AsyncClient() as client:
483536
response = await client.delete(
484-
f"{CONDUCTOR_URL}/api/{creds['organization']}/applications/{application_name}/workflows/{workflow_id}",
537+
f"{CONDUCTOR_URL}/api/{creds['organization']}/applications/{_path(application_name)}/workflows/{_path(workflow_id)}",
485538
params=params,
486539
headers={
487540
"Content-Type": "application/json",
@@ -559,7 +612,7 @@ async def get_workflow_aggregates(
559612

560613
async with httpx.AsyncClient() as client:
561614
response = await client.post(
562-
f"{CONDUCTOR_URL}/api/{creds['organization']}/applications/{application_name}/workflows/aggregates",
615+
f"{CONDUCTOR_URL}/api/{creds['organization']}/applications/{_path(application_name)}/workflows/aggregates",
563616
json=body,
564617
headers={
565618
"Content-Type": "application/json",
@@ -580,16 +633,15 @@ async def get_workflow_events(
580633
creds = _get_credentials()
581634
async with httpx.AsyncClient() as client:
582635
response = await client.get(
583-
f"{CONDUCTOR_URL}/api/{creds['organization']}/applications/{application_name}/workflows/{workflow_id}/events",
636+
f"{CONDUCTOR_URL}/api/{creds['organization']}/applications/{_path(application_name)}/workflows/{_path(workflow_id)}/events",
584637
headers={
585638
"Content-Type": "application/json",
586639
"Authorization": f"Bearer {creds['token']}",
587640
},
588641
timeout=30.0,
589642
)
590643
response.raise_for_status()
591-
data = response.json()
592-
result: list[dict[str, Any]] = data.get("events", [])
644+
result: list[dict[str, Any]] = response.json()
593645
return result
594646

595647

@@ -601,16 +653,15 @@ async def get_workflow_notifications(
601653
creds = _get_credentials()
602654
async with httpx.AsyncClient() as client:
603655
response = await client.get(
604-
f"{CONDUCTOR_URL}/api/{creds['organization']}/applications/{application_name}/workflows/{workflow_id}/notifications",
656+
f"{CONDUCTOR_URL}/api/{creds['organization']}/applications/{_path(application_name)}/workflows/{_path(workflow_id)}/notifications",
605657
headers={
606658
"Content-Type": "application/json",
607659
"Authorization": f"Bearer {creds['token']}",
608660
},
609661
timeout=30.0,
610662
)
611663
response.raise_for_status()
612-
data = response.json()
613-
result: list[dict[str, Any]] = data.get("notifications", [])
664+
result: list[dict[str, Any]] = response.json()
614665
return result
615666

616667

@@ -619,7 +670,6 @@ async def list_schedules(
619670
status: str | list[str] | None = None,
620671
workflow_name: str | list[str] | None = None,
621672
schedule_name_prefix: str | list[str] | None = None,
622-
load_context: bool | None = None,
623673
) -> list[dict[str, Any]]:
624674
"""List schedules for an application."""
625675
creds = _get_credentials()
@@ -630,12 +680,10 @@ async def list_schedules(
630680
body["workflow_name"] = workflow_name
631681
if schedule_name_prefix is not None:
632682
body["schedule_name_prefix"] = schedule_name_prefix
633-
if load_context is not None:
634-
body["load_context"] = load_context
635683

636684
async with httpx.AsyncClient() as client:
637685
response = await client.post(
638-
f"{CONDUCTOR_URL}/api/{creds['organization']}/applications/{application_name}/schedules/list",
686+
f"{CONDUCTOR_URL}/api/{creds['organization']}/applications/{_path(application_name)}/schedules/list",
639687
json=body,
640688
headers={
641689
"Content-Type": "application/json",
@@ -644,8 +692,7 @@ async def list_schedules(
644692
timeout=30.0,
645693
)
646694
response.raise_for_status()
647-
data = response.json()
648-
result: list[dict[str, Any]] = data.get("output", [])
695+
result: list[dict[str, Any]] = response.json()
649696
return result
650697

651698

@@ -657,16 +704,15 @@ async def get_schedule(
657704
creds = _get_credentials()
658705
async with httpx.AsyncClient() as client:
659706
response = await client.get(
660-
f"{CONDUCTOR_URL}/api/{creds['organization']}/applications/{application_name}/schedules/{schedule_name}",
707+
f"{CONDUCTOR_URL}/api/{creds['organization']}/applications/{_path(application_name)}/schedules/{_path(schedule_name)}",
661708
headers={
662709
"Content-Type": "application/json",
663710
"Authorization": f"Bearer {creds['token']}",
664711
},
665712
timeout=30.0,
666713
)
667714
response.raise_for_status()
668-
data = response.json()
669-
result: dict[str, Any] = data.get("output", {})
715+
result: dict[str, Any] = response.json()
670716
return result
671717

672718

@@ -678,7 +724,7 @@ async def pause_schedule(
678724
creds = _get_credentials()
679725
async with httpx.AsyncClient() as client:
680726
response = await client.post(
681-
f"{CONDUCTOR_URL}/api/{creds['organization']}/applications/{application_name}/schedules/{schedule_name}/pause",
727+
f"{CONDUCTOR_URL}/api/{creds['organization']}/applications/{_path(application_name)}/schedules/{_path(schedule_name)}/pause",
682728
headers={
683729
"Content-Type": "application/json",
684730
"Authorization": f"Bearer {creds['token']}",
@@ -696,7 +742,7 @@ async def resume_schedule(
696742
creds = _get_credentials()
697743
async with httpx.AsyncClient() as client:
698744
response = await client.post(
699-
f"{CONDUCTOR_URL}/api/{creds['organization']}/applications/{application_name}/schedules/{schedule_name}/resume",
745+
f"{CONDUCTOR_URL}/api/{creds['organization']}/applications/{_path(application_name)}/schedules/{_path(schedule_name)}/resume",
700746
headers={
701747
"Content-Type": "application/json",
702748
"Authorization": f"Bearer {creds['token']}",
@@ -714,7 +760,7 @@ async def trigger_schedule(
714760
creds = _get_credentials()
715761
async with httpx.AsyncClient() as client:
716762
response = await client.post(
717-
f"{CONDUCTOR_URL}/api/{creds['organization']}/applications/{application_name}/schedules/{schedule_name}/trigger",
763+
f"{CONDUCTOR_URL}/api/{creds['organization']}/applications/{_path(application_name)}/schedules/{_path(schedule_name)}/trigger",
718764
headers={
719765
"Content-Type": "application/json",
720766
"Authorization": f"Bearer {creds['token']}",
@@ -733,16 +779,15 @@ async def list_application_versions(
733779
creds = _get_credentials()
734780
async with httpx.AsyncClient() as client:
735781
response = await client.get(
736-
f"{CONDUCTOR_URL}/api/{creds['organization']}/applications/{application_name}/versions",
782+
f"{CONDUCTOR_URL}/api/{creds['organization']}/applications/{_path(application_name)}/versions",
737783
headers={
738784
"Content-Type": "application/json",
739785
"Authorization": f"Bearer {creds['token']}",
740786
},
741787
timeout=30.0,
742788
)
743789
response.raise_for_status()
744-
data = response.json()
745-
result: list[dict[str, Any]] = data.get("output", [])
790+
result: list[dict[str, Any]] = response.json()
746791
return result
747792

748793

@@ -754,7 +799,7 @@ async def set_latest_application_version(
754799
creds = _get_credentials()
755800
async with httpx.AsyncClient() as client:
756801
response = await client.post(
757-
f"{CONDUCTOR_URL}/api/{creds['organization']}/applications/{application_name}/versions/set-latest",
802+
f"{CONDUCTOR_URL}/api/{creds['organization']}/applications/{_path(application_name)}/versions/set-latest",
758803
json={"version_name": version_name},
759804
headers={
760805
"Content-Type": "application/json",

0 commit comments

Comments
 (0)