Skip to content

Commit 062f160

Browse files
committed
added the overloaded logic and integration test
1 parent 8cb5fc9 commit 062f160

File tree

2 files changed

+18
-3
lines changed

2 files changed

+18
-3
lines changed

pyatlan/client/workflow.py

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
ReRunRequest,
3030
ScheduleQueriesSearchRequest,
3131
Workflow,
32+
WorkflowMetadata,
3233
WorkflowResponse,
3334
WorkflowRunResponse,
3435
WorkflowSchedule,
@@ -407,20 +408,29 @@ def update_owner(self, workflow_name: str, username: str) -> WorkflowResponse:
407408

408409
@validate_arguments(config=dict(arbitrary_types_allowed=True))
409410
def monitor(
410-
self, workflow_response: WorkflowResponse, logger: Optional[Logger] = None
411+
self,
412+
workflow_response: Optional[WorkflowResponse] = None,
413+
workflow_name: Optional[str] = None,
414+
logger: Optional[Logger] = None,
411415
) -> Optional[AtlanWorkflowPhase]:
412416
"""
413417
Monitor the status of the workflow's run.
414418
415419
:param workflow_response: The workflow_response returned from running the workflow
420+
:param workflow_name: name of the workflow to be monitored
416421
:param logger: the logger to log status information
417422
(logging.INFO for summary info. logging:DEBUG for detail info)
418423
:returns: the status at completion or None if the workflow wasn't run
419424
:raises ValidationError: If the provided `workflow_response`, `logger` is invalid
420425
:raises AtlanError: on any API communication issue
421426
"""
422-
if workflow_response.metadata and workflow_response.metadata.name:
423-
name = workflow_response.metadata.name
427+
if workflow_name and not workflow_response:
428+
workflow_response = WorkflowResponse(
429+
metadata=WorkflowMetadata(name=workflow_name)
430+
)
431+
432+
if workflow_response.metadata and workflow_response.metadata.name: # type: ignore
433+
name = workflow_response.metadata.name # type: ignore
424434
status: Optional[AtlanWorkflowPhase] = None
425435
while status not in {
426436
AtlanWorkflowPhase.SUCCESS,

tests/integration/test_workflow_client.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -157,6 +157,11 @@ def test_workflow_get_runs_and_stop(client: AtlanClient, workflow: WorkflowRespo
157157
workflow_status = client.workflow.monitor(workflow_response=workflow)
158158
assert workflow_status == AtlanWorkflowPhase.FAILED
159159

160+
# Test workflow monitoring by providing workflow name directly
161+
workflow_name = workflow.metadata.name
162+
workflow_status = client.workflow.monitor(workflow_name=workflow_name)
163+
assert workflow_status == AtlanWorkflowPhase.FAILED
164+
160165
# Test find run by id
161166
workflow_run = client.workflow.find_run_by_id(id=run.id)
162167
assert (

0 commit comments

Comments
 (0)