Allow passing access_token directly

lamindb_setup/dev/_hub_client.py

@@ -60,17 +60,21 @@ def connect_hub(
 
 
 def connect_hub_with_auth(
-    fallback_env: bool = False, renew_token: bool = False
+    fallback_env: bool = False,
+    renew_token: bool = False,
+    access_token: Optional[str] = None,
 ) -> Client:
-    from lamindb_setup import settings
-
     hub = connect_hub(fallback_env=fallback_env)
-    if renew_token:
-        settings.user.access_token = get_access_token(
-            settings.user.email, settings.user.password
-        )
-    hub.postgrest.auth(settings.user.access_token)
-    hub.functions.set_auth(settings.user.access_token)
+    if access_token is None:
+        from lamindb_setup import settings
+
+        if renew_token:
+            settings.user.access_token = get_access_token(
+                settings.user.email, settings.user.password
+            )
+        access_token = settings.user.access_token
+    hub.postgrest.auth(access_token)
+    hub.functions.set_auth(access_token)
     return hub
 
 
@@ -98,6 +102,19 @@ def call_with_fallback_auth(
     callable,
     **kwargs,
 ):
+    access_token = kwargs.pop("access_token", None)
+
+    if access_token is not None:
+        try:
+            client = connect_hub_with_auth(access_token=access_token)
+            result = callable(**kwargs, client=client)
+        finally:
+            try:
+                client.auth.sign_out()
+            except NameError:
+                pass
+        return result
+
     for renew_token, fallback_env in [(False, False), (True, False), (False, True)]:
         try:
             client = connect_hub_with_auth(
@@ -109,7 +126,10 @@ def call_with_fallback_auth(
             if fallback_env:
                 raise e
         finally:
-            client.auth.sign_out()
+            try:
+                client.auth.sign_out()
+            except NameError:
+                pass
     return result
 
 

lamindb_setup/dev/_hub_core.py

@@ -208,11 +208,11 @@ def _load_instance(
     return instance, storage
 
 
-def access_aws() -> Dict[str, str]:
+def access_aws(access_token: Optional[str] = None) -> Dict[str, str]:
     from .._settings import settings
 
-    if settings.user.handle != "anonymous":
-        credentials = call_with_fallback_auth(_access_aws)
+    if settings.user.handle != "anonymous" or access_token is not None:
+        credentials = call_with_fallback_auth(_access_aws, access_token=access_token)
         logger.important("loaded AWS credentials")
         return credentials
     else:

lamindb_setup/dev/_settings_storage.py

@@ -116,6 +116,7 @@ def __init__(
         region: Optional[str] = None,
         uid: Optional[str] = None,
         uuid: Optional[UUID] = None,
+        access_token: Optional[str] = None,
     ):
         self._uid = uid
         self._uuid = uuid
@@ -141,6 +142,8 @@ def __init__(
             self._cache_dir = _process_cache_path(cache_path)
         else:
             self._cache_dir = None
+        # save access_token here for use in self.root
+        self.access_token = access_token
 
     @property
     def id(self) -> int:
@@ -185,7 +188,7 @@ def root(self) -> UPath:
         if self._root is None:
             # below also makes network requests to get credentials
             # right
-            root_path = create_path(self._root_init)
+            root_path = create_path(self._root_init, access_token=self.access_token)
             self._root = root_path
         return self._root
 

lamindb_setup/dev/upath.py

@@ -432,7 +432,7 @@ def convert_pathlike(pathlike: Union[str, Path, UPath]) -> UPath:
 AWS_CREDENTIALS_PRESENT = None
 
 
-def create_path(path: UPath) -> UPath:
+def create_path(path: UPath, access_token: Optional[str] = None) -> UPath:
     path = convert_pathlike(path)
     # test whether we have an AWS S3 path
     if not isinstance(path, S3Path):
@@ -465,18 +465,21 @@ def create_path(path: UPath) -> UPath:
     if not is_hosted_storage:
         # make anon request if no credentials present
         return UPath(path, cache_regions=True, anon=anon)
+
+    root_folder = path_str.replace("s3://", "").split("/")[0]
+
+    if access_token is None and root_folder in credentials_cache:
+        credentials = credentials_cache[root_folder]
     else:
-        root_folder = path_str.replace("s3://", "").split("/")[0]
-        if root_folder in credentials_cache:
-            credentials = credentials_cache[root_folder]
-        else:
-            from ._hub_core import access_aws
+        from ._hub_core import access_aws
 
-            credentials = access_aws()
+        credentials = access_aws(access_token=access_token)
+        if access_token is None:
             credentials_cache[root_folder] = credentials
-        return UPath(
-            path,
-            key=credentials["key"],
-            secret=credentials["secret"],
-            token=credentials["token"],
-        )
+
+    return UPath(
+        path,
+        key=credentials["key"],
+        secret=credentials["secret"],
+        token=credentials["token"],
+    )