Repository: PrefectHQ/prefect
Pull request: wait_for_flow_run PrefectClient function
<!-- 
Thanks for opening a pull request to Prefect! We've got a few requests to help us review contributions:

- Make sure that your title neatly summarizes the proposed changes.
- Provide a short overview of the change and the value it adds.
- Share an example to help us understand the change in user experience.
- Confirm that you've done common tasks so we can give a timely review.
- Review our contribution guidelines: https://docs.prefect.io/latest/contributing/overview/

Happy engineering!
-->

As described in issue #11177 , I would like to see this feature in production. The ability to use the PrefectClient to automatically wait for a flow run to be done before proceeding, using just the prefect client.

closes #11177 

### Example
Imagine you have an API endpoint that fires off a flow_run and then creates a background task that waits for the flow run to finish so it can use the results.
```python
async def my_endpoint(background_tasks, flow_input):
    async with get_client() as c:
        new_flow_run = await c.create_flow_run_from_deployment('<some UUID>', flow_input)
    background_tasks.append(wait_and_insert, new_flow_run)

async def wait_and_insert(flow_run):
    async with get_client() as c:
        completed_flow_run = await c.wait_for_flow_run(flow_run.id)
    if completed_flow_run.state.type = StateType.COMPLETED:
        result = completed_flow_run.state.result(fetch=True)
        insert_results(result)
    else:
        raise Failure
````
Suppose here that `insert_results` is some database write function.

### Checklist
<!-- These boxes may be checked after opening the pull request. -->

- [x] This pull request references any related issue by including "closes `<link to issue>`"
	- If no issue exists and your change is not a small fix, please [create an issue](https://github.com/PrefectHQ/prefect/issues/new/choose) first.
- [x] This pull request includes tests or only affects documentation.
- [x] This pull request includes a label categorizing the change e.g. `fix`, `feature`, `enhancement`, `docs`.
  <!-- If you do not have permission to add a label, a maintainer will add one for you -->

For documentation changes:

- [ ] This pull request includes redirect settings in `netlify.toml` for files that are removed or renamed


Problem statement:
client.wait_for_flow_run function proposal
### First check

- [X] I added a descriptive title to this issue.
- [X] I used the GitHub search to find a similar request and didn't find it.
- [X] I searched the Prefect documentation for this feature.

### Prefect Version

2.x

### Describe the current behavior

- I can use `prefect.deployments.run_deployment` to create a flow run and wait for it to finish, and get the results.
- I cannot (afaik) get a flow run from `client.create_flow_run_from_deployment` and then use some another function to later wait for that deployment and get the results

### Describe the proposed behavior

A function that lets me wait for a flow_run to finish and get the results when it is done, just like `prefect.deployments.run_deployment` but with an existing flow_run ID that I instantiated in a separate call.

### Example Use

A REST API where someone tells my API to kick off a prefect flow, i kick it off, and hand them back an ID. My endpoint then starts a background task to wait for the results of that flow to come in so it can write them to it's own database.

### Additional context

I am happy to put this feature in myself. Opening this issue first because IDK if this is a duplicate request or not, or maybe the feature exists and I couldn't find it. I looked through the docs and github issues and I haven't seen this. Thanks

Hints from issue discussion:
Thanks for opening this issue @Fizzizist!

As you said, there is presently no such method. Do I understand correctly that you are looking for something like the `wait()` method for task futures, but for flows?


@billpalombi I am thinking something that can be done with just the prefect client and a flow run ID. Here is how I've done it in my repo (taking most inspiration from `prefect.deployments.run_deployment`:
```python
async def wait_for_flow_run(
    flow_run: FlowRun, timeout: int = 10800, poll_interval: int = 5
) -> None:
    """Waits for the prefect flow to finish and return the results. If anything fails, an error is logged.

    Parameters
    ----------
    flow_run : FlowRun
        The flow run.
    timeout : int, optional
        The timeout in seconds, by default 10800 (3 hours).
    poll_interval : int, optional
        The poll interval in seconds, by default 5.
    """
    flow_running = True
    async with get_client() as client:
        with anyio.move_on_after(timeout):
            while flow_running:
                flow_run = await client.read_flow_run(flow_run.id)
                flow_state = flow_run.state
                if flow_state and flow_state.is_final():
                    flow_running = False
                await anyio.sleep(poll_interval)

    if flow_state is None or flow_state.type != StateType.COMPLETED:
        slogger.error(
            "flow run did not succeed.",
            flow_run_id=flow_run.id,
            state=flow_state.type if flow_state is not None else None,
            output=flow_run.json(),
        )
        return

    return await flow_state.result().get()
```
@Fizzizist this looks like a novel request and your example looks good! We'd be happy to review a contribution!
awesome. I'll put a PR together

Commit message: Merge branch 'main' into main

Patch to review:
diff --git a/src/prefect/client/orchestration.py b/src/prefect/client/orchestration.py
index 461f7b2531ff..d6d327b6773c 100644
--- a/src/prefect/client/orchestration.py
+++ b/src/prefect/client/orchestration.py
@@ -680,7 +680,7 @@ async def delete_flow_run(
             httpx.RequestError: If requests fails
         """
         try:
-            await self._client.delete(f"/flow_runs/{flow_run_id}"),
+            await self._client.delete(f"/flow_runs/{flow_run_id}")
         except httpx.HTTPStatusError as e:
             if e.response.status_code == status.HTTP_404_NOT_FOUND:
                 raise prefect.exceptions.ObjectNotFound(http_exc=e) from e
diff --git a/src/prefect/exceptions.py b/src/prefect/exceptions.py
index 983dc7b2027d..d957c437a225 100644
--- a/src/prefect/exceptions.py
+++ b/src/prefect/exceptions.py
@@ -401,3 +401,7 @@ class NotPausedError(PrefectException):
 
 class FlowPauseTimeout(PrefectException):
     """Raised when a flow pause times out"""
+
+
+class FlowWaitTimeout(PrefectException):
+    """Raised when a flow run takes longer than a given timeout"""
diff --git a/src/prefect/flow_runs.py b/src/prefect/flow_runs.py
new file mode 100644
index 000000000000..6195fcc5a76d
--- /dev/null
+++ b/src/prefect/flow_runs.py
@@ -0,0 +1,37 @@
+from uuid import UUID
+
+import anyio
+
+from prefect.client.orchestration import PrefectClient
+from prefect.client.schemas import FlowRun
+from prefect.exceptions import FlowWaitTimeout
+
+
+async def wait_for_flow_run(
+    client: PrefectClient,
+    flow_run_id: UUID,
+    timeout: int = 10800,
+    poll_interval: int = 5,
+) -> FlowRun:
+    """Waits for the prefect flow run to finish and returns the FlowRun
+
+    Ars:
+        flow_run_id: The flow run ID for the flow run to wait for.
+        timeout: The timeout in seconds, by default 10800 (3 hours).
+        poll_interval: The poll interval in seconds, by default 5.
+    Returns:
+        FlowRun: The finished flow run.
+    Raises:
+        prefect.exceptions.FlowWaitTimeout: If flow run goes over the timeout.
+    """
+    with anyio.move_on_after(timeout):
+        while True:
+            flow_run = await client.read_flow_run(flow_run_id)
+            flow_state = flow_run.state
+            if flow_state and flow_state.is_final():
+                return flow_run
+            await anyio.sleep(poll_interval)
+
+    raise FlowWaitTimeout(
+        f"FlowRun with ID {flow_run_id} is taking longer than {timeout} seconds"
+    )
diff --git a/tests/test_flow_runs.py b/tests/test_flow_runs.py
new file mode 100644
index 000000000000..43cbcd6bfd23
--- /dev/null
+++ b/tests/test_flow_runs.py
@@ -0,0 +1,40 @@
+import time
+
+import pytest
+
+import prefect.client.schemas as client_schemas
+from prefect import flow
+from prefect.exceptions import FlowWaitTimeout
+from prefect.flow_runs import wait_for_flow_run
+from prefect.states import Completed
+
+
+async def test_create_then_wait_for_flow_run(prefect_client):
+    @flow
+    def foo():
+        pass
+
+    flow_run = await prefect_client.create_flow_run(
+        foo, name="petes-flow-run", state=Completed()
+    )
+    assert isinstance(flow_run, client_schemas.FlowRun)
+
+    lookup = await wait_for_flow_run(prefect_client, flow_run.id, poll_interval=0)
+    # Estimates will not be equal since time has passed
+    assert lookup == flow_run
+    assert flow_run.state.is_final()
+
+
+async def test_create_then_wait_timeout(prefect_client):
+    @flow
+    def foo():
+        time.sleep(9999)
+
+    flow_run = await prefect_client.create_flow_run(
+        foo,
+        name="petes-flow-run",
+    )
+    assert isinstance(flow_run, client_schemas.FlowRun)
+
+    with pytest.raises(FlowWaitTimeout):
+        await wait_for_flow_run(prefect_client, flow_run.id, timeout=0)


Code review comments:
diff_hunk: @@ -0,0 +1,37 @@
+from uuid import UUID
+
+import anyio
+
+from prefect.client.orchestration import PrefectClient
+from prefect.client.schemas import FlowRun
+from prefect.exceptions import FlowWaitTimeout
+
+
+async def wait_for_flow_run(
+    client: PrefectClient,
+    flow_run_id: UUID,
+    timeout: int = 10800,
+    poll_interval: int = 5,
+) -> FlowRun:
+    """Waits for the prefect flow run to finish and returns the FlowRun
+
+    Ars:
+        flow_run_id: The flow run ID for the flow run to wait for.
+        timeout: The timeout in seconds, by default 10800 (3 hours).
+        poll_interval: The poll interval in seconds, by default 5.
+    Returns:
+        FlowRun: The finished flow run.
+    Raises:
+        prefect.exceptions.FlowWaitTimeout: If flow run goes over the timeout.
+    """
+    with anyio.move_on_after(timeout):
+        while True:
+            flow_run = await client.read_flow_run(flow_run_id)
+            flow_state = flow_run.state
+            if flow_state and flow_state.is_final():
+                return flow_run
+            await anyio.sleep(poll_interval)
+
+    raise FlowWaitTimeout(
+        f"FlowRun with ID {flow_run_id} is taking longer than {timeout} seconds"
path: src/prefect/flow_runs.py
text: @user1:
```suggestion
        f"Flow run with ID {flow_run_id} exceeded watch timeout of {timeout} seconds"
```

@author:
addressed

diff_hunk: @@ -0,0 +1,37 @@
+from uuid import UUID
+
+import anyio
+
+from prefect.client.orchestration import PrefectClient
+from prefect.client.schemas import FlowRun
+from prefect.exceptions import FlowWaitTimeout
+
+
+async def wait_for_flow_run(
+    client: PrefectClient,
+    flow_run_id: UUID,
+    timeout: int = 10800,
path: src/prefect/flow_runs.py
text: @user1:
It'd be nice to type this a `Optional` since passing `None` will disable the timeout.
```suggestion
    timeout: Optional[int] = 10800,
```

@author:
addressed

diff_hunk: @@ -401,3 +401,7 @@ class NotPausedError(PrefectException):
 
 class FlowPauseTimeout(PrefectException):
     """Raised when a flow pause times out"""
+
+
+class FlowWaitTimeout(PrefectException):
path: src/prefect/exceptions.py
text: @user1:
I think it'd be more accurate to name this `FlowRunWaitTimeout`.
```suggestion
class FlowRunWaitTimeout(PrefectException):
```

@author:
addressed