Repository: PrefectHQ/prefect

Issue:
Expose context arg in FlowTaskRun task
## Opened from the [Prefect Public Slack Community](https://join.slack.com/t/prefect-public/shared_invite/enQtNzE5OTU3OTQwNzc1LTQ5M2FkZmQzZjI0ODg1ZTBmOTc0ZjVjYWFjMWExZDAyYzBmYjVmMTE1NTQ1Y2IxZTllOTc4MmI3NzYxMDlhYWU)

**asandeep.me**: Hi all,
Is there a way to pass-on current context to next flow run triggered using `FlowRunTask` task?
Basically, I am sending some context parameter while creating a flow run using python client's `create_flow_run` method. The flow upon completion triggers other downstream flows.
Now, I want the context that I originally sent, available to downstream flows but I haven't had any success so far.

**laura**: Hi <@U013572S7RU>, looks like no since `FlowRunTask` doesn’t accept `context` right now (would eventually expect it to get passed along here <https://github.com/PrefectHQ/prefect/blob/master/src/prefect/tasks/prefect/flow_run.py#L109>), but I think that is a feature we can add pretty quick! I can make an issue from this ticket but first let me check if that fixes your issue. I think in that situation then your Flow A (kicked off with the python client) would have in it a node of `FlowRunTask` to kick off Flow B. You want to send context from Flow A (lets say `prefect.context.foo=bar`) to that Flow B. I just want to clarify that in practice, Flow A would need to have, or be able to re-derive just the added context (just `foo=bar`) since we probably don’t want to send the entire context object from Flow A straight into the `FlowRunTask` to ship to Flow B (as it has a bunch of other prefect sutff in it too). Would that achieve your use case?

**asandeep.me**: Thanks <@UTHEQ4F9R> for your reply.
Yes exactly this is what my use case is and selectively passing the required context to Flow B will work.

**laura**: <@ULVA73B9P> open “Expose context arg in FlowTaskRun task”

Original thread can be found [here](https://prefect-community.slack.com/archives/CL09KU1K7/p1594286544388000?thread_ts=1594286544.388000&cid=CL09KU1K7).




Discussion:

Resolving patch:
diff --git a/src/prefect/tasks/prefect/flow_run.py b/src/prefect/tasks/prefect/flow_run.py
--- a/src/prefect/tasks/prefect/flow_run.py
+++ b/src/prefect/tasks/prefect/flow_run.py
@@ -22,6 +22,7 @@ class FlowRunTask(Task):
             running with Prefect Core's server as the backend, this should not be provided.
         - parameters (dict, optional): the parameters to pass to the flow run being scheduled;
             this value may also be provided at run time
+        - new_flow_context (dict, optional): the optional run context for the new flow run
         - wait (bool, optional): whether to wait the triggered flow run's state; if True, this
             task will wait until the flow run is complete, and then reflect the corresponding
             state as the state of this task.  Defaults to `False`.
@@ -34,23 +35,26 @@ def __init__(
         project_name: str = None,
         parameters: dict = None,
         wait: bool = False,
+        new_flow_context: dict = None,
         **kwargs: Any,
     ):
         self.flow_name = flow_name
         self.project_name = project_name
         self.parameters = parameters
+        self.new_flow_context = new_flow_context
         self.wait = wait
         if flow_name:
             kwargs.setdefault("name", f"Flow {flow_name}")
         super().__init__(**kwargs)
 
-    @defaults_from_attrs("flow_name", "project_name", "parameters")
+    @defaults_from_attrs("flow_name", "project_name", "parameters", "new_flow_context")
     def run(
         self,
         flow_name: str = None,
         project_name: str = None,
         parameters: dict = None,
         idempotency_key: str = None,
+        new_flow_context: dict = None,
     ) -> str:
         """
         Run method for the task; responsible for scheduling the specified flow run.
@@ -67,6 +71,7 @@ def run(
             - idempotency_key (str, optional): an optional idempotency key for scheduling the
                 flow run; if provided, ensures that only run is created if this task is retried
                 or rerun with the same inputs.  If not provided, the current flow run ID will be used.
+            - new_flow_context (dict, optional): the optional run context for the new flow run
 
         Returns:
             - str: the ID of the newly-scheduled flow run
@@ -135,7 +140,10 @@ def run(
         # providing an idempotency key ensures that retries for this task
         # will not create additional flow runs
         flow_run_id = client.create_flow_run(
-            flow_id=flow_id, parameters=parameters, idempotency_key=idem_key,
+            flow_id=flow_id,
+            parameters=parameters,
+            idempotency_key=idem_key,
+            context=new_flow_context,
         )
 
         self.logger.debug(f"Flow Run {flow_run_id} created.")


Test patch:
diff --git a/tests/tasks/prefect/test_flow_run.py b/tests/tasks/prefect/test_flow_run.py
--- a/tests/tasks/prefect/test_flow_run.py
+++ b/tests/tasks/prefect/test_flow_run.py
@@ -29,12 +29,14 @@ def test_initialization(self, cloud_api):
             checkpoint=False,
             project_name="Test Project",
             flow_name="Test Flow",
+            new_flow_context={"foo": "bar"},
             parameters={"test": "ing"},
         )
         assert task.name == "My Flow Run Task"
         assert task.checkpoint is False
         assert task.project_name == "Test Project"
         assert task.flow_name == "Test Flow"
+        assert task.new_flow_context == {"foo": "bar"}
         assert task.parameters == {"test": "ing"}
 
     def test_flow_run_task(self, client, cloud_api):
@@ -53,7 +55,10 @@ def test_flow_run_task(self, client, cloud_api):
 
         # verify create_flow_run was called with the correct arguments
         client.create_flow_run.assert_called_once_with(
-            flow_id="abc123", parameters={"test": "ing"}, idempotency_key=None
+            flow_id="abc123",
+            parameters={"test": "ing"},
+            idempotency_key=None,
+            context=None,
         )
 
     def test_flow_run_task_with_flow_run_id(self, client, cloud_api):
@@ -70,7 +75,10 @@ def test_flow_run_task_with_flow_run_id(self, client, cloud_api):
 
         # verify create_flow_run was called with the correct arguments
         client.create_flow_run.assert_called_once_with(
-            flow_id="abc123", parameters={"test": "ing"}, idempotency_key="test-id"
+            flow_id="abc123",
+            parameters={"test": "ing"},
+            idempotency_key="test-id",
+            context=None,
         )
 
     def test_idempotency_key_uses_map_index_if_present(self, client, cloud_api):
@@ -83,7 +91,10 @@ def test_idempotency_key_uses_map_index_if_present(self, client, cloud_api):
 
         # verify create_flow_run was called with the correct arguments
         client.create_flow_run.assert_called_once_with(
-            flow_id="abc123", idempotency_key="test-id-4", parameters=None
+            flow_id="abc123",
+            idempotency_key="test-id-4",
+            parameters=None,
+            context=None,
         )
 
     def test_flow_run_task_without_flow_name(self, cloud_api):
@@ -113,11 +124,13 @@ def test_initialization(self, server_api):
             name="My Flow Run Task",
             checkpoint=False,
             flow_name="Test Flow",
+            new_flow_context={"foo": "bar"},
             parameters={"test": "ing"},
         )
         assert task.name == "My Flow Run Task"
         assert task.checkpoint is False
         assert task.flow_name == "Test Flow"
+        assert task.new_flow_context == {"foo": "bar"}
         assert task.parameters == {"test": "ing"}
 
     def test_flow_run_task(self, client, server_api):
@@ -131,7 +144,10 @@ def test_flow_run_task(self, client, server_api):
 
         # verify create_flow_run was called with the correct arguments
         client.create_flow_run.assert_called_once_with(
-            flow_id="abc123", parameters={"test": "ing"}, idempotency_key=None
+            flow_id="abc123",
+            parameters={"test": "ing"},
+            idempotency_key=None,
+            context=None,
         )
 
     def test_flow_run_task_without_flow_name(self, server_api):
@@ -161,5 +177,8 @@ def test_flow_run_task_with_flow_run_id(self, client, server_api):
 
         # verify create_flow_run was called with the correct arguments
         client.create_flow_run.assert_called_once_with(
-            flow_id="abc123", parameters={"test": "ing"}, idempotency_key=None
+            flow_id="abc123",
+            parameters={"test": "ing"},
+            idempotency_key=None,
+            context=None,
         )