import os
from typing import Union
from urllib.parse import urljoin
import requests
from requests.adapters import HTTPAdapter
from requests.packages.urllib3.util.retry import Retry
import openprotein.config as config
from openprotein.errors import APIError, AuthError, HTTPError
BACKEND = os.getenv("OPENPROTEIN_API_BACKEND", "https://api.openprotein.ai/api/")
class BearerAuth(requests.auth.AuthBase):
"""
See https://stackoverflow.com/a/58055668
"""
def __init__(self, token):
self.token = token
def __call__(self, r):
r.headers["authorization"] = "Bearer " + self.token
return r
[docs]
class APISession(requests.Session):
"""
A class to handle API sessions. This class provides a connection session to the OpenProtein API.
Parameters
----------
username : str
The username of the user.
password : str
The password of the user.
Examples
--------
>>> session = APISession("username", "password")
"""
[docs]
def __init__(
self,
username: str,
password: str,
backend: str = BACKEND,
timeout: int = 180,
):
super().__init__()
self.backend = backend
self.verify = True
self.timeout = timeout
# Custom retry strategies
# auto retry for pesky connection reset errors and others
# 503 will catch if BE is refreshing
retry = Retry(
total=4,
backoff_factor=3, # 0,1,4,13s
status_forcelist=[500, 502, 503, 504, 101, 104],
)
adapter = HTTPAdapter(max_retries=retry)
self.mount("https://", adapter)
self.login(username, password)
[docs]
def post(self, url, data=None, json=None, **kwargs):
r"""Sends a POST request. Returns :class:`Response` object.
:param url: URL for the new :class:`Request` object.
:param data: (optional) Dictionary, list of tuples, bytes, or file-like
object to send in the body of the :class:`Request`.
:param json: (optional) json to send in the body of the :class:`Request`.
:param \*\*kwargs: Optional arguments that ``request`` takes.
:rtype: requests.Response
"""
timeout = self.timeout
if "timeout" in kwargs:
timeout = kwargs.pop("timeout")
return self.request("POST", url, data=data, json=json, timeout=timeout, **kwargs)
[docs]
def login(self, username: str, password: str):
"""
Authenticate connection to OpenProtein with your credentials.
Parameters
-----------
username: str
username
password: str
password
"""
self.auth = self._get_auth_token(username, password)
def _get_auth_token(self, username: str, password: str):
endpoint = "v1/login/access-token"
url = urljoin(self.backend, endpoint)
try:
response = self.post(url, data={"username": username, "password": password}, timeout=3)
except HTTPError as e:
# if an error occured during auth, we raise an AuthError with reference to the HTTPError
raise AuthError(
f"Authentication failed. Please check your credentials and connection."
) from e
result = response.json()
token = result.get("access_token")
if token is None:
raise AuthError("Unable to authenticate with given credentials.")
return BearerAuth(token)
def request(self, method: Union[str, bytes], url: Union[str, bytes], *args, **kwargs):
full_url = urljoin(self.backend, url)
response = super().request(method, full_url, *args, **kwargs)
if not response.ok:
# raise custom exception that prints better error message than requests.HTTPError
raise HTTPError(response)
return response
class RestEndpoint:
pass