From 54d257562351042f53befa4d901b58a5a5ac21a5 Mon Sep 17 00:00:00 2001
From: dgaloop <david.gyulnazaryan@activeloop.dev>
Date: Sat, 3 Aug 2024 20:40:01 +0400
Subject: [PATCH] Ensure username fetching logic is not token dependent (#2923)

* Ensure username fetching logic is not token dependent
* Better handling of username from client, less errors when guessing username from an unvalidated token
---
 deeplake/client/client.py                     | 35 +++++++++++--------
 deeplake/core/dataset/dataset.py              |  2 +-
 .../core/dataset/deeplake_cloud_dataset.py    |  4 +++
 .../dataset_handlers/dataset_handler_base.py  |  3 ++
 4 files changed, 29 insertions(+), 15 deletions(-)

diff --git a/deeplake/client/client.py b/deeplake/client/client.py
index 47729bde39..5552ee6e21 100644
--- a/deeplake/client/client.py
+++ b/deeplake/client/client.py
@@ -1,6 +1,6 @@
 import deeplake
 import requests  # type: ignore
-from typing import Any, Optional, Dict
+from typing import Any, Optional, Dict, Tuple
 from deeplake.util.exceptions import (
     AgreementNotAcceptedError,
     AuthorizationException,
@@ -59,18 +59,25 @@ def __init__(self, token: Optional[str] = None):
 
         # remove public token, otherwise env var will be ignored
         # we can remove this after a while
-        orgs = self.get_user_organizations()
-        if orgs == ["public"]:
+        self._username, self._organizations = self._get_username_and_organizations()
+        if self._username == "public":
             self.token = token or self.get_token()
         else:
-            username = self.get_user_profile()["name"]
-            if get_reporting_config().get("username") != username:
-                save_reporting_config(True, username=username)
-                set_username(username)
+            if get_reporting_config().get("username") != self._username:
+                save_reporting_config(True, username=self._username)
+                set_username(self._username)
 
     def get_token(self):
         return self.auth_context.get_token()
 
+    @property
+    def username(self) -> str:
+        return self._username
+
+    @property
+    def organizations(self) -> list[str]:
+        return self._organizations
+
     def request(
         self,
         method: str,
@@ -335,25 +342,25 @@ def rename_dataset_entry(self, username, old_name, new_name):
             "PUT", suffix, endpoint=self.endpoint(), json={"basename": new_name}
         )
 
-    def get_user_organizations(self):
-        """Get list of user organizations from the backend. If user is not authenticated, returns ['public'].
+    def _get_username_and_organizations(self) -> Tuple[str, list[str]]:
+        """Get the username plus a list of user organizations from the backend. If user is not authenticated, returns ('public', ['public']).
 
         Returns:
-            list: user/organization names
+            (str, list[str]): user + organization namess
         """
+
         if self.auth_context.is_public_user():
-            return ["public"]
+            return "public", ["public"]
 
         response = self.request(
             "GET", GET_USER_PROFILE, endpoint=self.endpoint()
         ).json()
-        return response["organizations"]
+        return response["_id"], response["organizations"]
 
     def get_workspace_datasets(
         self, workspace: str, suffix_public: str, suffix_user: str
     ):
-        organizations = self.get_user_organizations()
-        if workspace in organizations:
+        if workspace in self.organizations:
             response = self.request(
                 "GET",
                 suffix_user,
diff --git a/deeplake/core/dataset/dataset.py b/deeplake/core/dataset/dataset.py
index 929c10cd84..5126e472fe 100644
--- a/deeplake/core/dataset/dataset.py
+++ b/deeplake/core/dataset/dataset.py
@@ -353,7 +353,7 @@ def username(self) -> str:
 
         try:
             return jwt.decode(self.token, options={"verify_signature": False})["id"]
-        except DecodeError:
+        except:
             return "public"
 
     @property
diff --git a/deeplake/core/dataset/deeplake_cloud_dataset.py b/deeplake/core/dataset/deeplake_cloud_dataset.py
index 5d187a009b..345325a209 100644
--- a/deeplake/core/dataset/deeplake_cloud_dataset.py
+++ b/deeplake/core/dataset/deeplake_cloud_dataset.py
@@ -46,6 +46,10 @@ def client(self):
             self.__dict__["_client"] = DeepLakeBackendClient(token=self._token)
         return self._client
 
+    @property
+    def username(self):
+        return self.client.username
+
     @property
     def is_actually_cloud(self) -> bool:
         """Datasets that are connected to Deep Lake cloud can still technically be stored anywhere.
diff --git a/deeplake/core/vectorstore/dataset_handlers/dataset_handler_base.py b/deeplake/core/vectorstore/dataset_handlers/dataset_handler_base.py
index c610eaa899..69b066ea57 100644
--- a/deeplake/core/vectorstore/dataset_handlers/dataset_handler_base.py
+++ b/deeplake/core/vectorstore/dataset_handlers/dataset_handler_base.py
@@ -101,6 +101,9 @@ def exec_option(self) -> str:
 
     @property
     def username(self) -> str:
+        if self.dataset:
+            return self.dataset.username
+
         username = "public"
         if self.token is not None:
             try: