Repository: PrefectHQ/prefect

Issue:
Missing key 'idempotency_key' in Server-mode GraphQL schema (found in v0.12.4)
## Description
When starting a registered flow from another flow using `prefect.tasks.prefect.flow_run.FlowRunTask` an error occurred: 
`requests.exceptions.HTTPError: 400 Client Error: Bad Request for url: http://localhost:4200/graphql`.  At the same time the Apollo log states

```
2020-07-20T13:21:31.790Z {"message":"Variable \"$input\" got invalid value { flow_id: \"80de1b39-fed7-47ab-a5a5-fc22ac6f87d2\", idempotency_key: \"6a83567d-7e10-4b3a-815c-a29fe3ecab02\" }; Field \"idempotency_key\" is not defined by type create_flow_run_input.","locations":[{"line":1,"column":10}],"extensions":{"code":"INTERNAL_SERVER_ERROR"}}
```

The problem exists in [schema.graphql line 114](https://github.com/PrefectHQ/prefect/blob/0.12.4/server/src/prefect_server/graphql/schema/schema.graphql#L114) . There is *no* `idempotency_key` key in the `create_flow_input` structure/input. The [Client Implementation](https://github.com/PrefectHQ/prefect/blob/0.12.4/src/prefect/client/client.py#L869) is not forcing the usage of `idempotency.py` because its optional. But the `FlowRunTask` is setting a default in [flow_run.py:122](https://github.com/PrefectHQ/prefect/blob/0.12.4/src/prefect/tasks/prefect/flow_run.py#L122). 

You can also find the key missing by using the (standalone/server-mode) UI using the interactive API under _DOCS_ -> _Mutations_ -> `create_flow` ->  _Arguments_ -> `input::CREATE_FLOW_INPUT`.



When this key is present, the start of a flow run must fail because of the unexpected key. While hunting down this bug I used my own `FlowRunTask` wich was an exact copy of Prefect's FlowRunTask, *except line 122 removed*. As a result I could initiate a flow run like expected. 

Maybe its sufficient to add `idempotency_key` to the _schema.graphql_ file?

## Expected Behavior
The registered flow should start. Usage of `idempotency_key` should be possible.

## Reproduction
Minimal example to reproduce:

### flow1.py
```python
#!/usr/bin/env python
# coding: utf-8

import prefect
from prefect import task, Flow
from prefect.environments.storage.local import Local


@task
def task1():
    prefect.context.get("logger").info("hello, task1")


with Flow("flow1", storage=Local(directory="/flows/.prefect/flows")) as flow:
    task1()

if __name__ == "__main__":
    flow.register()
```

### flow2.py

```python
#!/usr/bin/env python
# coding: utf-8

import prefect
from prefect import task, Flow
from prefect.tasks.prefect import FlowRunTask
from prefect.environments.storage.local import Local


@task
def task2():
    prefect.context.get("logger").info("hello, task2")


with Flow("flow2", storage=Local(directory="/flows/.prefect/flows")) as flow:
    # the task will generate an idempotency_key if not present (or None!)
    flow1_run = FlowRunTask(name="flow1-runner", flow_name="flow1")
    task2.set_downstream(flow1_run)

if __name__ == "__main__":
    flow.register()
```

### Test run (fail)
In the Agent's container in the path where I store my flows:
```
prefect  run server -n flow2 -w
```

## Environment
Server mode; Prefect Server and all other components are at v0.12.4.



Discussion:
Thanks for raising @realknorke that line in the FlowRunTask which sets the "default" should be optional/toggle 🙂 
FYI: `prefect run server …` CLI command does not know `idempotency_key`either.

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,17 +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
+        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.
@@ -58,6 +68,10 @@ def run(
             - parameters (dict, optional): the parameters to pass to the flow run being
                 scheduled; if not provided, this method will use the parameters provided at
                 initialization
+            - idempotency_key (str, optional): an optional idempotency key for scheduling the
+                flow run; if provided, ensures that only one 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
@@ -114,12 +128,22 @@ def run(
         # grab the ID for the most recent version
         flow_id = flow[0].id
 
+        if is_hosted_backend and context.get("flow_run_id"):
+            map_index = context.get("map_index")
+            default = context.get("flow_run_id") + (
+                f"-{map_index}" if map_index else ""
+            )
+            idem_key = idempotency_key or default
+        else:
+            idem_key = None
+
         # 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=context.get("flow_run_id"),
+            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,26 @@ 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):
+        # verify that create_flow_run was called
+        task = FlowRunTask(project_name="Test Project", flow_name="Test Flow",)
+
+        # verify that run returns the new flow run ID
+        with prefect.context(flow_run_id="test-id", map_index=4):
+            assert task.run() == "xyz890"
+
+        # 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,
+            context=None,
         )
 
     def test_flow_run_task_without_flow_name(self, cloud_api):
@@ -100,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):
@@ -118,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):
@@ -133,3 +162,23 @@ def test_flow_run_task_with_no_matching_flow(self, client, server_api):
         client.graphql = MagicMock(return_value=MagicMock(data=MagicMock(flow=[])))
         with pytest.raises(ValueError, match="Flow 'Test Flow' not found."):
             task.run()
+
+    def test_flow_run_task_with_flow_run_id(self, client, server_api):
+        # verify that create_flow_run was called
+        task = FlowRunTask(
+            project_name="Test Project",
+            flow_name="Test Flow",
+            parameters={"test": "ing"},
+        )
+
+        # verify that run returns the new flow run ID
+        with prefect.context(flow_run_id="test-id"):
+            assert task.run() == "xyz890"
+
+        # 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,
+            context=None,
+        )