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)