def _request(
self, method: str, path: str, params: dict = None, server: str = None
) -> "requests.models.Response":
"""
Runs any specified request (GET, POST, DELETE) against the server
Args:
- method (str): The type of request to be made (GET, POST, DELETE)
- path (str): Path of the API URL
- params (dict, optional): Parameters used for the request
- server (str, optional): The server to make requests against, base API
server is used if not specified
Returns:
- requests.models.Response: The response returned from the request
Raises:
- ClientError: if the client token is not in the context (due to not being logged in)
- ValueError: if a method is specified outside of the accepted GET, POST, DELETE
- requests.HTTPError: if a status code is returned that is not `200` or `401`
"""
# lazy import for performance
import requests
if server is None:
server = self.graphql_server
assert isinstance(server, str) # mypy assert
if self.token is None:
raise AuthorizationError("Call Client.login() to set the client token.")
url = os.path.join(server, path.lstrip("/")).rstrip("/")
params = params or {}
# 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)}
if method == "GET":
response = requests.get(url, headers=headers, params=params)
elif method == "POST":
response = requests.post(url, headers=headers, json=params)
elif method == "DELETE":
response = requests.delete(url, headers=headers)
else:
raise ValueError("Invalid method: {}".format(method))
# Check if request returned a successful status
response.raise_for_status()
return response
# If a 401 status code is returned, refresh the login token
try:
return request_fn()
except requests.HTTPError as err:
if err.response.status_code == 401:
self.refresh_token()
return request_fn()
raise
def _request(self,
method: str,
path: str,
params: dict = None,
server: str = None) -> "requests.models.Response":
"""
Runs any specified request (GET, POST, DELETE) against the server
Args:
- method (str): The type of request to be made (GET, POST, DELETE)
- path (str): Path of the API URL
- params (dict, optional): Parameters used for the request
- server (str, optional): The server to make requests against, base API
server is used if not specified
Returns:
- requests.models.Response: The response returned from the request
Raises:
- ClientError: if the client token is not in the context (due to not being logged in)
- ValueError: if a method is specified outside of the accepted GET, POST, DELETE
- requests.HTTPError: if a status code is returned that is not `200` or `401`
"""
# lazy import for performance
import requests
if server is None:
server = self.graphql_server
assert isinstance(server, str) # mypy assert
if self.token is None:
raise AuthorizationError(
"No token found; call Client.login() to set one.")
url = os.path.join(server, path.lstrip("/")).rstrip("/")
params = params or {}
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 = session.get(url, headers=headers, params=params)
elif method == "POST":
response = session.post(url, headers=headers, json=params)
elif method == "DELETE":
response = session.delete(url, headers=headers)
else:
raise ValueError("Invalid method: {}".format(method))
# Check if request returned a successful status
response.raise_for_status()
return response
def graphql(
self,
query: Any,
raise_on_error: bool = True,
headers: Dict[str, str] = None,
variables: Dict[str, JSONLike] = None,
token: str = None,
) -> GraphQLResult:
"""
Convenience function for running queries against the Prefect GraphQL API
Args:
- query (Any): A representation of a graphql query to be executed. It will be
parsed by prefect.utilities.graphql.parse_graphql().
- raise_on_error (bool): if True, a `ClientError` will be raised if the GraphQL
returns any `errors`.
- headers (dict): any additional headers that should be passed as part of the
request
- variables (dict): Variables to be filled into a query with the key being
equivalent to the variables that are accepted by the query
- token (str): an auth token. If not supplied, the `client.access_token` is used.
Returns:
- dict: Data returned from the GraphQL query
Raises:
- ClientError if there are errors raised by the GraphQL mutation
"""
result = self.post(
path="",
server=self.api_server,
headers=headers,
params=dict(query=parse_graphql(query),
variables=json.dumps(variables)),
token=token,
)
if raise_on_error and "errors" in result:
if "UNAUTHENTICATED" in str(result["errors"]):
raise AuthorizationError(result["errors"])
elif "Malformed Authorization header" in str(result["errors"]):
raise AuthorizationError(result["errors"])
raise ClientError(result["errors"])
else:
return GraphQLResult(result) # type: ignore
def _verify_token(self, token: str) -> None:
"""
Checks whether a token with a `RUNNER` scope was provided
Args:
- token (str): The provided agent token to verify
Raises:
- AuthorizationError: if token is empty or does not have a RUNNER role
"""
if not token:
raise AuthorizationError("No agent API token provided.")
# Check if RUNNER role
result = self.client.graphql(query="query { auth_info { api_token_scope } }")
if (
not result.data # type: ignore
or result.data.auth_info.api_token_scope != "RUNNER" # type: ignore
):
raise AuthorizationError("Provided token does not have a RUNNER scope.")
def create_prefect_project(environment: str, prefect_agent_token: str) -> None:
"""
Get the Prefect Agent definition for an environment that run workflows on AWS ECS Fargate
Parameters:
environment [str] -- environment to create the prefect project
prefect_agent_token [str] -- prefect token
"""
client = Client(api_token=prefect_agent_token)
try:
client.create_project(
project_name=f"{environment}_dataflow_automation")
except AuthorizationError:
error_message = (
"Invalid API token provided. Check that the secret "
"PREFECT_AGENT_TOKEN is set on the Github repository configuration"
)
logger.error(error_message)
raise AuthorizationError(error_message)
def login(
self,
email: str,
password: str,
account_slug: str = None,
account_id: str = None,
) -> None:
"""
Login to the server in order to gain access
Args:
- email (str): User's email on the platform
- password (str): User's password on the platform
- account_slug (str, optional): Slug that is unique to the user
- account_id (str, optional): Specific Account ID for this user to use
Raises:
- AuthorizationError if unable to login to the server (request does not return `200`)
"""
# lazy import for performance
import requests
# TODO: This needs to call the main graphql server and be adjusted for auth0
url = os.path.join(self.graphql_server, "login_email") # type: ignore
response = requests.post(
url,
auth=(email, password),
json=dict(account_id=account_id, account_slug=account_slug),
)
# Load the current auth token if able to login
if not response.ok:
raise AuthorizationError("Could not log in.")
self.token = response.json().get("token")
if self.token:
creds_path = os.path.expanduser("~/.prefect/.credentials")
if not os.path.exists(creds_path):
os.makedirs(creds_path)
with open(os.path.join(creds_path, "auth_token"), "w+") as f:
f.write(self.token)