Merge pull request #1187 from PrefectHQ/robust-hits

Robust Client calls

CHANGELOG.md

@@ -17,6 +17,7 @@ These changes are available in the [master branch](https://github.com/PrefectHQ/
 - Validate configuration objects on initial load - [#1136](https://github.com/PrefectHQ/prefect/pull/1136)
 - Add `auto_generated` property to Tasks for convenient filtering - [#1135](https://github.com/PrefectHQ/prefect/pull/1135)
 - Disable dask work-stealing in kubernetes via scheduler config - [#1166](https://github.com/PrefectHQ/prefect/pull/1166)
+- Implement backoff retry settings on Client calls - [#1187](https://github.com/PrefectHQ/prefect/pull/1187)
 
 ### Task Library
 

src/prefect/client/client.py

@@ -3,6 +3,9 @@
 import json
 import logging
 import os
+
+from requests.adapters import HTTPAdapter
+from requests.packages.urllib3.util.retry import Retry
 from typing import TYPE_CHECKING, Any, Dict, List, NamedTuple, Optional, Union
 
 import pendulum
@@ -209,12 +212,20 @@ def _request(
         # write this as a function to allow reuse in next try/except block
         def request_fn() -> "requests.models.Response":
             headers = {"Authorization": "Bearer {}".format(self.token)}
+            session = requests.Session()
+            retries = Retry(
+                total=6,
+                backoff_factor=1,
+                status_forcelist=[500, 502, 503, 504],
+                method_whitelist=["DELETE", "GET", "POST"],
+            )
+            session.mount("https://", HTTPAdapter(max_retries=retries))
             if method == "GET":
-                response = requests.get(url, headers=headers, params=params)
+                response = session.get(url, headers=headers, params=params)
             elif method == "POST":
-                response = requests.post(url, headers=headers, json=params)
+                response = session.post(url, headers=headers, json=params)
             elif method == "DELETE":
-                response = requests.delete(url, headers=headers)
+                response = session.delete(url, headers=headers)
             else:
                 raise ValueError("Invalid method: {}".format(method))
 

tests/cli/test_auth.py

@@ -47,7 +47,9 @@ def test_auth_add(monkeypatch):
                 json=MagicMock(return_value=dict(data=dict(hello="hi")))
             )
         )
-        monkeypatch.setattr("requests.post", post)
+        session = MagicMock()
+        session.return_value.post = post
+        monkeypatch.setattr("requests.Session", session)
 
         with set_temporary_config(
             {"cloud.graphql": "http://my-cloud.foo", "cloud.auth_token": "secret_token"}
@@ -60,7 +62,7 @@ def test_auth_add(monkeypatch):
             assert "Auth token added to Prefect config" in result.output
 
 
-def test_auth_add_failes_query(monkeypatch):
+def test_auth_add_failed_query(monkeypatch):
     with tempfile.TemporaryDirectory() as temp_dir:
 
         file = "{}/temp_config.toml".format(temp_dir)
@@ -73,7 +75,9 @@ def test_auth_add_failes_query(monkeypatch):
                 json=MagicMock(return_value=dict(data=dict(hello=None)))
             )
         )
-        monkeypatch.setattr("requests.post", post)
+        session = MagicMock()
+        session.return_value.post = post
+        monkeypatch.setattr("requests.Session", session)
 
         with set_temporary_config(
             {"cloud.graphql": "http://my-cloud.foo", "cloud.auth_token": "secret_token"}

tests/cli/test_describe.py

@@ -40,7 +40,9 @@ def test_describe_flows(monkeypatch):
             json=MagicMock(return_value=dict(data=dict(flow=[{"name": "flow"}])))
         )
     )
-    monkeypatch.setattr("requests.post", post)
+    session = MagicMock()
+    session.return_value.post = post
+    monkeypatch.setattr("requests.Session", session)
 
     with set_temporary_config(
         {"cloud.graphql": "http://my-cloud.foo", "cloud.auth_token": "secret_token"}
@@ -76,7 +78,9 @@ def test_describe_flows_not_found(monkeypatch):
     post = MagicMock(
         return_value=MagicMock(json=MagicMock(return_value=dict(data=dict(flow=[]))))
     )
-    monkeypatch.setattr("requests.post", post)
+    session = MagicMock()
+    session.return_value.post = post
+    monkeypatch.setattr("requests.Session", session)
 
     with set_temporary_config(
         {"cloud.graphql": "http://my-cloud.foo", "cloud.auth_token": "secret_token"}
@@ -96,7 +100,9 @@ def test_describe_flows_populated(monkeypatch):
             json=MagicMock(return_value=dict(data=dict(flow=[{"name": "flow"}])))
         )
     )
-    monkeypatch.setattr("requests.post", post)
+    session = MagicMock()
+    session.return_value.post = post
+    monkeypatch.setattr("requests.Session", session)
 
     with set_temporary_config(
         {"cloud.graphql": "http://my-cloud.foo", "cloud.auth_token": "secret_token"}
@@ -141,7 +147,9 @@ def test_describe_tasks(monkeypatch):
             )
         )
     )
-    monkeypatch.setattr("requests.post", post)
+    session = MagicMock()
+    session.return_value.post = post
+    monkeypatch.setattr("requests.Session", session)
 
     with set_temporary_config(
         {"cloud.graphql": "http://my-cloud.foo", "cloud.auth_token": "secret_token"}
@@ -176,7 +184,9 @@ def test_describe_tasks_flow_not_found(monkeypatch):
     post = MagicMock(
         return_value=MagicMock(json=MagicMock(return_value=dict(data=dict(flow=[]))))
     )
-    monkeypatch.setattr("requests.post", post)
+    session = MagicMock()
+    session.return_value.post = post
+    monkeypatch.setattr("requests.Session", session)
 
     with set_temporary_config(
         {"cloud.graphql": "http://my-cloud.foo", "cloud.auth_token": "secret_token"}
@@ -193,7 +203,9 @@ def test_describe_tasks_not_found(monkeypatch):
             json=MagicMock(return_value=dict(data=dict(flow=[{"tasks": []}])))
         )
     )
-    monkeypatch.setattr("requests.post", post)
+    session = MagicMock()
+    session.return_value.post = post
+    monkeypatch.setattr("requests.Session", session)
 
     with set_temporary_config(
         {"cloud.graphql": "http://my-cloud.foo", "cloud.auth_token": "secret_token"}
@@ -215,7 +227,9 @@ def test_describe_flow_runs(monkeypatch):
             )
         )
     )
-    monkeypatch.setattr("requests.post", post)
+    session = MagicMock()
+    session.return_value.post = post
+    monkeypatch.setattr("requests.Session", session)
 
     with set_temporary_config(
         {"cloud.graphql": "http://my-cloud.foo", "cloud.auth_token": "secret_token"}
@@ -255,7 +269,9 @@ def test_describe_flow_runs_not_found(monkeypatch):
             json=MagicMock(return_value=dict(data=dict(flow_run=[])))
         )
     )
-    monkeypatch.setattr("requests.post", post)
+    session = MagicMock()
+    session.return_value.post = post
+    monkeypatch.setattr("requests.Session", session)
 
     with set_temporary_config(
         {"cloud.graphql": "http://my-cloud.foo", "cloud.auth_token": "secret_token"}
@@ -277,7 +293,9 @@ def test_describe_flow_runs_populated(monkeypatch):
             )
         )
     )
-    monkeypatch.setattr("requests.post", post)
+    session = MagicMock()
+    session.return_value.post = post
+    monkeypatch.setattr("requests.Session", session)
 
     with set_temporary_config(
         {"cloud.graphql": "http://my-cloud.foo", "cloud.auth_token": "secret_token"}

tests/cli/test_execute.py

@@ -39,7 +39,9 @@ def test_execute_cloud_flow_not_found(monkeypatch):
             json=MagicMock(return_value=dict(data=dict(flow_run=[])))
         )
     )
-    monkeypatch.setattr("requests.post", post)
+    session = MagicMock()
+    session.return_value.post = post
+    monkeypatch.setattr("requests.Session", session)
 
     with set_temporary_config(
         {"cloud.graphql": "http://my-cloud.foo", "cloud.auth_token": "secret_token"}

tests/cli/test_get.py

@@ -36,7 +36,9 @@ def test_get_flows(monkeypatch):
     post = MagicMock(
         return_value=MagicMock(json=MagicMock(return_value=dict(data=dict(flow=[]))))
     )
-    monkeypatch.setattr("requests.post", post)
+    session = MagicMock()
+    session.return_value.post = post
+    monkeypatch.setattr("requests.Session", session)
 
     with set_temporary_config(
         {"cloud.graphql": "http://my-cloud.foo", "cloud.auth_token": "secret_token"}
@@ -75,7 +77,9 @@ def test_get_flows_populated(monkeypatch):
     post = MagicMock(
         return_value=MagicMock(json=MagicMock(return_value=dict(data=dict(flow=[]))))
     )
-    monkeypatch.setattr("requests.post", post)
+    session = MagicMock()
+    session.return_value.post = post
+    monkeypatch.setattr("requests.Session", session)
 
     with set_temporary_config(
         {"cloud.graphql": "http://my-cloud.foo", "cloud.auth_token": "secret_token"}
@@ -122,7 +126,9 @@ def test_get_projects(monkeypatch):
     post = MagicMock(
         return_value=MagicMock(json=MagicMock(return_value=dict(data=dict(project=[]))))
     )
-    monkeypatch.setattr("requests.post", post)
+    session = MagicMock()
+    session.return_value.post = post
+    monkeypatch.setattr("requests.Session", session)
 
     with set_temporary_config(
         {"cloud.graphql": "http://my-cloud.foo", "cloud.auth_token": "secret_token"}
@@ -163,7 +169,9 @@ def test_get_projects_populated(monkeypatch):
     post = MagicMock(
         return_value=MagicMock(json=MagicMock(return_value=dict(data=dict(project=[]))))
     )
-    monkeypatch.setattr("requests.post", post)
+    session = MagicMock()
+    session.return_value.post = post
+    monkeypatch.setattr("requests.Session", session)
 
     with set_temporary_config(
         {"cloud.graphql": "http://my-cloud.foo", "cloud.auth_token": "secret_token"}
@@ -200,7 +208,9 @@ def test_get_flow_runs(monkeypatch):
             json=MagicMock(return_value=dict(data=dict(flow_run=[])))
         )
     )
-    monkeypatch.setattr("requests.post", post)
+    session = MagicMock()
+    session.return_value.post = post
+    monkeypatch.setattr("requests.Session", session)
 
     with set_temporary_config(
         {"cloud.graphql": "http://my-cloud.foo", "cloud.auth_token": "secret_token"}
@@ -245,7 +255,9 @@ def test_get_flow_runs_populated(monkeypatch):
             json=MagicMock(return_value=dict(data=dict(flow_run=[])))
         )
     )
-    monkeypatch.setattr("requests.post", post)
+    session = MagicMock()
+    session.return_value.post = post
+    monkeypatch.setattr("requests.Session", session)
 
     with set_temporary_config(
         {"cloud.graphql": "http://my-cloud.foo", "cloud.auth_token": "secret_token"}
@@ -292,7 +304,9 @@ def test_get_tasks(monkeypatch):
     post = MagicMock(
         return_value=MagicMock(json=MagicMock(return_value=dict(data=dict(task=[]))))
     )
-    monkeypatch.setattr("requests.post", post)
+    session = MagicMock()
+    session.return_value.post = post
+    monkeypatch.setattr("requests.Session", session)
 
     with set_temporary_config(
         {"cloud.graphql": "http://my-cloud.foo", "cloud.auth_token": "secret_token"}
@@ -335,7 +349,9 @@ def test_get_tasks_populated(monkeypatch):
     post = MagicMock(
         return_value=MagicMock(json=MagicMock(return_value=dict(data=dict(task=[]))))
     )
-    monkeypatch.setattr("requests.post", post)
+    session = MagicMock()
+    session.return_value.post = post
+    monkeypatch.setattr("requests.Session", session)
 
     with set_temporary_config(
         {"cloud.graphql": "http://my-cloud.foo", "cloud.auth_token": "secret_token"}
@@ -403,7 +419,9 @@ def test_get_logs(monkeypatch):
             )
         )
     )
-    monkeypatch.setattr("requests.post", post)
+    session = MagicMock()
+    session.return_value.post = post
+    monkeypatch.setattr("requests.Session", session)
 
     with set_temporary_config(
         {"cloud.graphql": "http://my-cloud.foo", "cloud.auth_token": "secret_token"}
@@ -446,7 +464,9 @@ def test_get_logs_info(monkeypatch):
             )
         )
     )
-    monkeypatch.setattr("requests.post", post)
+    session = MagicMock()
+    session.return_value.post = post
+    monkeypatch.setattr("requests.Session", session)
 
     with set_temporary_config(
         {"cloud.graphql": "http://my-cloud.foo", "cloud.auth_token": "secret_token"}
@@ -478,7 +498,9 @@ def test_get_logs_fails(monkeypatch):
             json=MagicMock(return_value=dict(data=dict(flow_run=[])))
         )
     )
-    monkeypatch.setattr("requests.post", post)
+    session = MagicMock()
+    session.return_value.post = post
+    monkeypatch.setattr("requests.Session", session)
 
     with set_temporary_config(
         {"cloud.graphql": "http://my-cloud.foo", "cloud.auth_token": "secret_token"}

tests/cli/test_run.py

@@ -34,7 +34,9 @@ def test_run_cloud(monkeypatch):
             json=MagicMock(return_value=dict(data=dict(flow=[{"id": "flow"}])))
         )
     )
-    monkeypatch.setattr("requests.post", post)
+    session = MagicMock()
+    session.return_value.post = post
+    monkeypatch.setattr("requests.Session", session)
 
     create_flow_run = MagicMock(resurn_value="id")
     monkeypatch.setattr(
@@ -67,7 +69,9 @@ def test_run_cloud_fails(monkeypatch):
     post = MagicMock(
         return_value=MagicMock(json=MagicMock(return_value=dict(data=dict(flow=[]))))
     )
-    monkeypatch.setattr("requests.post", post)
+    session = MagicMock()
+    session.return_value.post = post
+    monkeypatch.setattr("requests.Session", session)
 
     with set_temporary_config(
         {"cloud.graphql": "http://my-cloud.foo", "cloud.auth_token": "secret_token"}