async def get_url(self, *, route: str, data: Union[None, dict] = None) -> httpx.Response: """ Attempt to get something from the API, and return the raw response object if an error response wasn't received. If an error response was received, raises an error. Parameters ---------- route : str Path relative to API host to get data : dict, optional JSON data to send in the request body (optional) Returns ------- requests.Response """ logger.debug(f"Getting {self.url}/api/{self.api_version}/{route}") try: response = await self.session.request( "GET", route, follow_redirects=False, json=data, ) except RequestError as e: error_msg = f"Unable to connect to FlowKit API at {self.url}: {e}" logger.info(error_msg) raise FlowclientConnectionError(error_msg) if response.status_code in {202, 200, 303}: return response elif response.status_code == 404: raise FileNotFoundError( f"{self.url}/api/{self.api_version}/{route} not found.") elif response.status_code in {401, 403}: try: error = response.json()["msg"] except (ValueError, KeyError): error = "Unknown access denied error" raise FlowclientConnectionError(error) else: try: error = response.json()["msg"] except (ValueError, KeyError): error = "Unknown error" try: status = response.json()["status"] except (ValueError, KeyError): status = "Unknown status" raise FlowclientConnectionError( f"Something went wrong: {error}. API returned with status code: {response.status_code} and status '{status}'" )
async def post_json(self, *, route: str, data: dict) -> httpx.Response: """ Attempt to post json to the API, and return the raw response object if an error response wasn't received. If an error response was received, raises an error. Parameters ---------- route : str Path relative to API host to post_json to data: dict Dictionary of json-encodeable data to post_json Returns ------- requests.Response """ logger.debug(f"Posting {data} to {self.url}/api/{self.api_version}/{route}") try: response = await self.session.post( f"{self.url}/api/{self.api_version}/{route}", json=data ) except RequestError as e: error_msg = f"Unable to connect to FlowKit API at {self.url}: {e}" logger.info(error_msg) raise FlowclientConnectionError(error_msg) if response.status_code == 202: return response elif response.status_code == 404: raise FileNotFoundError( f"{self.url}/api/{self.api_version}/{route} not found." ) elif response.status_code in {401, 403}: try: error_msg = response.json()["msg"] except ValueError: error_msg = "Unknown access denied error" raise FlowclientConnectionError(error_msg) else: try: error_msg = response.json()["msg"] try: returned_payload = response.json()["payload"] payload_info = ( "" if not returned_payload else f" Payload: {returned_payload}" ) except KeyError: payload_info = "" except ValueError: # Happens if the response body does not contain valid JSON # (see http://docs.python-requests.org/en/master/api/#requests.Response.json) error_msg = f"the response did not contain valid JSON" payload_info = "" raise FlowclientConnectionError( f"Something went wrong. API returned with status code {response.status_code}. Error message: '{error_msg}'.{payload_info}" )
def run_query(*, connection: Connection, query_spec: dict) -> str: """ Run a query of a specified kind with parameters and get the identifier for it. Parameters ---------- connection : Connection API connection to use query_spec : dict Query specification to run Returns ------- str Identifier of the query """ logger.info( f"Requesting run of {query_spec} at {connection.url}/api/{connection.api_version}" ) r = connection.post_json(route="run", data=query_spec) if r.status_code == 202: query_id = r.headers["Location"].split("/").pop() logger.info( f"Accepted {query_spec} at {connection.url}/api/{connection.api_version} with id {query_id}" ) return query_id else: try: error = r.json()["msg"] except (ValueError, KeyError): error = "Unknown error" raise FlowclientConnectionError( f"Error running the query: {error}. Status code: {r.status_code}.")
def get_status(*, connection: Connection, query_id: str) -> str: """ Check the status of a query. Parameters ---------- connection : Connection API connection to use query_id : str Identifier of the query to retrieve Returns ------- str Query status Raises ------ FlowclientConnectionError if response does not contain the query status """ try: ready, reply = query_is_ready(connection=connection, query_id=query_id) except FileNotFoundError: # Can't distinguish 'known', 'cancelled', 'resetting' and 'awol' from the error, # so return generic 'not_running' status. return "not_running" if ready: return "completed" else: try: return reply.json()["status"] except (KeyError, TypeError): raise FlowclientConnectionError(f"No status reported.")
def get_geography(*, connection: Connection, aggregation_unit: str) -> dict: """ Get geography data from the database. Parameters ---------- connection : Connection API connection to use aggregation_unit : str aggregation unit, e.g. 'admin3' Returns ------- dict geography data as a GeoJSON FeatureCollection """ logger.info( f"Getting {connection.url}/api/{connection.api_version}/geography/{aggregation_unit}" ) response = connection.get_url(route=f"geography/{aggregation_unit}") if response.status_code != 200: try: msg = response.json()["msg"] more_info = f" Reason: {msg}" except KeyError: more_info = "" raise FlowclientConnectionError( f"Could not get result. API returned with status code: {response.status_code}.{more_info}" ) result = response.json() logger.info( f"Got {connection.url}/api/{connection.api_version}/geography/{aggregation_unit}" ) return result
def get_json_dataframe(*, connection: Connection, location: str) -> pd.DataFrame: """ Get a dataframe from a json source. Parameters ---------- connection : Connection API connection to use location : str API enpoint to retrieve json from Returns ------- pandas.DataFrame Dataframe containing the result """ response = connection.get_url(route=location) if response.status_code != 200: try: msg = response.json()["msg"] more_info = f" Reason: {msg}" except KeyError: more_info = "" raise FlowclientConnectionError( f"Could not get result. API returned with status code: {response.status_code}.{more_info}" ) result = response.json() logger.info( f"Got {connection.url}/api/{connection.api_version}/{location}") return pd.DataFrame.from_records(result["query_result"])
def update_token(self, token: str) -> None: """ Replace this connection's API token with a new one. Parameters ---------- token : str JSON Web Token for this API server """ try: self.user = jwt.decode(token, verify=False)["identity"] except jwt.DecodeError: raise FlowclientConnectionError(f"Unable to decode token: '{token}'") except KeyError: raise FlowclientConnectionError(f"Token does not contain user identity.") self.token = token self.session.headers["Authorization"] = f"Bearer {self.token}"
def get_geojson_result_by_query_id( *, connection: Connection, query_id: str, poll_interval: int = 1, disable_progress: Optional[bool] = None, ) -> dict: """ Get a query by id, and return it as a geojson dict Parameters ---------- connection : Connection API connection to use query_id : str Identifier of the query to retrieve poll_interval : int Number of seconds to wait between checks for the query being ready disable_progress : bool, default None Set to True to disable progress bar display entirely, None to disable on non-TTY, or False to always enable Returns ------- dict geojson """ result_endpoint = get_result_location_from_id_when_ready( connection=connection, query_id=query_id, poll_interval=poll_interval, disable_progress=disable_progress, ) response = connection.get_url(route=f"{result_endpoint}.geojson") if response.status_code != 200: try: msg = response.json()["msg"] more_info = f" Reason: {msg}" except KeyError: more_info = "" raise FlowclientConnectionError( f"Could not get result. API returned with status code: {response.status_code}.{more_info}" ) return response.json()
def query_is_ready(*, connection: Connection, query_id: str) -> Tuple[bool, requests.Response]: """ Check if a query id has results available. Parameters ---------- connection : Connection API connection to use query_id : str Identifier of the query to retrieve Returns ------- Tuple[bool, requests.Response] True if the query result is available Raises ------ FlowclientConnectionError if query has errored """ logger.info( f"Polling server on {connection.url}/api/{connection.api_version}/poll/{query_id}" ) reply = connection.get_url(route=f"poll/{query_id}") if reply.status_code == 303: logger.info( f"{connection.url}/api/{connection.api_version}/poll/{query_id} ready." ) return True, reply # Query is ready, so exit the loop elif reply.status_code == 202: logger.info( "{eligible} parts to run, {queued} in queue and {running} running." .format(**reply.json()["progress"])) return False, reply else: raise FlowclientConnectionError( f"Something went wrong: {reply}. API returned with status code: {reply.status_code}" )
def get_available_dates(*, connection: Connection, event_types: Union[None, List[str]] = None) -> dict: """ Get available dates for different event types from the database. Parameters ---------- connection : Connection API connection to use event_types : list of str, optional The event types for which to return available dates (for example: ["calls", "sms"]). If None, return available dates for all available event types. Returns ------- dict Available dates in the format {event_type: [list of dates]} """ logger.info( f"Getting {connection.url}/api/{connection.api_version}/available_dates" ) response = connection.get_url(route=f"available_dates") if response.status_code != 200: try: msg = response.json()["msg"] more_info = f" Reason: {msg}" except KeyError: more_info = "" raise FlowclientConnectionError( f"Could not get available dates. API returned with status code: {response.status_code}.{more_info}" ) result = response.json()["available_dates"] logger.info( f"Got {connection.url}/api/{connection.api_version}/available_dates") if event_types is None: return result else: return {k: v for k, v in result.items() if k in event_types}