PYROS-IRAP / pyros

Commit 0e62158466e57c5c8fec8af8a68479c4f97a8e82

Authored by Alexis Koralewski
1 parent cd21eb38
Exists in dev

adding host option in pyros_api.py and a better handling of requests exception

Showing 1 changed file with 58 additions and 27 deletions   Show diff stats
Inline Side-by-side
pyros_api.py
  Diff comments   View file @0e62158
1 1 #!/usr/bin/env python3
2   -from email import header
3 2 import os
  3 +from urllib3.exceptions import ProtocolError
  4 +from requests.exceptions import ConnectionError
4 5 import requests
5 6 import sys
6 7 import yaml
... ... @@ -12,10 +13,9 @@ class PyrosAPI:
12 13 """
13 14 Request Pyros API with an PyrosUser account linked to the username and password
14 15 """
15   - BASE_PYROS_URL = "http://localhost:8000/api/"
16   - #BASE_PYROS_URL = "http://pyros.irap.omp.eu/api/"
  16 + #host = "http://pyros.irap.omp.eu/api/"
17 17  
18   - def __init__(self) -> None:
  18 + def __init__(self, host) -> None:
19 19 """
20 20 Initialize PyrosAPI class by getting the authentification token required to use the API
21 21  
... ... @@ -23,6 +23,11 @@ class PyrosAPI:
23 23 username : Username of the user (mail adress)
24 24 password : Password of the user
25 25 """
  26 + if host == None:
  27 + self.host = "http://localhost:8000"
  28 + else:
  29 + self.host = host
  30 + self.host = self.host+"/api/"
26 31 self.get_token()
27 32  
28 33 def get_token(self):
... ... @@ -38,23 +43,27 @@ class PyrosAPI:
38 43 self.token = token_file.read()
39 44 return
40 45 else:
41   - url = f"{self.BASE_PYROS_URL}api-token-auth/"
  46 + url = f"{self.host}api-token-auth/"
42 47 print("Token needs to be generated, you need to authentificate to PyROS :")
43 48 username = input("Enter your username: ")
44   -
45 49 try:
46 50 password = getpass.getpass()
47 51 except Exception as error:
48 52 print('ERROR', error)
49   - request = requests.post(
50   - url, data={"username": username, "password": password})
51   - if request.status_code != 200:
52   - print("Authentification failed, please retry")
53   - exit(1)
54   - json_response = request.json()
55   - self.token = json_response["token"]
56   - with open("TOKEN", "w+") as token_file:
57   - token_file.write(self.token)
  53 + try:
  54 + request = requests.post(
  55 + url, data={"username": username, "password": password})
  56 + except (ConnectionError, ConnectionResetError, ProtocolError):
  57 + print(
  58 + f"Server {self.host} doesn't respond. The website might be down.")
  59 + exit(1)
  60 + if request.status_code != 200:
  61 + print("Authentification failed, please retry")
  62 + exit(1)
  63 + json_response = request.json()
  64 + self.token = json_response["token"]
  65 + with open("TOKEN", "w+") as token_file:
  66 + token_file.write(self.token)
58 67  
59 68 def get_url(self, url: str) -> dict:
60 69 """
... ... @@ -67,8 +76,12 @@ class PyrosAPI:
67 76 dict : Json object that represents the response of the API
68 77 """
69 78 headers = {"Authorization": f"Token {self.token}"}
70   -
71   - request = requests.get(self.BASE_PYROS_URL+url, headers=headers)
  79 + try:
  80 + request = requests.get(self.host+url, headers=headers)
  81 + except (ConnectionError, ConnectionResetError, ProtocolError):
  82 + print(
  83 + f"Server {self.host} doesn't respond. The website might be down.")
  84 + exit(1)
72 85 return request.json()
73 86  
74 87 def submit_sequence_file(self, file: str) -> dict:
... ... @@ -85,12 +98,17 @@ class PyrosAPI:
85 98 "Authorization": f"Token {self.token}",
86 99 'Content-type': 'application/json',
87 100 }
88   - url = f"{self.BASE_PYROS_URL}submit_sequence"
  101 + url = f"{self.host}submit_sequence"
89 102 yaml_file = open(file, "r")
90 103 sequence_file_as_dict = yaml.safe_load(yaml_file)
91 104 print(f"File content : {sequence_file_as_dict}")
92   - request = requests.put(url, headers=headers,
93   - json=sequence_file_as_dict)
  105 + try:
  106 + request = requests.put(url, headers=headers,
  107 + json=sequence_file_as_dict)
  108 + except (ConnectionError, ConnectionResetError, ProtocolError):
  109 + print(
  110 + f"Server {self.host} doesn't respond. The website might be down.")
  111 + exit(1)
94 112 print(f"Request status code {request.status_code}")
95 113 return request.json()
96 114  
... ... @@ -103,23 +121,36 @@ class PyrosAPI:
103 121 Returns:
104 122 dict : Json object that represents the response of the API
105 123 """
106   - url = f"{self.BASE_PYROS_URL}full_sequences/get_sequences_for_period/"
  124 + url = f"{self.host}full_sequences/get_sequences_for_period/"
107 125 headers = {"Authorization": f"Token {self.token}"}
108   - response = requests.get(
109   - url, params={"start_date": start_date, "end_date": end_date}, headers=headers)
  126 + try:
  127 + response = requests.get(
  128 + url, params={"start_date": start_date, "end_date": end_date}, headers=headers)
  129 + except (ConnectionError, ConnectionResetError, ProtocolError):
  130 + print(
  131 + f"Server {self.host} doesn't respond. The website might be down.")
  132 + exit(1)
110 133 return response.json()
111 134  
112 135 def logout(self):
113   - url = f"{self.BASE_PYROS_URL}logout/"
  136 + url = f"{self.host}logout/"
114 137 header = {"Authorization": f"Token {self.token}"}
115   - response = requests.get(url, headers=header)
  138 + try:
  139 + response = requests.get(url, headers=header)
  140 + except (ConnectionError, ConnectionResetError, ProtocolError):
  141 + print(
  142 + f"Server {self.host} doesn't respond. The website might be down.")
  143 + exit(1)
116 144 return response.content.decode("utf-8")
117 145  
118 146  
119 147 @click.group()
  148 +@click.option('--host', '-h', help='host name (example: http://localhost:8000 or http://pyros.omp.eu) without last slash')
120 149 @click.pass_context
121   -def cli(ctx):
122   - ctx.obj = PyrosAPI()
  150 +def cli(ctx, host):
  151 + if not host:
  152 + host = None
  153 + ctx.obj = PyrosAPI(host)
123 154  
124 155  
125 156 @cli.command("get_sequences_for_period", help="Get sequences for a period of date (start/end)")
... ...