Spaces:
Running
Running
| import datetime | |
| import logging | |
| import os | |
| import types | |
| import uuid | |
| from stat import S_ISDIR, S_ISLNK | |
| import paramiko | |
| from .. import AbstractFileSystem | |
| from ..utils import infer_storage_options | |
| logger = logging.getLogger("fsspec.sftp") | |
| class SFTPFileSystem(AbstractFileSystem): | |
| """Files over SFTP/SSH | |
| Peer-to-peer filesystem over SSH using paramiko. | |
| Note: if using this with the ``open`` or ``open_files``, with full URLs, | |
| there is no way to tell if a path is relative, so all paths are assumed | |
| to be absolute. | |
| """ | |
| protocol = "sftp", "ssh" | |
| def __init__(self, host, **ssh_kwargs): | |
| """ | |
| Parameters | |
| ---------- | |
| host: str | |
| Hostname or IP as a string | |
| temppath: str | |
| Location on the server to put files, when within a transaction | |
| ssh_kwargs: dict | |
| Parameters passed on to connection. See details in | |
| https://docs.paramiko.org/en/3.3/api/client.html#paramiko.client.SSHClient.connect | |
| May include port, username, password... | |
| """ | |
| if self._cached: | |
| return | |
| super().__init__(**ssh_kwargs) | |
| self.temppath = ssh_kwargs.pop("temppath", "/tmp") # remote temp directory | |
| self.host = host | |
| self.ssh_kwargs = ssh_kwargs | |
| self._connect() | |
| def _connect(self): | |
| logger.debug("Connecting to SFTP server %s", self.host) | |
| self.client = paramiko.SSHClient() | |
| self.client.set_missing_host_key_policy(paramiko.AutoAddPolicy()) | |
| self.client.connect(self.host, **self.ssh_kwargs) | |
| self.ftp = self.client.open_sftp() | |
| def _strip_protocol(cls, path): | |
| return infer_storage_options(path)["path"] | |
| def _get_kwargs_from_urls(urlpath): | |
| out = infer_storage_options(urlpath) | |
| out.pop("path", None) | |
| out.pop("protocol", None) | |
| return out | |
| def mkdir(self, path, create_parents=True, mode=511): | |
| logger.debug("Creating folder %s", path) | |
| if self.exists(path): | |
| raise FileExistsError(f"File exists: {path}") | |
| if create_parents: | |
| self.makedirs(path) | |
| else: | |
| self.ftp.mkdir(path, mode) | |
| def makedirs(self, path, exist_ok=False, mode=511): | |
| if self.exists(path) and not exist_ok: | |
| raise FileExistsError(f"File exists: {path}") | |
| parts = path.split("/") | |
| new_path = "/" if path[:1] == "/" else "" | |
| for part in parts: | |
| if part: | |
| new_path = f"{new_path}/{part}" if new_path else part | |
| if not self.exists(new_path): | |
| self.ftp.mkdir(new_path, mode) | |
| def rmdir(self, path): | |
| logger.debug("Removing folder %s", path) | |
| self.ftp.rmdir(path) | |
| def info(self, path): | |
| stat = self._decode_stat(self.ftp.stat(path)) | |
| stat["name"] = path | |
| return stat | |
| def _decode_stat(stat, parent_path=None): | |
| if S_ISDIR(stat.st_mode): | |
| t = "directory" | |
| elif S_ISLNK(stat.st_mode): | |
| t = "link" | |
| else: | |
| t = "file" | |
| out = { | |
| "name": "", | |
| "size": stat.st_size, | |
| "type": t, | |
| "uid": stat.st_uid, | |
| "gid": stat.st_gid, | |
| "time": datetime.datetime.fromtimestamp( | |
| stat.st_atime, tz=datetime.timezone.utc | |
| ), | |
| "mtime": datetime.datetime.fromtimestamp( | |
| stat.st_mtime, tz=datetime.timezone.utc | |
| ), | |
| } | |
| if parent_path: | |
| out["name"] = "/".join([parent_path.rstrip("/"), stat.filename]) | |
| return out | |
| def ls(self, path, detail=False): | |
| logger.debug("Listing folder %s", path) | |
| stats = [self._decode_stat(stat, path) for stat in self.ftp.listdir_iter(path)] | |
| if detail: | |
| return stats | |
| else: | |
| paths = [stat["name"] for stat in stats] | |
| return sorted(paths) | |
| def put(self, lpath, rpath, callback=None, **kwargs): | |
| logger.debug("Put file %s into %s", lpath, rpath) | |
| self.ftp.put(lpath, rpath) | |
| def get_file(self, rpath, lpath, **kwargs): | |
| if self.isdir(rpath): | |
| os.makedirs(lpath, exist_ok=True) | |
| else: | |
| self.ftp.get(self._strip_protocol(rpath), lpath) | |
| def _open(self, path, mode="rb", block_size=None, **kwargs): | |
| """ | |
| block_size: int or None | |
| If 0, no buffering, if 1, line buffering, if >1, buffer that many | |
| bytes, if None use default from paramiko. | |
| """ | |
| logger.debug("Opening file %s", path) | |
| if kwargs.get("autocommit", True) is False: | |
| # writes to temporary file, move on commit | |
| path2 = "/".join([self.temppath, str(uuid.uuid4())]) | |
| f = self.ftp.open(path2, mode, bufsize=block_size if block_size else -1) | |
| f.temppath = path2 | |
| f.targetpath = path | |
| f.fs = self | |
| f.commit = types.MethodType(commit_a_file, f) | |
| f.discard = types.MethodType(discard_a_file, f) | |
| else: | |
| f = self.ftp.open(path, mode, bufsize=block_size if block_size else -1) | |
| return f | |
| def _rm(self, path): | |
| if self.isdir(path): | |
| self.ftp.rmdir(path) | |
| else: | |
| self.ftp.remove(path) | |
| def mv(self, old, new): | |
| logger.debug("Renaming %s into %s", old, new) | |
| self.ftp.posix_rename(old, new) | |
| def commit_a_file(self): | |
| self.fs.mv(self.temppath, self.targetpath) | |
| def discard_a_file(self): | |
| self.fs._rm(self.temppath) | |