Repository: PrefectHQ/prefect

Issue:
Allow idempotency keys with Server
In https://github.com/PrefectHQ/prefect/pull/3005/files#diff-7836972d3e8bafd3361be68b33c62fb3R126 we introduced a trap to ensure that Server doesn't receive idempotency keys. Once a new version of Server is released that does support them, we should remove this trap.


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
@@ -1,8 +1,7 @@
 import time
 from typing import Any
-from urllib.parse import urlparse
 
-from prefect import config, context, Task
+from prefect import context, Task
 from prefect.client import Client
 from prefect.engine.signals import signal_from_state
 from prefect.utilities.graphql import EnumValue, with_args
@@ -88,22 +87,19 @@ def run(
             ```
 
         """
-        is_hosted_backend = "prefect.io" in urlparse(config.cloud.api).netloc
 
         # verify that flow and project names were passed where necessary
         if flow_name is None:
             raise ValueError("Must provide a flow name.")
-        if project_name is None and is_hosted_backend:
+        if project_name is None:
             raise ValueError("Must provide a project name.")
 
         where_clause = {
             "name": {"_eq": flow_name},
             "archived": {"_eq": False},
+            "project": {"name": {"_eq": project_name}},
         }
 
-        if project_name:
-            where_clause["project"] = {"name": {"_eq": project_name}}
-
         # find the flow ID to schedule
         query = {
             "query": {
@@ -128,21 +124,20 @@ 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"):
+        idem_key = None
+        if 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=idem_key,
+            idempotency_key=idem_key or idempotency_key,
             context=new_flow_context,
         )
 


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
@@ -122,6 +122,7 @@ def test_initialization(self, server_api):
         # verify that the task is initialized as expected
         task = FlowRunTask(
             name="My Flow Run Task",
+            project_name="Demo",
             checkpoint=False,
             flow_name="Test Flow",
             new_flow_context={"foo": "bar"},
@@ -135,7 +136,9 @@ def test_initialization(self, server_api):
 
     def test_flow_run_task(self, client, server_api):
         # verify that create_flow_run was called
-        task = FlowRunTask(flow_name="Test Flow", parameters={"test": "ing"},)
+        task = FlowRunTask(
+            flow_name="Test Flow", project_name="Demo", parameters={"test": "ing"},
+        )
         # verify that run returns the new flow run ID
         assert task.run() == "xyz890"
         # verify the GraphQL query was called with the correct arguments
@@ -158,7 +161,7 @@ def test_flow_run_task_without_flow_name(self, server_api):
 
     def test_flow_run_task_with_no_matching_flow(self, client, server_api):
         # verify a ValueError is raised if the client returns no flows
-        task = FlowRunTask(flow_name="Test Flow")
+        task = FlowRunTask(flow_name="Test Flow", project_name="Demo")
         client.graphql = MagicMock(return_value=MagicMock(data=MagicMock(flow=[])))
         with pytest.raises(ValueError, match="Flow 'Test Flow' not found."):
             task.run()
@@ -179,6 +182,6 @@ def test_flow_run_task_with_flow_run_id(self, client, server_api):
         client.create_flow_run.assert_called_once_with(
             flow_id="abc123",
             parameters={"test": "ing"},
-            idempotency_key=None,
+            idempotency_key="test-id",
             context=None,
         )