diff --git a/dcfs/app/sftp/handler.py b/dcfs/app/sftp/handler.py index 43bdfac..90ee2bc 100644 --- a/dcfs/app/sftp/handler.py +++ b/dcfs/app/sftp/handler.py @@ -4,7 +4,7 @@ import os import stat import time -from typing import Any, AsyncIterator, Optional +from typing import Any, AsyncGenerator, AsyncIterator, Optional, cast import asyncssh @@ -330,22 +330,51 @@ def __init__(self, ops: Ops, path: str, mode: str, client_name: str): self.buffer = bytearray() self.closed = False + # Read streaming state + self._read_stream: Optional[AsyncIterator[bytes]] = None + self._read_iter: Optional[AsyncIterator[bytes]] = None + self._read_pos = 0 + self._read_buf = bytearray() + self._read_lock = asyncio.Lock() + self._cached_attrs: Optional[asyncssh.SFTPAttrs] = None + async def read(self, offset: int, size: int) -> bytes: if "r" not in self.mode: raise asyncssh.SFTPPermissionDenied("File not open for reading") - end = -1 if size < 0 else offset + size - 1 - stream = await self.ops.download( - self.path, - offset, - end, - os.path.basename(self.path), - validate=False, - ) - data = bytearray() - async for chunk in stream: - data.extend(chunk) - return bytes(data) + async with self._read_lock: + # If we don't have a stream or it's at the wrong position, start a new one. + if self._read_stream is None or offset != self._read_pos: + if self._read_stream is not None: + await cast(AsyncGenerator[bytes, None], self._read_stream).aclose() + self._read_stream = None + self._read_iter = None + + self._read_buf = bytearray() + self._read_pos = offset + self._read_stream = await self.ops.download( + self.path, + offset, + -1, + os.path.basename(self.path), + validate=False, + ) + self._read_iter = self._read_stream.__aiter__() + + # it is impossible for self._read_iter to be None here given the logic above, + # but we cast to satisfy mypy. + it = cast(AsyncIterator[bytes], self._read_iter) + while len(self._read_buf) < size: + try: + chunk = await anext(it) + self._read_buf.extend(chunk) + except StopAsyncIteration: + break + + data = self._read_buf[:size] + self._read_buf = self._read_buf[size:] + self._read_pos += len(data) + return bytes(data) async def write(self, offset: int, data: bytes) -> int: if "w" not in self.mode: @@ -371,6 +400,11 @@ async def close(self) -> None: self.closed = True + if self._read_stream is not None: + await cast(AsyncGenerator[bytes, None], self._read_stream).aclose() + self._read_stream = None + self._read_iter = None + if "w" in self.mode and self.buffer: data = bytes(self.buffer) self.buffer.clear() @@ -415,16 +449,20 @@ async def fstat(self) -> asyncssh.SFTPAttrs: atime=now, ) + if self._cached_attrs: + return self._cached_attrs + fd = await self.ops.desc(self.path, validate=False) fv = fd.get_latest_version() size = await self.ops._client.fc_repo.content_length(fv) mtime = int(fv.updated_at_timestamp / 1000) - return asyncssh.SFTPAttrs( + self._cached_attrs = asyncssh.SFTPAttrs( permissions=stat.S_IFREG | 0o644, size=size, mtime=mtime, atime=mtime, ) + return self._cached_attrs class DCFSSFTPStreamingFile(DCFSSFTPFileBase):