Coverage for crunch/client/connections.py: 100.00%
131 statements
« prev ^ index » next coverage.py v6.5.0, created at 2023-10-01 13:43 +0000
« prev ^ index » next coverage.py v6.5.0, created at 2023-10-01 13:43 +0000
1from os import getenv
2from typing import Union, Dict
3from unicodedata import decimal
4import requests
5from rest_framework import status as drf_status
6import enum
7from datetime import datetime, date
8from django.core.validators import URLValidator
9from django.core.exceptions import ValidationError
10from rich.console import Console
11from crunch.django.app.enums import Stage, State
13console = Console()
15from . import diagnostics
17CRUNCH_URL_KEY = "CRUNCH_URL"
18CRUNCH_TOKEN_KEY = "CRUNCH_TOKEN"
20class CrunchAPIException(Exception):
21 """ Raised when there is an error getting information from the API of a crunch site. """
22 pass
25class Connection():
26 """
27 An object to manage calls to the REST API of a crunch hosted site.
28 """
29 def __init__(self, base_url:str = None, token:str = None, verbose:bool = False):
30 """
31 An object to manage calls to the REST API of a crunch hosted site.
33 Args:
34 base_url (str, optional): The URL for the endpoint for the project on the crunch hosted site.
35 If not provided then it attempts to use the 'CRUNCH_URL' environment variable.
36 token (str, optional): An access token for a user on the crunch hosted site.
37 If not provided then it attempts to use the 'CRUNCH_TOKEN' environment variable.
39 Raises:
40 CrunchAPIException: If the `base_url` is not provided and it is not available using the 'CRUNCH_URL' environment variable.
41 CrunchAPIException: If the `token` is not provided and it is not available using the 'CRUNCH_TOKEN' environment variable.
42 """
43 self.base_url = base_url or getenv(CRUNCH_URL_KEY, None)
44 if not self.base_url:
45 raise CrunchAPIException(f"Please provide a base URL to a crunch hosted site. This can be set using the '{CRUNCH_URL_KEY}' environment variable.")
47 self.token = token or getenv(CRUNCH_TOKEN_KEY, None)
48 if not self.token:
49 raise CrunchAPIException(f"Please provide an authentication token. This can be set using the '{CRUNCH_TOKEN_KEY}' environment variable.")
51 self.verbose = verbose
53 def get_headers(self) -> dict:
54 """
55 Creates the headers needed to API calls to the REST API on a crunch hosted site.
57 Used internally when making GET and POST requests using this class.
59 Raises:
60 CrunchAPIException: Raised if no valid token is available.
62 Returns:
63 dict: The headers for API calls as a Python dictionary.
64 """
65 headers = {"Authorization": f"Token {self.token}" }
66 return headers
68 def absolute_url(self, relative_url:str) -> str:
69 if self.base_url.endswith("/"):
70 self.base_url = self.base_url[:-1]
71 if relative_url.startswith("/"):
72 relative_url = relative_url[1:]
74 url = f"{self.base_url}/{relative_url}"
75 return url
77 def post(self, relative_url:str, **kwargs) -> requests.Response:
78 url = self.absolute_url(relative_url)
79 if self.verbose:
80 console.print(f"Posting to {url} : {kwargs}")
82 result = requests.post(url, headers=self.get_headers(), data=kwargs)
83 if self.verbose:
84 console.print(f"Response {result.status_code}: {result.reason}")
86 if result.status_code >= 400:
87 console.print(f"Failed posting to {url} : {kwargs}")
88 console.print(result.json())
89 return result
91 def add_project(self, project:str, description:str="", details:str="") -> requests.Response:
92 """
93 Creates a new project on a hosted django-crunch site.
95 Args:
96 project (str): The name of the new crunch project.
97 description (str, optional): A brief description of this new project. Defaults to "".
98 details (str, optional): A long description of this project in Markdown format. Defaults to "".
100 Returns:
101 requests.Response: The request object from posting to the crunch API.
102 """
103 if self.verbose:
104 console.print(f"Adding project '{project}' on the site {self.base_url}")
106 return self.post(
107 "api/projects/",
108 name=project,
109 description=description,
110 details=details,
111 )
113 def add_dataset(self, project:str, dataset:str, description:str="", details:str="") -> requests.Response:
114 """
115 Creates a new dataset and adds it to a project on a hosted django-crunch site.
117 Args:
118 project (str): The slug of the project that this dataset is to be added to.
119 dataset (str): The name of the new dataset.
120 description (str, optional): A brief description of this new dataset. Defaults to "".
121 details (str, optional): A long description of this dataset in Markdown format. Defaults to "".
123 Returns:
124 requests.Response: The request object from posting to the crunch API.
125 """
126 if self.verbose:
127 console.print(f"Adding dataset '{dataset}' to project '{project}' on the site {self.base_url}")
129 return self.post(
130 "api/datasets/",
131 parent=project,
132 name=dataset,
133 description=description,
134 details=details,
135 )
137 def add_item(self, parent:str, item:str, description:str="", details:str="") -> requests.Response:
138 """
139 Creates a new item on a hosted django-crunch site.
141 Args:
142 parent (str): The slug of the parent item that this item is to be added to.
143 item (str): The name of the new item.
144 description (str, optional): A brief description of this new dataset. Defaults to "".
145 details (str, optional): A long description of this dataset in Markdown format. Defaults to "".
147 Returns:
148 requests.Response: The request object from posting to the crunch API.
149 """
150 if self.verbose:
151 console.print(f"Adding item '{item}' to parent '{parent}' on the site {self.base_url}")
153 return self.post(
154 "api/items/",
155 parent=parent,
156 name=item,
157 description=description,
158 details=details,
159 )
161 def add_key_value_attribute(self, url:str, item:str, key:str, value) -> requests.Response:
162 """
163 Adds an attribute as a key/value pair on an item.
165 This is mainly used by other methods on this class to add attributes with specific types.
167 Args:
168 url (str): The relative URL for adding this type of attribute on the crunch site. For this, see urls.py in the crunch Django app.
169 item (str): The slug for the item.
170 key (str): The key for this attribute.
171 value: The data to be used for this attribute. The object needs to be serializable.
173 Returns:
174 requests.Response: The request object from posting to the crunch API.
175 """
176 if self.verbose:
177 console.print(f"Adding attribute '{key}'->'{value}' to item '{item}' on the hosted site {self.base_url}")
179 return self.post(
180 url,
181 item=item,
182 key=key,
183 value=value,
184 )
186 def add_char_attribute(self, item:str, key:str, value:str) -> requests.Response:
187 """
188 Adds an attribute as a key/value pair on a dataset when the value is a string of characters.
190 Args:
191 project (str): The slug for the project.
192 dataset (str): The slug for the dataset.
193 key (str): The key for this attribute.
194 value (str): The string of characters for this attribute.
196 Returns:
197 requests.Response: The request object from posting to the crunch API.
198 """
199 return self.add_key_value_attribute(
200 url="api/attributes/char/",
201 item=item,
202 key=key,
203 value=value,
204 )
206 def add_attributes(self, item:str, **kwargs):
207 """
208 Adds multiple attributes as a key/value pairs on a dataset.
210 Each type is inferred from the type of the value.
212 Args:
213 item (str): The slug for the item.
214 **kwargs: key/value pairs to add as char attributes.
215 """
216 url_validator = URLValidator()
217 for key, value in kwargs.items():
218 if isinstance(value, str):
219 try:
220 url_validator(value)
221 self.add_url_attribute(item=item, key=key, value=value)
222 except ValidationError:
223 self.add_char_attribute(item=item, key=key, value=value)
224 elif isinstance(value, float):
225 self.add_float_attribute(item=item, key=key, value=value)
226 elif isinstance(value, bool): # This needs to be before the 'int' one because bools are also integers in python
227 self.add_boolean_attribute(item=item, key=key, value=value)
228 elif isinstance(value, int):
229 self.add_integer_attribute(item=item, key=key, value=value)
230 elif isinstance(value, datetime):
231 self.add_datetime_attribute(item=item, key=key, value=value)
232 elif isinstance(value, date):
233 self.add_date_attribute(item=item, key=key, value=value)
234 else:
235 raise CrunchAPIException(f"Cannot infer type of value '{value}' ({type(value).__name__}). (The key was '{key}')")
237 def add_float_attribute(self, item:str, key:str, value:float) -> requests.Response:
238 """
239 Adds an attribute as a key/value pair on a dataset when the value is a float.
241 Args:
242 item (str): The slug for the item.
243 key (str): The key for this attribute.
244 value (str): The float value for this attribute.
246 Returns:
247 requests.Response: The request object from posting to the crunch API.
248 """
249 return self.add_key_value_attribute(
250 url="api/attributes/float/",
251 item=item,
252 key=key,
253 value=value,
254 )
256 def add_datetime_attribute(self, item:str, key:str, value:Union[datetime,str], format:str="") -> requests.Response:
257 """
258 Adds an attribute as a key/value pair on a dataset when the value is a datetime.
260 Args:
261 item (str): The slug for the item.
262 key (str): The key for this attribute.
263 value (Union[datetime,str]): The value for this attribute as a datetime or a string.
264 format (str): If the `value` is a string then this format string can be used with datetime.strptime to convert to a datetime object. If no format is given then the string is interpreted using dateutil.parser.
266 Returns:
267 requests.Response: The request object from posting to the crunch API.
268 """
269 if isinstance(value, str):
270 if format:
271 value = datetime.strptime(value, format)
272 else:
273 from dateutil import parser
274 value = parser.parse(value)
276 return self.add_key_value_attribute(
277 url="api/attributes/datetime/",
278 item=item,
279 key=key,
280 value=value,
281 )
283 def add_date_attribute(self, item:str, key:str, value:Union[date,str], format:str="") -> requests.Response:
284 """
285 Adds an attribute as a key/value pair on a dataset when the value is a datetime.
287 Args:
288 item (str): The slug for the item.
289 key (str): The key for this attribute.
290 value (Union[date,str]): The value for this attribute as a date or a string.
291 format (str): If the `value` is a string then this format string can be used with datetime.strptime to convert to a date object. If no format is given then the string is interpreted using dateutil.parser.
293 Returns:
294 requests.Response: The request object from posting to the crunch API.
295 """
296 if isinstance(value, str):
297 if format:
298 value = datetime.strptime(value, format)
299 else:
300 from dateutil import parser
301 value = parser.parse(value)
303 if isinstance(value, datetime):
304 value = value.date()
306 return self.add_key_value_attribute(
307 url="api/attributes/date/",
308 item=item,
309 key=key,
310 value=value,
311 )
313 def add_integer_attribute(self, item:str, key:str, value:int) -> requests.Response:
314 """
315 Adds an attribute as a key/value pair on a dataset when the value is an integer.
317 Args:
318 item (str): The slug for the item.
319 key (str): The key for this attribute.
320 value (int): The integer value for this attribute.
322 Returns:
323 requests.Response: The request object from posting to the crunch API.
324 """
325 return self.add_key_value_attribute(
326 url="api/attributes/int/",
327 item=item,
328 key=key,
329 value=value,
330 )
332 def add_filesize_attribute(self, item:str, key:str, value:int) -> requests.Response:
333 """
334 Adds an attribute as a key/value pair on a dataset when the value is an filesize.
336 Args:
337 item (str): The slug for the item.
338 key (str): The key for this attribute.
339 value (int): The size of the file in bytes.
341 Returns:
342 requests.Response: The request object from posting to the crunch API.
343 """
344 return self.add_key_value_attribute(
345 url="api/attributes/filesize/",
346 item=item,
347 key=key,
348 value=value,
349 )
351 def add_boolean_attribute(self, item:str, key:str, value:bool) -> requests.Response:
352 """
353 Adds an attribute as a key/value pair on a dataset when the value is a boolean.
355 Args:
356 item (str): The slug for the item.
357 key (str): The key for this attribute.
358 value (bool): The integer value for this attribute.
360 Returns:
361 requests.Response: The request object from posting to the crunch API.
362 """
363 return self.add_key_value_attribute(
364 url="api/attributes/bool/",
365 item=item,
366 key=key,
367 value=value,
368 )
370 def add_url_attribute(self, item:str, key:str, value:str) -> requests.Response:
371 """
372 Adds an attribute as a key/value pair on a dataset when the value is a URL.
374 Args:
375 item (str): The slug for the item.
376 key (str): The key for this attribute.
377 value (str): The str value for this attribute.
379 Returns:
380 requests.Response: The request object from posting to the crunch API.
381 """
382 return self.add_key_value_attribute(
383 url="api/attributes/url/",
384 item=item,
385 key=key,
386 value=value,
387 )
389 def add_lat_long_attribute(self, item:str, key:str, latitude:Union[str,float,decimal], longitude:Union[str,float,decimal]) -> requests.Response:
390 """
391 Adds an attribute as a key/value pair on a dataset when the value is a coordinate with latitude and longitude.
393 Args:
394 item (str): The slug for the item.
395 key (str): The key for this attribute.
396 latitude (Union[str,float,decimal]): The latitude for this coordinate.
397 longitude (Union[str,float,decimal]): The longitude for this coordinate.
399 Returns:
400 requests.Response: The request object from posting to the crunch API.
401 """
402 if self.verbose:
403 console.print(f"Adding attribute '{key}'->'{latitude},{longitude}' to item '{item}' on the hosted site {self.base_url}")
405 return self.post(
406 relative_url="api/attributes/lat-long/",
407 item=item,
408 key=key,
409 latitude=latitude,
410 longitude=longitude,
411 )
413 def send_status(self, dataset_id:str, stage:Stage, state:State, note:str="") -> requests.Response:
414 """
415 Sends an update of the status of one stage in processing a dataset.
417 Args:
418 dataset_id (str): The ID of the dataset for this status update.
419 stage (Stage): The stage of this status update.
420 state (State): The state of this status update.
421 note (str, optional): A note which gives more information to this status update. Defaults to "".
423 Raises:
424 CrunchAPIException: If there was an error posting this status update to the API.
426 Returns:
427 requests.Response: The resulting response from the request to the API.
428 """
429 data = dict(
430 dataset=dataset_id,
431 stage=stage.value,
432 state=state.value,
433 note=note,
434 )
435 data.update( diagnostics.get_diagnostics() )
436 result = self.post("api/statuses/", **data)
437 if result.status_code >= 400:
438 raise CrunchAPIException(f"Failed sending status.\n{result.status_code}: {result.reason}\nData: {data}")
440 return result
442 def get_request(self, relative_url:str):
443 url = self.absolute_url(relative_url)
444 return requests.get(url, headers=self.get_headers())
446 def get_json_response( self, relative_url:str ) -> Dict:
447 """
448 Requests JSON data from the API of a crunch hosted site and returns it as a dictionary.
450 Args:
451 relative_url (str): The URL path relative to the base URL of the endpoint for the project on the crunch hosted site.
453 Raises:
454 CrunchAPIException: Raises exception if there is an error getting a JSON response from the API.
456 Returns:
457 Dict: The JSON data from the API encoded as a dictionary.
458 """
459 response = self.get_request(relative_url)
460 json_response = response.json()
462 if len(json_response.keys()) == 1 and "detail" in json_response:
463 raise CrunchAPIException(
464 f"Error getting JSON response from URL '{self.absolute_url(relative_url)}':\n" +
465 f"{json_response['detail']}"
466 )
468 return json_response