perf(cli): resolve a teamspace for drive reads from the org and memberships only

python/lightning_sdk/api/org_api.py

@@ -6,6 +6,7 @@
     V1CreateProjectRequest,
     V1Organization,
 )
+from lightning_sdk.lightning_cloud.openapi.rest import ApiException
 
 # The billing summary API only supports querying up to 2 years of history.
 _MAX_DURATION = timedelta(days=730)
@@ -50,6 +51,26 @@ def get_org(self, name: str) -> V1Organization:
             raise ValueError(f"Org {name} does not exist")
         return res
 
+    def find_org_id(self, name: str) -> Optional[str]:
+        """Return the ID of the organisation with this name, or ``None`` if there is none.
+
+        Args:
+            name: The organisation name to look up.
+
+        Returns:
+            Optional[str]: The organisation's ID, or ``None`` if no organisation has that name.
+
+        Raises:
+            ApiException: If the API returns an error other than not found.
+        """
+        try:
+            res = self._client.organizations_service_get_organization(id="", name=name)
+        except ApiException as ex:
+            if ex.status == 404:
+                return None
+            raise
+        return res.id if res else None
+
     def _get_org_by_id(self, org_id: str) -> V1Organization:
         """Gets the organization from the given ID.
 

python/lightning_sdk/api/teamspace_api.py

@@ -178,6 +178,33 @@ def list_teamspaces(self, owner_id: str, name: Optional[str] = None) -> Optional
                 teamspaces.append(self._get_teamspace_by_id(teamspace.project_id))
         return teamspaces
 
+    def find_teamspace_id(self, name: str, owner_id: str) -> Optional[str]:
+        """Return the ID of the authed user's teamspace with this name and owner.
+
+        Unlike ``get_teamspace``, reads only the user's memberships, without
+        fetching the teamspace itself.
+
+        Args:
+            name: Name or display name of the teamspace.
+            owner_id: User or organization ID that owns the teamspace.
+
+        Returns:
+            Optional[str]: The teamspace's ID, or ``None`` if the user is a member of
+            no teamspace with that name and owner.
+
+        Raises:
+            RuntimeError: If the name matches more than one teamspace.
+        """
+        res = self._client.projects_service_list_memberships(filter_by_user_id=True)
+        ids = [
+            membership.project_id
+            for membership in res.memberships
+            if membership.owner_id == owner_id and name in (membership.name, membership.display_name)
+        ]
+        if len(ids) > 1:
+            raise RuntimeError(f"{name} is no unique name for a Teamspace")
+        return ids[0] if ids else None
+
     def list_studios(self, teamspace_id: str, cloud_account: str = "") -> List[V1CloudSpace]:
         """List studios in teamspace.
 

python/lightning_sdk/cli/utils/filesystem.py

@@ -1,9 +1,22 @@
+from dataclasses import dataclass
 from typing import Optional, Tuple
 
+from lightning_sdk.api.org_api import OrgApi
+from lightning_sdk.api.teamspace_api import TeamspaceApi
 from lightning_sdk.cli.utils.resource_resolution import join_teamspace_slug
 from lightning_sdk.cli.utils.resource_resolution import resolve_teamspace as resolve_cli_teamspace
 from lightning_sdk.teamspace import Teamspace
+from lightning_sdk.utils.config import Config, DefaultConfigKeys
 from lightning_sdk.utils.filesystem import parse_lit_url
+from lightning_sdk.utils.resolve import _resolve_org_name, _resolve_teamspace_name
+
+
+@dataclass(frozen=True)
+class TeamspaceRef:
+    """The parts of a teamspace that reading its drive needs."""
+
+    id: str
+    name: str
 
 
 def resolve_teamspace(teamspace: Optional[str], owner: Optional[str]) -> Teamspace:
@@ -19,3 +32,46 @@ def resolve_lit_url(url: str) -> Tuple[Teamspace, str]:
     """
     parsed = parse_lit_url(url)
     return resolve_teamspace(parsed["teamspace"], parsed["owner"]), parsed["destination"] or ""
+
+
+def resolve_teamspace_ref(teamspace: Optional[str], owner: Optional[str]) -> TeamspaceRef:
+    """Like :func:`resolve_teamspace`, for operations that need only the teamspace's ID.
+
+    Teamspaces are owned by organizations, so a teamspace resolves from two
+    requests, the organization and the user's memberships, instead of building a
+    full :class:`Teamspace`. A missing name or owner is taken from the current
+    teamspace's configuration, as :class:`Teamspace` does.
+
+    Raises:
+        ValueError: If no teamspace is given or configured, the organization does
+            not exist, or the user is not a member of a teamspace with that name in it.
+    """
+    teamspace = _resolve_teamspace_name(teamspace)
+    owner = owner or _current_org_name()
+    if not (teamspace and owner):
+        raise ValueError("Could not resolve a teamspace. Use lit://<org>/<teamspace>/<path>.")
+
+    org_id = OrgApi().find_org_id(owner)
+    if not org_id:
+        raise ValueError(f"Org {owner} does not exist")
+    teamspace_id = TeamspaceApi().find_teamspace_id(name=teamspace, owner_id=org_id)
+    if not teamspace_id:
+        raise ValueError(f"Teamspace {teamspace} does not exist")
+    return TeamspaceRef(id=teamspace_id, name=teamspace)
+
+
+def _current_org_name() -> Optional[str]:
+    """The organization of the current teamspace, from the environment or the config."""
+    org = _resolve_org_name(None)
+    if org:
+        return org
+    config = Config()
+    if (config.get_value(DefaultConfigKeys.teamspace_owner_type) or "").lower() == "organization":
+        return config.get_value(DefaultConfigKeys.teamspace_owner)
+    return None
+
+
+def resolve_lit_url_ref(url: str) -> Tuple[TeamspaceRef, str]:
+    """Like :func:`resolve_lit_url`, resolving the teamspace with :func:`resolve_teamspace_ref`."""
+    parsed = parse_lit_url(url)
+    return resolve_teamspace_ref(parsed["teamspace"], parsed["owner"]), parsed["destination"] or ""

python/lightning_sdk/filesystem.py

@@ -6,7 +6,7 @@
 
 from lightning_sdk.api.filesystem_api import FilesystemApi
 from lightning_sdk.api.utils import _RemoteApiError, _tree_path_info
-from lightning_sdk.cli.utils.filesystem import resolve_lit_url
+from lightning_sdk.cli.utils.filesystem import resolve_lit_url, resolve_lit_url_ref
 from lightning_sdk.teamspace import Teamspace
 from lightning_sdk.utils.logging import TrackCallsMeta
 
@@ -35,8 +35,8 @@ def listdir(self, uri: str) -> List[str]:
         Returns:
             List[str]: Basenames of the entries directly inside the given directory.
         """
-        selected_teamspace, remote_path = resolve_lit_url(uri)
-        output = self._filesystem_api.list_files(teamspace_id=selected_teamspace.id, path=remote_path, recursive=False)
+        teamspace_ref, remote_path = resolve_lit_url_ref(uri)
+        output = self._filesystem_api.list_files(teamspace_id=teamspace_ref.id, path=remote_path, recursive=False)
         return [os.path.basename(item["path"]) for item in output]
 
     def walk(self, url: str) -> Generator[Tuple[str, List[str], List[str]], None, None]:
@@ -51,8 +51,8 @@ def walk(self, url: str) -> Generator[Tuple[str, List[str], List[str]], None, No
             current directory path, a list of its immediate subdirectory names, and a list of
             its immediate file names — mirroring the behaviour of :func:`os.walk`.
         """
-        selected_teamspace, remote_path = resolve_lit_url(url)
-        output = self._filesystem_api.list_files(teamspace_id=selected_teamspace.id, path=remote_path, recursive=True)
+        teamspace_ref, remote_path = resolve_lit_url_ref(url)
+        output = self._filesystem_api.list_files(teamspace_id=teamspace_ref.id, path=remote_path, recursive=True)
 
         dirs: dict[str, list[str]] = {}
         files: dict[str, list[str]] = {}
@@ -148,18 +148,18 @@ def copy(
         if not source_is_lit and not dest_is_lit:
             raise ValueError("At least one path must be a lit://")
 
-        selected_teamspace, remote_path = resolve_lit_url(source if source_is_lit else destination)
         local_path = destination if source_is_lit else source
         if source_is_lit:
             # download
+            teamspace_ref, remote_path = resolve_lit_url_ref(source)
             parent = os.path.dirname(remote_path.strip("/"))
             try:
-                entries = self._filesystem_api.list_files(selected_teamspace.id, parent, recursive=False)
+                entries = self._filesystem_api.list_files(teamspace_ref.id, parent, recursive=False)
             except _RemoteApiError as e:
                 # A missing parent means the path can't exist either.
                 if e.status_code != 404:
                     raise
-                raise ValueError(f"File {remote_path} does not exist in teamspace {selected_teamspace.name}") from e
+                raise ValueError(f"File {remote_path} does not exist in teamspace {teamspace_ref.name}") from e
             found = False
             is_directory = False
 
@@ -170,7 +170,7 @@ def copy(
                     break
 
             if not found:
-                raise ValueError(f"File {remote_path} does not exist in teamspace {selected_teamspace.name}")
+                raise ValueError(f"File {remote_path} does not exist in teamspace {teamspace_ref.name}")
 
             if is_directory:
                 if not recursive:
@@ -180,21 +180,22 @@ def copy(
                 local_folder_name = os.path.basename(remote_path.rstrip("/"))
                 if local_path in ("./", "."):
                     if local_folder_name == "":
-                        local_folder_name = f"{selected_teamspace.name}_downloads"
+                        local_folder_name = f"{teamspace_ref.name}_downloads"
                     target_path = os.path.join(local_path, local_folder_name)
                 else:
                     target_path = local_path
-                self._filesystem_api.download_folder(remote_path, target_path, selected_teamspace.id, progress_bar)
+                self._filesystem_api.download_folder(remote_path, target_path, teamspace_ref.id, progress_bar)
             else:
                 if os.path.isdir(local_path) or local_path.endswith(("/", "\\")):
                     # if local_path ends with / or \ or is a directory, treat it as a directory
                     file_name = os.path.basename(remote_path)
                     target_path = os.path.join(local_path, file_name)
                 else:
                     target_path = local_path
-                self._filesystem_api.download_file(remote_path, target_path, selected_teamspace.id, progress_bar)
+                self._filesystem_api.download_file(remote_path, target_path, teamspace_ref.id, progress_bar)
         else:
             # upload
+            selected_teamspace, remote_path = resolve_lit_url(destination)
             if not os.path.exists(local_path):
                 raise FileNotFoundError(f"The provided path does not exist: {local_path}")
 

python/tests/core/filesystem/test_filesystem_copy.py

@@ -32,7 +32,7 @@ def fake_path_result():
 @mock.patch("lightning_sdk.api.filesystem_api.requests.get")
 @mock.patch("lightning_sdk.api.utils.LightningClient")
 @mock.patch("lightning_sdk.api.filesystem_api._authenticate_and_get_auth_headers")
-@mock.patch("lightning_sdk.cli.utils.filesystem.resolve_teamspace")
+@mock.patch("lightning_sdk.cli.utils.filesystem.resolve_teamspace_ref")
 @mock.patch("lightning_sdk.cli.utils.filesystem.parse_lit_url")
 def test_copy_download_file(
     mock_parse_lit_url, mock_resolve, mock_authenticate, mock_client_cls, mock_get, fake_teamspace, tmp_path
@@ -74,7 +74,7 @@ def fake_get(url, **kwargs):
 @mock.patch("lightning_sdk.api.filesystem_api.requests.get")
 @mock.patch("lightning_sdk.api.utils.LightningClient")
 @mock.patch("lightning_sdk.api.filesystem_api._authenticate_and_get_auth_headers")
-@mock.patch("lightning_sdk.cli.utils.filesystem.resolve_teamspace")
+@mock.patch("lightning_sdk.cli.utils.filesystem.resolve_teamspace_ref")
 @mock.patch("lightning_sdk.cli.utils.filesystem.parse_lit_url")
 def test_copy_download_folder(
     mock_parse_lit_url,
@@ -121,7 +121,7 @@ def fake_get(url, **kwargs):
 @mock.patch("lightning_sdk.api.filesystem_api.requests.get")
 @mock.patch("lightning_sdk.api.utils.LightningClient")
 @mock.patch("lightning_sdk.api.filesystem_api._authenticate_and_get_auth_headers", return_value=FAKE_AUTH_HEADERS)
-@mock.patch("lightning_sdk.cli.utils.filesystem.resolve_teamspace")
+@mock.patch("lightning_sdk.cli.utils.filesystem.resolve_teamspace_ref")
 @mock.patch("lightning_sdk.cli.utils.filesystem.parse_lit_url")
 def test_copy_raises_if_directory_without_recursive(
     mock_parse_lit_url, mock_resolve, _mock_authenticate, mock_client_cls, mock_get, fake_teamspace
@@ -420,5 +420,9 @@ def test_copy_download_maps_missing_parent_to_does_not_exist(fake_teamspace):
     fs, patches = _upload_fs(fake_teamspace, "r2_connections/nope/file.txt")
     fs._filesystem_api.list_files.side_effect = _RemoteApiError("Failed to list files: 404", status_code=404)
 
-    with patches[0], patches[1], pytest.raises(ValueError, match="does not exist in teamspace"):
+    with (
+        patches[0],
+        mock.patch("lightning_sdk.cli.utils.filesystem.resolve_teamspace_ref", return_value=fake_teamspace),
+        pytest.raises(ValueError, match="does not exist in teamspace"),
+    ):
         fs.copy("lit://my-org/my-teamspace/r2_connections/nope/file.txt", "/tmp/out")

python/tests/core/filesystem/test_filesystem_listdir.py

@@ -34,7 +34,7 @@ def fake_path_result():
     "lightning_sdk.api.filesystem_api._authenticate_and_get_auth_headers",
     new=mock.MagicMock(return_value=FAKE_AUTH_HEADERS),
 )
-@mock.patch("lightning_sdk.cli.utils.filesystem.resolve_teamspace")
+@mock.patch("lightning_sdk.cli.utils.filesystem.resolve_teamspace_ref")
 @mock.patch("lightning_sdk.cli.utils.filesystem.parse_lit_url")
 def test_listdir_returns_files(
     mock_parse_lit_url, mock_resolve, mock_client_cls, mock_get, fake_teamspace, fake_path_result
@@ -61,7 +61,7 @@ def test_listdir_returns_files(
     "lightning_sdk.api.filesystem_api._authenticate_and_get_auth_headers",
     new=mock.MagicMock(return_value=FAKE_AUTH_HEADERS),
 )
-@mock.patch("lightning_sdk.cli.utils.filesystem.resolve_teamspace")
+@mock.patch("lightning_sdk.cli.utils.filesystem.resolve_teamspace_ref")
 @mock.patch("lightning_sdk.cli.utils.filesystem.parse_lit_url")
 def test_listdir_passes_correct_teamspace_id(mock_parse_lit_url, mock_resolve, mock_client_cls, mock_get):
     mock_parse_lit_url.return_value = {"teamspace": "my-teamspace", "owner": "my-org", "destination": REMOTE_PATH}
@@ -85,7 +85,7 @@ def test_listdir_passes_correct_teamspace_id(mock_parse_lit_url, mock_resolve, m
     "lightning_sdk.api.filesystem_api._authenticate_and_get_auth_headers",
     new=mock.MagicMock(return_value=FAKE_AUTH_HEADERS),
 )
-@mock.patch("lightning_sdk.cli.utils.filesystem.resolve_teamspace")
+@mock.patch("lightning_sdk.cli.utils.filesystem.resolve_teamspace_ref")
 @mock.patch("lightning_sdk.cli.utils.filesystem.parse_lit_url")
 def test_listdir_non_recursive(
     mock_parse_lit_url, mock_resolve, mock_client_cls, mock_get, fake_teamspace, fake_path_result