def test_tokens_are_passed_to_post(self, monkeypatch):
post = MagicMock()
session = MagicMock()
session.return_value.post = post
monkeypatch.setattr("requests.Session", session)
with set_temporary_config({"cloud.graphql": "http://my-cloud.foo"}):
client = Client()
client.post("/foo/bar", token="secret_token")
assert post.called
assert post.call_args[1]["headers"] == {"Authorization": "Bearer secret_token"}
def test_headers_are_passed_to_post(self, monkeypatch):
post = MagicMock()
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"}
):
client = Client()
client.post("/foo/bar", headers={"x": "y", "Authorization": "z"})
assert post.called
assert post.call_args[1]["headers"] == {
"x": "y",
"Authorization": "Bearer secret_token",
"X-PREFECT-CORE-VERSION": str(prefect.__version__),
}
def test_client_posts_raises_with_no_token(monkeypatch):
post = MagicMock()
monkeypatch.setattr("requests.post", post)
with set_temporary_config(
{"cloud.graphql": "http://my-cloud.foo", "cloud.auth_token": None}
):
client = Client()
with pytest.raises(AuthorizationError) as exc:
result = client.post("/foo/bar")
assert "Client.login" in str(exc.value)
def test_client_posts_to_api_server(patch_post):
post = patch_post(dict(success=True))
with set_temporary_config(
{"cloud.graphql": "http://my-cloud.foo", "cloud.auth_token": "secret_token"}
):
client = Client()
result = client.post("/foo/bar")
assert result == {"success": True}
assert post.called
assert post.call_args[0][0] == "http://my-cloud.foo/foo/bar"
def test_client_posts_raises_with_no_token(monkeypatch):
post = MagicMock()
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": None
}):
client = Client()
with pytest.raises(AuthorizationError, match="Client.login"):
result = client.post("/foo/bar")
def test_client_posts_to_graphql_server(monkeypatch):
post = MagicMock(return_value=MagicMock(json=MagicMock(return_value=dict(
success=True))))
monkeypatch.setattr("requests.post", post)
with set_temporary_config({
"cloud.graphql": "http://my-cloud.foo",
"cloud.auth_token": "secret_token"
}):
client = Client()
result = client.post("/foo/bar")
assert result == {"success": True}
assert post.called
assert post.call_args[0][0] == "http://my-cloud.foo/foo/bar"
def test_client_posts_retries_if_token_needs_refreshing(monkeypatch):
error = requests.HTTPError()
error.response = MagicMock(status_code=401) # unauthorized
post = MagicMock(return_value=MagicMock(
raise_for_status=MagicMock(side_effect=error),
json=MagicMock(return_value=dict(token="new-token")),
))
monkeypatch.setattr("requests.post", post)
with set_temporary_config({
"cloud.graphql": "http://my-cloud.foo",
"cloud.auth_token": "secret_token"
}):
client = Client()
with pytest.raises(requests.HTTPError) as exc:
result = client.post("/foo/bar")
assert exc.value is error
assert post.call_count == 3 # first call -> refresh token -> last call
assert post.call_args[0][0] == "http://my-cloud.foo/foo/bar"
assert client.token == "new-token"
class CloudResultHandler(ResultHandler):
"""
Hook for storing and retrieving task results from Prefect cloud storage.
Args:
- result_handler_service (str, optional): the location of the service
which will further process and store the results; if not provided, will default to
the value of `cloud.result_handler` in your config file
"""
def __init__(self, result_handler_service: str = None) -> None:
self._client = None
if result_handler_service is None:
self.result_handler_service = config.cloud.result_handler
else:
self.result_handler_service = result_handler_service
super().__init__()
def _initialize_client(self) -> None:
"""
Helper method for ensuring that CloudHandlers which are initialized locally
do not attempt to start a Client. This is important because CloudHandlers are
currently attached to `Flow` objects which need to be serialized / deserialized
independently of cloud settings.
This will instantiate a Client upon the first call to (de)serialize.
"""
if self._client is None:
self._client = Client() # type: ignore
def read(self, uri: str) -> Any:
"""
Read a result from the given URI location.
Args:
- uri (str): the path to the location of a result
Returns:
- the deserialized result from the provided URI
"""
self._initialize_client()
self.logger.debug("Starting to read result from {}...".format(uri))
res = self._client.get( # type: ignore
"/",
server=self.result_handler_service,
**{"uri": uri})
try:
return_val = cloudpickle.loads(
base64.b64decode(res.get("result", "")))
except EOFError:
return_val = None
self.logger.debug("Finished reading result from {}...".format(uri))
return return_val
def write(self, result: Any) -> str:
"""
Write the provided result to Prefect Cloud.
Args:
- result (Any): the result to store
Returns:
- str: the URI path to the result in Cloud storage
"""
self._initialize_client()
binary_data = base64.b64encode(cloudpickle.dumps(result)).decode()
self.logger.debug("Starting to upload result to {}...".format(
self.result_handler_service))
res = self._client.post( # type: ignore
"/",
server=self.result_handler_service,
**{"result": binary_data})
self.logger.debug("Finished uploading result to {}...".format(
self.result_handler_service))
return res.get("uri", "")
