diff --git a/singlestoredb/fusion/handlers/files.py b/singlestoredb/fusion/handlers/files.py index d823aeae..7f848611 100644 --- a/singlestoredb/fusion/handlers/files.py +++ b/singlestoredb/fusion/handlers/files.py @@ -32,6 +32,7 @@ def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: for x in file_space.listdir( params['at_path'] or '/', recursive=params['recursive'], + return_objects=False, ): info = file_space.info(x) files.append( @@ -47,6 +48,7 @@ def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]: res.set_rows([(x,) for x in file_space.listdir( params['at_path'] or '/', recursive=params['recursive'], + return_objects=False, )]) if params['like']: diff --git a/singlestoredb/management/files.py b/singlestoredb/management/files.py index ac3561ad..593f7e39 100644 --- a/singlestoredb/management/files.py +++ b/singlestoredb/management/files.py @@ -10,9 +10,12 @@ from abc import ABC from abc import abstractmethod from typing import Any +from typing import cast from typing import Dict from typing import List +from typing import Literal from typing import Optional +from typing import overload from typing import Union from .. import config @@ -421,15 +424,36 @@ def is_dir(self, path: PathLike) -> bool: def is_file(self, path: PathLike) -> bool: pass - @abstractmethod + @overload def listdir( self, path: PathLike = '/', *, recursive: bool = False, + return_objects: Literal[True], + ) -> List[FilesObject]: + pass + + @overload + def listdir( + self, + path: PathLike = '/', + *, + recursive: bool = False, + return_objects: Literal[False] = False, ) -> List[str]: pass + @abstractmethod + def listdir( + self, + path: PathLike = '/', + *, + recursive: bool = False, + return_objects: bool = False, + ) -> Union[List[str], List[FilesObject]]: + pass + @abstractmethod def download_file( self, @@ -908,9 +932,13 @@ def is_file(self, path: PathLike) -> bool: return False raise - def _listdir(self, path: PathLike, *, recursive: bool = False) -> List[str]: + def _listdir( + self, path: PathLike, *, + recursive: bool = False, + return_objects: bool = False, + ) -> List[Union[str, FilesObject]]: """ - Return the names of files in a directory. + Return the names (or FilesObject instances) of files in a directory. Parameters ---------- @@ -918,28 +946,64 @@ def _listdir(self, path: PathLike, *, recursive: bool = False) -> List[str]: Path to the folder recursive : bool, optional Should folders be listed recursively? - + return_objects : bool, optional + If True, return list of FilesObject instances. Otherwise just paths. """ res = self._manager._get( f'files/fs/{self._location}/{path}', ).json() if recursive: - out = [] - for item in res['content'] or []: - out.append(item['path']) + out: List[Union[str, FilesObject]] = [] + for item in res.get('content') or []: + if return_objects: + out.append(FilesObject.from_dict(item, self)) + else: + out.append(item['path']) if item['type'] == 'directory': - out.extend(self._listdir(item['path'], recursive=recursive)) + out.extend( + self._listdir( + item['path'], + recursive=recursive, + return_objects=return_objects, + ), + ) return out - return [x['path'] for x in res['content'] or []] + if return_objects: + return [ + FilesObject.from_dict(x, self) + for x in (res.get('content') or []) + ] + return [x['path'] for x in (res.get('content') or [])] + @overload def listdir( self, path: PathLike = '/', *, recursive: bool = False, + return_objects: Literal[True], + ) -> List[FilesObject]: + ... + + @overload + def listdir( + self, + path: PathLike = '/', + *, + recursive: bool = False, + return_objects: Literal[False] = False, ) -> List[str]: + ... + + def listdir( + self, + path: PathLike = '/', + *, + recursive: bool = False, + return_objects: bool = False, + ) -> Union[List[str], List[FilesObject]]: """ List the files / folders at the given path. @@ -948,22 +1012,40 @@ def listdir( path : Path or str, optional Path to the file location + return_objects : bool, optional + If True, return list of FilesObject instances. Otherwise just paths. + Returns ------- - List[str] + List[str] or List[FilesObject] """ path = re.sub(r'^(\./|/)+', r'', str(path)) path = re.sub(r'/+$', r'', path) + '/' - if not self.is_dir(path): - raise NotADirectoryError(f'path is not a directory: {path}') + # Validate via listing GET; if response lacks 'content', it's not a directory + try: + out = self._listdir(path, recursive=recursive, return_objects=return_objects) + except (ManagementError, KeyError) as exc: + # If the path doesn't exist or isn't a directory, _listdir will fail + raise NotADirectoryError(f'path is not a directory: {path}') from exc - out = self._listdir(path, recursive=recursive) if path != '/': path_n = len(path.split('/')) - 1 - out = ['/'.join(x.split('/')[path_n:]) for x in out] - return out + if return_objects: + result: List[FilesObject] = [] + for item in out: + if isinstance(item, FilesObject): + rel = '/'.join(item.path.split('/')[path_n:]) + item.path = rel + result.append(item) + return result + return ['/'.join(str(x).split('/')[path_n:]) for x in out] + + # _listdir guarantees homogeneous type based on return_objects + if return_objects: + return cast(List[FilesObject], out) + return cast(List[str], out) def download_file( self, @@ -992,10 +1074,49 @@ def download_file( bytes or str - ``local_path`` is None None - ``local_path`` is a Path or str + """ + return self._download_file( + path, + local_path=local_path, + overwrite=overwrite, + encoding=encoding, + _skip_dir_check=False, + ) + + def _download_file( + self, + path: PathLike, + local_path: Optional[PathLike] = None, + *, + overwrite: bool = False, + encoding: Optional[str] = None, + _skip_dir_check: bool = False, + ) -> Optional[Union[bytes, str]]: + """ + Internal method to download the content of a file path. + + Parameters + ---------- + path : Path or str + Path to the file + local_path : Path or str + Path to local file target location + overwrite : bool, optional + Should an existing file be overwritten if it exists? + encoding : str, optional + Encoding used to convert the resulting data + _skip_dir_check : bool, optional + Skip the directory check (internal use only) + + Returns + ------- + bytes or str - ``local_path`` is None + None - ``local_path`` is a Path or str + """ if local_path is not None and not overwrite and os.path.exists(local_path): raise OSError('target file already exists; use overwrite=True to replace') - if self.is_dir(path): + if not _skip_dir_check and self.is_dir(path): raise IsADirectoryError(f'file path is a directory: {path}') out = self._manager._get( @@ -1036,17 +1157,27 @@ def download_folder( if local_path is not None and not overwrite and os.path.exists(local_path): raise OSError('target path already exists; use overwrite=True to replace') - if not self.is_dir(path): - raise NotADirectoryError(f'path is not a directory: {path}') - - files = self.listdir(path, recursive=True) - for f in files: - remote_path = os.path.join(path, f) - if self.is_dir(remote_path): + # listdir validates directory; no extra info call needed + entries = self.listdir(path, recursive=True, return_objects=True) + for entry in entries: + # Each entry is a FilesObject with path relative to root and type + if not isinstance(entry, FilesObject): # defensive: skip unexpected + continue + rel_path = entry.path + if entry.type == 'directory': + # Ensure local directory exists; no remote call needed + target_dir = os.path.normpath(os.path.join(local_path, rel_path)) + os.makedirs(target_dir, exist_ok=True) continue - target = os.path.normpath(os.path.join(local_path, f)) - os.makedirs(os.path.dirname(target), exist_ok=True) - self.download_file(remote_path, target, overwrite=overwrite) + remote_path = os.path.join(path, rel_path) + target_file = os.path.normpath( + os.path.join(local_path, rel_path), + ) + os.makedirs(os.path.dirname(target_file), exist_ok=True) + self._download_file( + remote_path, target_file, + overwrite=overwrite, _skip_dir_check=True, + ) def remove(self, path: PathLike) -> None: """ diff --git a/singlestoredb/management/workspace.py b/singlestoredb/management/workspace.py index 6df98a12..1b5d7c27 100644 --- a/singlestoredb/management/workspace.py +++ b/singlestoredb/management/workspace.py @@ -10,9 +10,12 @@ import time from collections.abc import Mapping from typing import Any +from typing import cast from typing import Dict from typing import List +from typing import Literal from typing import Optional +from typing import overload from typing import Union from .. import config @@ -447,9 +450,13 @@ def is_file(self, stage_path: PathLike) -> bool: return False raise - def _listdir(self, stage_path: PathLike, *, recursive: bool = False) -> List[str]: + def _listdir( + self, stage_path: PathLike, *, + recursive: bool = False, + return_objects: bool = False, + ) -> List[Union[str, 'FilesObject']]: """ - Return the names of files in a directory. + Return the names (or FilesObject instances) of files in a directory. Parameters ---------- @@ -457,26 +464,64 @@ def _listdir(self, stage_path: PathLike, *, recursive: bool = False) -> List[str Path to the folder in Stage recursive : bool, optional Should folders be listed recursively? + return_objects : bool, optional + If True, return list of FilesObject instances. Otherwise just paths. """ + from .files import FilesObject res = self._manager._get( re.sub(r'/+$', r'/', f'stage/{self._deployment_id}/fs/{stage_path}'), ).json() if recursive: - out = [] + out: List[Union[str, FilesObject]] = [] for item in res['content'] or []: - out.append(item['path']) + if return_objects: + out.append(FilesObject.from_dict(item, self)) + else: + out.append(item['path']) if item['type'] == 'directory': - out.extend(self._listdir(item['path'], recursive=recursive)) + out.extend( + self._listdir( + item['path'], + recursive=recursive, + return_objects=return_objects, + ), + ) return out + if return_objects: + return [ + FilesObject.from_dict(x, self) + for x in res['content'] or [] + ] return [x['path'] for x in res['content'] or []] + @overload def listdir( self, stage_path: PathLike = '/', *, recursive: bool = False, + return_objects: Literal[True], + ) -> List['FilesObject']: + ... + + @overload + def listdir( + self, + stage_path: PathLike = '/', + *, + recursive: bool = False, + return_objects: Literal[False] = False, ) -> List[str]: + ... + + def listdir( + self, + stage_path: PathLike = '/', + *, + recursive: bool = False, + return_objects: bool = False, + ) -> Union[List[str], List['FilesObject']]: """ List the files / folders at the given path. @@ -484,21 +529,40 @@ def listdir( ---------- stage_path : Path or str, optional Path to the stage location + recursive : bool, optional + If True, recursively list all files and folders + return_objects : bool, optional + If True, return list of FilesObject instances. Otherwise just paths. Returns ------- - List[str] + List[str] or List[FilesObject] """ + from .files import FilesObject stage_path = re.sub(r'^(\./|/)+', r'', str(stage_path)) stage_path = re.sub(r'/+$', r'', stage_path) + '/' if self.is_dir(stage_path): - out = self._listdir(stage_path, recursive=recursive) + out = self._listdir( + stage_path, + recursive=recursive, + return_objects=return_objects, + ) if stage_path != '/': stage_path_n = len(stage_path.split('/')) - 1 - out = ['/'.join(x.split('/')[stage_path_n:]) for x in out] - return out + if return_objects: + result: List[FilesObject] = [] + for item in out: + if isinstance(item, FilesObject): + rel = '/'.join(item.path.split('/')[stage_path_n:]) + item.path = rel + result.append(item) + return result + out = ['/'.join(str(x).split('/')[stage_path_n:]) for x in out] + if return_objects: + return cast(List[FilesObject], out) + return cast(List[str], out) raise NotADirectoryError(f'stage path is not a directory: {stage_path}') @@ -577,7 +641,7 @@ def download_folder( if not self.is_dir(stage_path): raise NotADirectoryError(f'stage path is not a directory: {stage_path}') - for f in self.listdir(stage_path, recursive=True): + for f in self.listdir(stage_path, recursive=True, return_objects=False): if self.is_dir(f): continue target = os.path.normpath(os.path.join(local_path, f)) diff --git a/singlestoredb/tests/test_management.py b/singlestoredb/tests/test_management.py index dd176f3d..f450e1f1 100755 --- a/singlestoredb/tests/test_management.py +++ b/singlestoredb/tests/test_management.py @@ -757,6 +757,60 @@ def test_os_directories(self): with self.assertRaises(s2.ManagementError): st.removedirs(mkdir_test_sql) + def test_listdir_return_objects(self): + st = self.wg.stage + + listdir_test_dir = f'listdir_test_{id(self)}' + listdir_test_sql = f'listdir_test_{id(self)}.sql' + + # Create test directory structure + st.mkdir(listdir_test_dir) + st.mkdir(f'{listdir_test_dir}/nest_1') + st.upload_file(TEST_DIR / 'test.sql', listdir_test_sql) + st.upload_file( + TEST_DIR / 'test.sql', + f'{listdir_test_dir}/nested_test.sql', + ) + + # Test return_objects=False (default behavior) + out = st.listdir('/') + assert isinstance(out, list) + assert all(isinstance(item, str) for item in out) + assert f'{listdir_test_dir}/' in out + assert listdir_test_sql in out + + # Test return_objects=True + out_objs = st.listdir('/', return_objects=True) + assert isinstance(out_objs, list) + assert all(hasattr(item, 'path') for item in out_objs) + assert all(hasattr(item, 'type') for item in out_objs) + + # Verify we have the expected items + obj_paths = [obj.path for obj in out_objs] + assert f'{listdir_test_dir}/' in obj_paths + assert listdir_test_sql in obj_paths + + # Verify object types + for obj in out_objs: + if obj.path == f'{listdir_test_dir}/': + assert obj.type == 'directory' + elif obj.path == listdir_test_sql: + assert obj.type == 'file' + + # Test with subdirectory and return_objects=True + out_objs_sub = st.listdir(listdir_test_dir, return_objects=True) + assert isinstance(out_objs_sub, list) + obj_paths_sub = [obj.path for obj in out_objs_sub] + assert 'nest_1/' in obj_paths_sub + assert 'nested_test.sql' in obj_paths_sub + + # Test recursive with return_objects=True + out_objs_rec = st.listdir('/', recursive=True, return_objects=True) + obj_paths_rec = [obj.path for obj in out_objs_rec] + assert f'{listdir_test_dir}/' in obj_paths_rec + assert f'{listdir_test_dir}/nest_1/' in obj_paths_rec + assert f'{listdir_test_dir}/nested_test.sql' in obj_paths_rec + def test_os_files(self): st = self.wg.stage