def test_graphql_uses_access_token_after_login(self, patch_post):
tenant_id = str(uuid.uuid4())
post = patch_post(
{
"data": {
"tenant": [{"id": tenant_id}],
"switch_tenant": {
"access_token": "ACCESS_TOKEN",
"expires_at": "2100-01-01",
"refresh_token": "REFRESH_TOKEN",
},
}
}
)
client = Client(api_token="api")
client.graphql({})
assert client.get_auth_token() == "api"
assert post.call_args[1]["headers"] == {
"Authorization": "Bearer api",
"X-PREFECT-CORE-VERSION": str(prefect.__version__),
}
client.login_to_tenant(tenant_id=tenant_id)
client.graphql({})
assert client.get_auth_token() == "ACCESS_TOKEN"
assert post.call_args[1]["headers"] == {
"Authorization": "Bearer ACCESS_TOKEN",
"X-PREFECT-CORE-VERSION": str(prefect.__version__),
}
def test_graphql_errors_get_raised(patch_post):
patch_post(dict(data="42", errors="GraphQL issue!"))
with set_temporary_config(
{"cloud.api": "http://my-cloud.foo", "cloud.auth_token": "secret_token"}
):
client = Client()
with pytest.raises(ClientError, match="GraphQL issue!"):
client.graphql("query: {}")
def test_tokens_are_passed_to_graphql(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.graphql("query {}", token="secret_token")
assert post.called
assert post.call_args[1]["headers"] == {"Authorization": "Bearer secret_token"}
def test_graphql_errors_get_raised(monkeypatch):
post = MagicMock(return_value=MagicMock(json=MagicMock(
return_value=dict(data="42", errors="GraphQL issue!"))))
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()
with pytest.raises(ClientError, match="GraphQL issue!"):
client.graphql("query: {}")
def test_headers_are_passed_to_graphql(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.graphql("query {}", 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 get(self) -> Optional[Any]:
"""
Retrieve the secret value. If not found, returns `None`.
If using local secrets, `Secret.get()` will attempt to call `json.loads` on the
value pulled from context. For this reason it is recommended to store local secrets as
JSON documents to avoid ambiguous behavior.
Returns:
- Any: the value of the secret; if not found, returns `None`
Raises:
- ValueError: if `use_local_secrets=False` and the Client fails to retrieve your secret
"""
if prefect.config.cloud.use_local_secrets is True:
secrets = prefect.context.get("secrets", {})
value = secrets.get(self.name)
try:
return json.loads(value)
except (json.JSONDecodeError, TypeError):
return value
else:
client = Client()
result = client.graphql(
"""
query($name: String!) {
secretValue(name: $name)
}
""",
name=self.name,
) # type: Any
return as_nested_dict(result.data.secretValue, dict)
def get(self) -> Optional[Any]:
"""
Retrieve the secret value.
If not found, returns `None`.
Returns:
- Any: the value of the secret; if not found, returns `None`
Raises:
- ValueError: if `use_local_secrets=False` and the Client fails to retrieve your secret
"""
if prefect.config.cloud.use_local_secrets is True:
secrets = prefect.context.get("secrets", {})
return secrets.get(self.name)
else:
client = Client()
result = client.graphql(
"""
query($name: String!) {
secretValue(name: $name)
}
""",
name=self.name,
) # type: Any
return as_nested_dict(result.data.secretValue, dict)
def test_client_posts_graphql_to_api_server(patch_post):
post = patch_post(dict(data=dict(success=True)))
with set_temporary_config(
{"cloud.graphql": "http://my-cloud.foo", "cloud.auth_token": "secret_token"}
):
client = Client()
result = client.graphql("{projects{name}}")
assert result.data == {"success": True}
assert post.called
assert post.call_args[0][0] == "http://my-cloud.foo"
def test_graphql_errors_get_raised(monkeypatch):
post = MagicMock(return_value=MagicMock(json=MagicMock(
return_value=dict(data="42", errors="GraphQL issue!"))))
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(ClientError) as exc:
res = client.graphql("query: {}")
assert "GraphQL issue!" in str(exc.value)
def test_client_posts_graphql_to_graphql_server(monkeypatch):
post = MagicMock(return_value=MagicMock(json=MagicMock(return_value=dict(
data=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.graphql("{projects{name}}")
assert result.data == {"success": True}
assert post.called
assert post.call_args[0][0] == "http://my-cloud.foo"
def get(self) -> Optional[Any]:
"""
Retrieve the secret value. If not found, returns `None`.
If using local secrets, `Secret.get()` will attempt to call `json.loads` on the
value pulled from context. For this reason it is recommended to store local secrets as
JSON documents to avoid ambiguous behavior.
Returns:
- Any: the value of the secret; if not found, raises an error
Raises:
- ValueError: if `.get()` is called within a Flow building context, or if `use_local_secrets=True`
and your Secret doesn't exist
- ClientError: if `use_local_secrets=False` and the Client fails to retrieve your secret
"""
if isinstance(prefect.context.get("flow"), prefect.core.flow.Flow):
raise ValueError(
"Secrets should only be retrieved during a Flow run, not while building a Flow."
)
if prefect.context.config.cloud.use_local_secrets is True:
secrets = prefect.context.get("secrets", {})
try:
value = secrets[self.name]
except KeyError:
raise ValueError('Local Secret "{}" was not found.'.format(
self.name)) from None
try:
return json.loads(value)
except (json.JSONDecodeError, TypeError):
return value
else:
client = Client()
result = client.graphql(
"""
query($name: String!) {
secretValue(name: $name)
}
""",
variables=dict(name=self.name),
)
# the result object is a Box, so we recursively restore builtin
# dict/list classes
result_dict = result.to_dict()
value = result_dict["data"]["secretValue"]
try:
return json.loads(value)
except (json.JSONDecodeError, TypeError):
return value
def test_graphql_uses_access_token_after_login(self, patch_post):
tenant_id = str(uuid.uuid4())
post = patch_post(
{
"data": {
"tenant": [{"id": tenant_id}],
"switchTenant": {
"accessToken": "ACCESS_TOKEN",
"expiresIn": 600,
"refreshToken": "REFRESH_TOKEN",
},
}
}
)
client = Client(api_token="api")
client.graphql({})
assert client.get_auth_token() == "api"
assert post.call_args[1]["headers"] == dict(Authorization="Bearer api")
client.login_to_tenant(tenant_id=tenant_id)
client.graphql({})
assert client.get_auth_token() == "ACCESS_TOKEN"
assert post.call_args[1]["headers"] == dict(Authorization="Bearer ACCESS_TOKEN")
def test_client_graphql_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.graphql("{}")
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"
assert client.token == "new-token"
