Repository: PrefectHQ/prefect

Issue:
Ability to pass **kwargs to on_completion and on_failure hooks
### 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

New features in version 2.8.3 allow for hooks on task and flow completion, and task / flow failure.

@jakekaplan shared the link below to the source for the task implementation:
https://github.com/PrefectHQ/prefect/blob/main/src/prefect/engine.py#L1998-L2028

This unfortunately does not allow accessing any of the task / flow context parameters, or allow passing of any variables outside of the defined `task`, `task_run`, `state`. 

A contrived example of the relevance here - 
I run a flow, and upon completion, or failure, for change management purposes I would like to update a Jira or ServiceNow ticket. In order to do so, I need the flow name, a callback url, a secret, and a piece of flow identifying information that was an input parameter to start the flow.

I call this function like:
`def update_api(callback_url, step_id, execution_id, secret)`.

This would be the same call for success or failure, with the only difference being the result / state update. 

Using `on_completion` or `on_failure` however, do not allow me to pass these variables, or access them. E.g.
`@flow(on_completion=[update_api])` - I have all the variables accessible in flow_context , and within the flow, but I don't have a way to pass them into the hook.



### Describe the proposed behavior

Ideally, a modification to allow **kwargs to be passed into the hooks which can be unpacked with whatever inputs are necessary for the hook.

### Example Use

Using the given example above (with callback_url etc.)

```
def update_api(callback_url, step_id, execution_id, secret):
...
...

@flow(on_completion=[update_api(callback_url, step_id, execution_id, secret)]
```

Alternatively, a dictionary example:
```
def update_api(params):
...
...

@flow(on_completion=[update_api(params)]
params['callback_url'] = 'https://my-callback.com'
params['step_id'] = 'aab1-eeac'
params['execution_id'] = '867-5309'
params['secret'] = 'abc123'
```


### Additional context

_No response_


Discussion:
This issue is stale because it has been open 30 days with no activity. To keep this issue open remove stale label or comment.
This issue was closed because it has been stale for 14 days with no activity. If this issue is important or you have more to add feel free to re-open it.
@billpalombi I'm wondering if we should support something like this?
I've seen two requests for this enhancement. I like it. As long as it doesn't add a lot of complexity, I fully support it.

[Related discourse thread](https://discourse.prefect.io/t/passing-and-retrieving-data-in-the-on-failure-hook/2541)

Resolving patch:
diff --git a/src/prefect/engine.py b/src/prefect/engine.py
--- a/src/prefect/engine.py
+++ b/src/prefect/engine.py
@@ -90,7 +90,18 @@
 import time
 from contextlib import AsyncExitStack, asynccontextmanager
 from functools import partial
-from typing import Any, Awaitable, Dict, Iterable, List, Optional, Set, TypeVar, Union
+from typing import (
+    Any,
+    Awaitable,
+    Callable,
+    Dict,
+    Iterable,
+    List,
+    Optional,
+    Set,
+    TypeVar,
+    Union,
+)
 from uuid import UUID, uuid4
 
 import anyio
@@ -2324,6 +2335,16 @@ def _resolve_custom_task_run_name(task: Task, parameters: Dict[str, Any]) -> str
     return task_run_name
 
 
+def _get_hook_name(hook: Callable) -> str:
+    return (
+        hook.__name__
+        if hasattr(hook, "__name__")
+        else (
+            hook.func.__name__ if isinstance(hook, partial) else hook.__class__.__name__
+        )
+    )
+
+
 async def _run_task_hooks(task: Task, task_run: TaskRun, state: State) -> None:
     """Run the on_failure and on_completion hooks for a task, making sure to
     catch and log any errors that occur.
@@ -2337,9 +2358,10 @@ async def _run_task_hooks(task: Task, task_run: TaskRun, state: State) -> None:
     if hooks:
         logger = task_run_logger(task_run)
         for hook in hooks:
+            hook_name = _get_hook_name(hook)
             try:
                 logger.info(
-                    f"Running hook {hook.__name__!r} in response to entering state"
+                    f"Running hook {hook_name!r} in response to entering state"
                     f" {state.name!r}"
                 )
                 if is_async_fn(hook):
@@ -2350,11 +2372,11 @@ async def _run_task_hooks(task: Task, task_run: TaskRun, state: State) -> None:
                     )
             except Exception:
                 logger.error(
-                    f"An error was encountered while running hook {hook.__name__!r}",
+                    f"An error was encountered while running hook {hook_name!r}",
                     exc_info=True,
                 )
             else:
-                logger.info(f"Hook {hook.__name__!r} finished running successfully")
+                logger.info(f"Hook {hook_name!r} finished running successfully")
 
 
 async def _run_flow_hooks(flow: Flow, flow_run: FlowRun, state: State) -> None:
@@ -2374,9 +2396,10 @@ async def _run_flow_hooks(flow: Flow, flow_run: FlowRun, state: State) -> None:
     if hooks:
         logger = flow_run_logger(flow_run)
         for hook in hooks:
+            hook_name = _get_hook_name(hook)
             try:
                 logger.info(
-                    f"Running hook {hook.__name__!r} in response to entering state"
+                    f"Running hook {hook_name!r} in response to entering state"
                     f" {state.name!r}"
                 )
                 if is_async_fn(hook):
@@ -2387,11 +2410,11 @@ async def _run_flow_hooks(flow: Flow, flow_run: FlowRun, state: State) -> None:
                     )
             except Exception:
                 logger.error(
-                    f"An error was encountered while running hook {hook.__name__!r}",
+                    f"An error was encountered while running hook {hook_name!r}",
                     exc_info=True,
                 )
             else:
-                logger.info(f"Hook {hook.__name__!r} finished running successfully")
+                logger.info(f"Hook {hook_name!r} finished running successfully")
 
 
 async def check_api_reachable(client: PrefectClient, fail_message: str):
diff --git a/src/prefect/utilities/callables.py b/src/prefect/utilities/callables.py
--- a/src/prefect/utilities/callables.py
+++ b/src/prefect/utilities/callables.py
@@ -268,10 +268,10 @@ def parameter_schema(fn: Callable) -> ParameterSchema:
         - additional constraints (like possible enum values)
 
     Args:
-        fn (function): The function whose arguments will be serialized
+        fn (Callable): The function whose arguments will be serialized
 
     Returns:
-        dict: the argument schema
+        ParameterSchema: the argument schema
     """
     signature = inspect.signature(fn)
     model_fields = {}


Test patch:
diff --git a/tests/test_flows.py b/tests/test_flows.py
--- a/tests/test_flows.py
+++ b/tests/test_flows.py
@@ -6,6 +6,7 @@
 import signal
 import sys
 import time
+from functools import partial
 from itertools import combinations
 from pathlib import Path
 from textwrap import dedent
@@ -2417,6 +2418,38 @@ async def my_hook(flow, flow_run, state):
     return my_hook
 
 
+class TestFlowHooksWithKwargs:
+    def test_hook_with_extra_default_arg(self):
+        data = {}
+
+        def hook(flow, flow_run, state, foo=42):
+            data.update(name=hook.__name__, state=state, foo=foo)
+
+        @flow(on_completion=[hook])
+        def foo_flow():
+            pass
+
+        state = foo_flow(return_state=True)
+
+        assert data == dict(name="hook", state=state, foo=42)
+
+    def test_hook_with_bound_kwargs(self):
+        data = {}
+
+        def hook(flow, flow_run, state, **kwargs):
+            data.update(name=hook.__name__, state=state, kwargs=kwargs)
+
+        hook_with_kwargs = partial(hook, foo=42)
+
+        @flow(on_completion=[hook_with_kwargs])
+        def foo_flow():
+            pass
+
+        state = foo_flow(return_state=True)
+
+        assert data == dict(name="hook", state=state, kwargs={"foo": 42})
+
+
 class TestFlowHooksOnCompletion:
     def test_noniterable_hook_raises(self):
         def completion_hook():
diff --git a/tests/test_tasks.py b/tests/test_tasks.py
--- a/tests/test_tasks.py
+++ b/tests/test_tasks.py
@@ -3,6 +3,7 @@
 import time
 import warnings
 from asyncio import Event, sleep
+from functools import partial
 from typing import Any, Dict, List
 from unittest.mock import MagicMock, call
 from uuid import UUID
@@ -3423,6 +3424,46 @@ async def my_hook(task, task_run, state):
     return my_hook
 
 
+class TestTaskHooksWithKwargs:
+    def test_hook_with_extra_default_arg(self):
+        data = {}
+
+        def hook(task, task_run, state, foo=42):
+            data.update(name=hook.__name__, state=state, foo=foo)
+
+        @task(on_completion=[hook])
+        def foo_task():
+            pass
+
+        @flow
+        def foo_flow():
+            return foo_task(return_state=True)
+
+        state = foo_flow()
+
+        assert data == dict(name="hook", state=state, foo=42)
+
+    def test_hook_with_bound_kwargs(self):
+        data = {}
+
+        def hook(task, task_run, state, **kwargs):
+            data.update(name=hook.__name__, state=state, kwargs=kwargs)
+
+        hook_with_kwargs = partial(hook, foo=42)
+
+        @task(on_completion=[hook_with_kwargs])
+        def foo_task():
+            pass
+
+        @flow
+        def foo_flow():
+            return foo_task(return_state=True)
+
+        state = foo_flow()
+
+        assert data == dict(name="hook", state=state, kwargs={"foo": 42})
+
+
 class TestTaskHooksOnCompletion:
     def test_noniterable_hook_raises(self):
         def completion_hook():