diff --git a/deeplake/client/client.py b/deeplake/client/client.py index be57443be5..f7d36e89d0 100644 --- a/deeplake/client/client.py +++ b/deeplake/client/client.py @@ -40,8 +40,6 @@ from deeplake.client.auth import initialize_auth_context import jwt # should add it to requirements.txt -from deeplake.util.storage import get_dataset_credentials - # for these codes, we will retry requests upto 3 times retry_status_codes = {502} @@ -128,19 +126,29 @@ def request( status_code = None tries = 0 + last_exception = None while status_code is None or (status_code in retry_status_codes and tries < 3): - response = requests.request( - method, - request_url, - params=params, - data=data, - json=json, - headers=headers, - files=files, - timeout=timeout, - ) + last_exception = None + try: + response = requests.request( + method, + request_url, + params=params, + data=data, + json=json, + headers=headers, + files=files, + timeout=timeout, + ) + except requests.exceptions.ConnectionError as e: + tries += 1 + last_exception = e + continue status_code = response.status_code tries += 1 + if last_exception: + raise last_exception + check_response_status(response) return response @@ -186,59 +194,48 @@ def get_dataset_credentials( """ import json - last_try_error = None - for try_num in range(0, 5): - db_engine = db_engine or {} - relative_url = GET_DATASET_CREDENTIALS_SUFFIX.format(org_id, ds_name) - try: - response = self.request( - "GET", - relative_url, - endpoint=self.endpoint(), - params={ - "mode": mode, - "no_cache": no_cache, - "db_engine": json.dumps(db_engine), - }, - ).json() - except Exception as e: - if isinstance(e, AuthorizationException): - code = -1 - if e.response is not None: - response_data = e.response.json() - code = response_data.get("code") - - if code == 1: - agreements = response_data["agreements"] - agreements = [agreement["text"] for agreement in agreements] - raise AgreementNotAcceptedError(agreements) from e - elif code == 2: - raise NotLoggedInAgreementError from e - else: - try: - jwt.decode( - self.get_token(), options={"verify_signature": False} - ) - except Exception: - raise InvalidTokenException - - raise TokenPermissionError(e.original_message) - if isinstance(e, requests.exceptions.ConnectionError): - print("Connection error getting credentials. Retrying") - last_try_error = e - continue - raise - full_url = response.get("path") - repository = response.get("repository") - creds = response["creds"] - mode = response["mode"] - expiration = creds["expiration"] if creds else None - return full_url, creds, mode, expiration, repository - - if last_try_error: - raise last_try_error - else: - raise "Unexpected error fetching credentials" + db_engine = db_engine or {} + relative_url = GET_DATASET_CREDENTIALS_SUFFIX.format(org_id, ds_name) + try: + response = self.request( + "GET", + relative_url, + endpoint=self.endpoint(), + params={ + "mode": mode, + "no_cache": no_cache, + "db_engine": json.dumps(db_engine), + }, + ).json() + except Exception as e: + if isinstance(e, AuthorizationException): + code = -1 + if e.response is not None: + response_data = e.response.json() + code = response_data.get("code") + + if code == 1: + agreements = response_data["agreements"] + agreements = [agreement["text"] for agreement in agreements] + raise AgreementNotAcceptedError(agreements) from e + elif code == 2: + raise NotLoggedInAgreementError from e + else: + try: + jwt.decode( + self.get_token(), options={"verify_signature": False} + ) + except Exception: + raise InvalidTokenException + + raise TokenPermissionError(e.original_message) + raise + full_url = response.get("path") + repository = response.get("repository") + creds = response["creds"] + mode = response["mode"] + expiration = creds["expiration"] if creds else None + return full_url, creds, mode, expiration, repository def send_event(self, event_json: dict): """Sends an event to the backend.