Repository: PrefectHQ/prefect

Issue:
Expose idempotency_key in orion rest api
### 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

The only way to find by the Orios rest API if a flow_run with a given idempotency_key exists is to create following request:

```
import requests
request.post("xxx/api/flow_runs", json={"idempotency_key":"aabb", "flow_id": "12345678-abcd-abcd-abcd-123456781234"})
```

Nevertheless there are situations where we just want to know if there already was a flow_run with a given idempotency key without triggering new one.

### Describe the proposed behavior

Possible solutions:

1. Add idempotency_key to `https://docs.prefect.io/api/flow_runs/filter` for the `flow_runs`.
2. Add a new get method with possibility to filter by idempotency_key and flow_id.


### Example Use

_No response_

### Additional context

This issue came as discussion from https://github.com/PrefectHQ/prefect/issues/5534


Discussion:
Hi @PiotrSiejda , thanks for the issue! This seems pretty reasonable to me, so I've gone ahead and accepted it. 

We have a pretty big back log of time-sensitive work at the moment, so I'm putting a `priority:low` tag on this issue for the sake of house keeping. 

The first solution that you suggested - adding the filter option - is the right way to go, and should be quick to implement and test. 

Hey,Can i work on this issue?I would like to contribute to this project
@kishorhange111 definitely! This one's on the harder end of first issues, but we're happy to help out.
Can you provide me some resources to resolve this issue
Hey there! So step one is making sure that you've got your development environment up and running. I would go through our contributing guide first and make sure that you can get that working. Make sure that you're able to run the test suite as well.
https://docs.prefect.io/contributing/overview/

Once that's done you'll want to familiarize yourself with our codebase. Madkinsz is right that this one could be a bit challenging as a first issue, mostly because it's context heavy. I think it could be a great learning experience though. 

Familiarize yourself with our API code a bit first. Here's the flow_runs file: https://github.com/PrefectHQ/prefect/blob/main/src/prefect/orion/api/flow_runs.py

Then checkout the `/filter` endpoint and see how options to filter are currently setup. We have filter objects `schemas.filters.<MyFilterType>`. Familiarize yourself with the filters here: https://github.com/PrefectHQ/prefect/blob/main/src/prefect/orion/schemas/filters.py . They add SQLAlchemy code to a query to filter down the number of responses you're getting. You'll need to be able to filter down by idempotency key, which a field on the FlowRun (see https://github.com/PrefectHQ/prefect/blob/32b70259b579572e75105a57dbfe11f0e90129e3/src/prefect/orion/schemas/core.py#L189).

That should get you pointed in the right direction. You should join the #prefect-contributors channel on our slack community if you're not already a member. The link to join can be found here: https://www.prefect.io/community/

Resolving patch:
diff --git a/src/prefect/client/orchestration.py b/src/prefect/client/orchestration.py
--- a/src/prefect/client/orchestration.py
+++ b/src/prefect/client/orchestration.py
@@ -358,7 +358,9 @@ async def read_flows(
         body = {
             "flows": flow_filter.dict(json_compatible=True) if flow_filter else None,
             "flow_runs": (
-                flow_run_filter.dict(json_compatible=True) if flow_run_filter else None
+                flow_run_filter.dict(json_compatible=True, exclude_unset=True)
+                if flow_run_filter
+                else None
             ),
             "task_runs": (
                 task_run_filter.dict(json_compatible=True) if task_run_filter else None
@@ -1544,7 +1546,9 @@ async def read_deployments(
         body = {
             "flows": flow_filter.dict(json_compatible=True) if flow_filter else None,
             "flow_runs": (
-                flow_run_filter.dict(json_compatible=True) if flow_run_filter else None
+                flow_run_filter.dict(json_compatible=True, exclude_unset=True)
+                if flow_run_filter
+                else None
             ),
             "task_runs": (
                 task_run_filter.dict(json_compatible=True) if task_run_filter else None
@@ -1666,7 +1670,9 @@ async def read_flow_runs(
         body = {
             "flows": flow_filter.dict(json_compatible=True) if flow_filter else None,
             "flow_runs": (
-                flow_run_filter.dict(json_compatible=True) if flow_run_filter else None
+                flow_run_filter.dict(json_compatible=True, exclude_unset=True)
+                if flow_run_filter
+                else None
             ),
             "task_runs": (
                 task_run_filter.dict(json_compatible=True) if task_run_filter else None
@@ -1858,7 +1864,9 @@ async def read_task_runs(
         body = {
             "flows": flow_filter.dict(json_compatible=True) if flow_filter else None,
             "flow_runs": (
-                flow_run_filter.dict(json_compatible=True) if flow_run_filter else None
+                flow_run_filter.dict(json_compatible=True, exclude_unset=True)
+                if flow_run_filter
+                else None
             ),
             "task_runs": (
                 task_run_filter.dict(json_compatible=True) if task_run_filter else None
diff --git a/src/prefect/server/schemas/filters.py b/src/prefect/server/schemas/filters.py
--- a/src/prefect/server/schemas/filters.py
+++ b/src/prefect/server/schemas/filters.py
@@ -440,6 +440,25 @@ def _get_filter_list(self, db: "PrefectDBInterface") -> List:
         return filters
 
 
+class FlowRunFilterIdempotencyKey(PrefectFilterBaseModel):
+    """Filter by FlowRun.idempotency_key."""
+
+    any_: Optional[List[str]] = Field(
+        default=None, description="A list of flow run idempotency keys to include"
+    )
+    not_any_: Optional[List[str]] = Field(
+        default=None, description="A list of flow run idempotency keys to exclude"
+    )
+
+    def _get_filter_list(self, db: "PrefectDBInterface") -> List:
+        filters = []
+        if self.any_ is not None:
+            filters.append(db.FlowRun.idempotency_key.in_(self.any_))
+        if self.not_any_ is not None:
+            filters.append(db.FlowRun.idempotency_key.not_in(self.not_any_))
+        return filters
+
+
 class FlowRunFilter(PrefectOperatorFilterBaseModel):
     """Filter flow runs. Only flow runs matching all criteria will be returned"""
 
@@ -477,6 +496,9 @@ class FlowRunFilter(PrefectOperatorFilterBaseModel):
     parent_task_run_id: Optional[FlowRunFilterParentTaskRunId] = Field(
         default=None, description="Filter criteria for `FlowRun.parent_task_run_id`"
     )
+    idempotency_key: Optional[FlowRunFilterIdempotencyKey] = Field(
+        default=None, description="Filter criteria for `FlowRun.idempotency_key`"
+    )
 
     def _get_filter_list(self, db: "PrefectDBInterface") -> List:
         filters = []
@@ -503,6 +525,8 @@ def _get_filter_list(self, db: "PrefectDBInterface") -> List:
             filters.append(self.next_scheduled_start_time.as_sql_filter(db))
         if self.parent_task_run_id is not None:
             filters.append(self.parent_task_run_id.as_sql_filter(db))
+        if self.idempotency_key is not None:
+            filters.append(self.idempotency_key.as_sql_filter(db))
 
         return filters
 


Test patch:
diff --git a/tests/server/api/test_flow_runs.py b/tests/server/api/test_flow_runs.py
--- a/tests/server/api/test_flow_runs.py
+++ b/tests/server/api/test_flow_runs.py
@@ -305,6 +305,39 @@ async def flow_runs(self, flow, work_queue_1, session):
         await session.commit()
         return [flow_run_1, flow_run_2, flow_run_3]
 
+    @pytest.fixture
+    async def flow_runs_with_idempotency_key(
+        self, flow, work_queue_1, session
+    ) -> List[core.FlowRun]:
+        """
+        Return a list of two `core.FlowRun`'s with different idempotency keys.
+        """
+        flow_run_1_with_idempotency_key = await models.flow_runs.create_flow_run(
+            session=session,
+            flow_run=actions.FlowRunCreate(
+                flow_id=flow.id,
+                name="fr1",
+                tags=["red"],
+                idempotency_key="my-idempotency-key",
+            ),
+        )
+        flow_run_2_with_a_different_idempotency_key = (
+            await models.flow_runs.create_flow_run(
+                session=session,
+                flow_run=actions.FlowRunCreate(
+                    flow_id=flow.id,
+                    name="fr2",
+                    tags=["blue"],
+                    idempotency_key="a-different-idempotency-key",
+                ),
+            )
+        )
+        await session.commit()
+        return [
+            flow_run_1_with_idempotency_key,
+            flow_run_2_with_a_different_idempotency_key,
+        ]
+
     async def test_read_flow_runs(self, flow_runs, client):
         response = await client.post("/flow_runs/filter")
         assert response.status_code == status.HTTP_200_OK
@@ -356,6 +389,77 @@ async def test_read_flow_runs_applies_flow_run_filter(
         assert len(response.json()) == 1
         assert response.json()[0]["id"] == str(flow_runs[0].id)
 
+    async def test_read_flow_runs_applies_flow_run_idempotency_key_filter(
+        self, flow_runs_with_idempotency_key, client
+    ):
+        """
+        This test tests that when we pass a value for idempotency key to the flow run
+        filter, we get back the flow run with the matching idempotency key.
+        """
+        idempotency_key_of_flow_run_we_want_to_retrieve = (
+            flow_runs_with_idempotency_key[0].idempotency_key
+        )
+        flow_run_idempotency_key_filter = dict(
+            flow_runs=schemas.filters.FlowRunFilter(
+                idempotency_key=schemas.filters.FlowRunFilterIdempotencyKey(
+                    any_=[idempotency_key_of_flow_run_we_want_to_retrieve]
+                )
+            ).dict(json_compatible=True)
+        )
+        response = await client.post(
+            "/flow_runs/filter", json=flow_run_idempotency_key_filter
+        )
+        flow_run_filter_results = response.json()
+        assert response.status_code == status.HTTP_200_OK
+        assert (
+            len(flow_run_filter_results) == 1
+            and len(flow_runs_with_idempotency_key) == 2
+        )
+        assert flow_run_filter_results[0]["idempotency_key"] == str(
+            idempotency_key_of_flow_run_we_want_to_retrieve
+        )
+
+    async def test_read_flow_runs_idempotency_key_filter_excludes_idempotency_key(
+        self, flow_runs_with_idempotency_key, client
+    ):
+        """
+        This test tests to make sure that when you pass idempotency keys to the not_any_ argument
+        of the filter, the filter excludes flow runs having that value for idempotency key
+        """
+        idempotency_key_of_flow_run_to_exclude: str = flow_runs_with_idempotency_key[
+            0
+        ].idempotency_key
+        idempotency_key_of_flow_run_that_should_be_included: str = (
+            flow_runs_with_idempotency_key[1].idempotency_key
+        )
+        flow_run_idempotency_key_exclude_filter = dict(
+            flow_runs=schemas.filters.FlowRunFilter(
+                idempotency_key=schemas.filters.FlowRunFilterIdempotencyKey(
+                    not_any_=[idempotency_key_of_flow_run_to_exclude]
+                )
+            ).dict(json_compatible=True)
+        )
+        response = await client.post(
+            "/flow_runs/filter", json=flow_run_idempotency_key_exclude_filter
+        )
+        flow_run_filter_results = response.json()
+        assert response.status_code == status.HTTP_200_OK
+
+        # assert we started with two flow runs from fixture
+        assert len(flow_runs_with_idempotency_key) == 2
+
+        # filtering the fixture should result in a single element
+        assert len(flow_run_filter_results) == 1
+
+        # make sure the idempotency key we're excluding is not included in the results
+        for result in flow_run_filter_results:
+            assert result["idempotency_key"] != idempotency_key_of_flow_run_to_exclude
+
+        # make sure the idempotency key we did not exclude is still in the results
+        assert flow_run_filter_results[0]["idempotency_key"] == str(
+            idempotency_key_of_flow_run_that_should_be_included
+        )
+
     async def test_read_flow_runs_applies_task_run_filter(
         self, flow, flow_runs, client, session
     ):