Fix bug with matching cached states when cache_key is None

Rogell · Aug 16, 2020 · eae74eb · eae74eb
1 parent a88c861
commit eae74eb
Showing 1 changed file with 11 additions and 11 deletions.
diff --git a/src/prefect/client/client.py b/src/prefect/client/client.py
@@ -1154,23 +1154,23 @@ def get_latest_cached_states(
         Returns:
             - List[State]: a list of Cached states created after the given date
         """
-        where_clause = {
+        args = {
             "where": {
                 "state": {"_eq": "Cached"},
-                "_or": [
-                    {
-                        "_and": [
-                            {"cache_key": {"_eq": cache_key}},
-                            {"cache_key": {"_is_null": False}},
-                        ]
-                    },
-                    {"task_id": {"_eq": task_id}},
-                ],
                 "state_timestamp": {"_gte": created_after.isoformat()},
             },
             "order_by": {"state_timestamp": EnumValue("desc")},
+            "limit": 100,
         }
-        query = {"query": {with_args("task_run", where_clause): "serialized_state"}}
+
+        # if a cache key was provided, match it against all tasks
+        if cache_key is not None:
+            args["where"].update({"cache_key": {"_eq": cache_key}})
+        # otherwise match against only this task, across all cache keys
+        else:
+            args["where"].update({"task_id": {"_eq": task_id}})
+
+        query = {"query": {with_args("task_run", args): "serialized_state"}}
         result = self.graphql(query)  # type: Any
         deserializer = prefect.engine.state.State.deserialize
         valid_states = [