Enable direct streaming for SFTP uploads

dcfs/app/sftp/handler.py

@@ -188,7 +188,7 @@ async def open(self, path: bytes, flags: int, attrs: asyncssh.SFTPAttrs) -> 'DCF
             if expected_size is not None and expected_size < 0:
                 expected_size = None
 
-        if mode == "w" and expected_size is not None:
+        if mode == "w":
             return DCFSSFTPStreamingFile(ops, sub_path, expected_size, client_name)
         return DCFSSFTPBufferedFile(ops, sub_path, mode, client_name)
 
@@ -428,10 +428,16 @@ async def fstat(self) -> asyncssh.SFTPAttrs:
 
 
 class DCFSSFTPStreamingFile(DCFSSFTPFileBase):
-    def __init__(self, ops: Ops, path: str, expected_size: int, client_name: str):
+    def __init__(
+        self,
+        ops: Ops,
+        path: str,
+        expected_size: Optional[int],
+        client_name: str,
+    ):
         self.ops = ops
         self.path = path
-        self.expected_size = expected_size
+        self.expected_size = expected_size if expected_size is not None else -1
         self.client_name = client_name
         self.closed = False
 
@@ -518,7 +524,7 @@ async def write(self, offset: int, data: bytes) -> int:
             return 0
 
         end = offset + len(data)
-        if end > self.expected_size:
+        if self.expected_size >= 0 and end > self.expected_size:
             raise asyncssh.SFTPFailure("Write exceeds expected file size")
 
         existing = self._pending.get(offset)
@@ -542,7 +548,7 @@ async def close(self) -> None:
 
         self.closed = True
 
-        if self._next_offset != self.expected_size:
+        if self.expected_size >= 0 and self._next_offset != self.expected_size:
             raise asyncssh.SFTPFailure(
                 f"Upload incomplete: received {self._next_offset} of {self.expected_size} bytes"
             )
@@ -555,7 +561,7 @@ async def fstat(self) -> asyncssh.SFTPAttrs:
         now = int(time.time())
         return asyncssh.SFTPAttrs(
             permissions=stat.S_IFREG | 0o644,
-            size=self.expected_size,
+            size=max(0, self.expected_size),
             mtime=now,
             atime=now,
         )

dcfs/core/repository/impl/file_content/__init__.py

@@ -85,42 +85,70 @@ async def save(self, file_msg: UploadableFileMessage) -> List[SentFileMessage]:
         tasks = []
 
         max_part_size = discord_max_file_size_bytes()
-        for i, part_size in enumerate(self._partition(total_size, max_part_size)):
-            # Acquire semaphore BEFORE buffering to keep memory usage under control.
-            # Only N parts will be buffered and uploading concurrently.
-            await self._upload_semaphore.acquire()
-
-            # Update file_msg for the current part
-            file_msg.name = f"[part{i+1}]{original_name}"
-            file_msg.size = part_size
-
-            # Create an uploader for this part. Using next_bot to distribute load
-            # across multiple bots if available.
-            uploader = FileUploader(self._message_api.discord_api.next_bot, file_msg)
-
-            # Read and buffer the part content sequentially from the source stream.
-            await uploader.upload()
-
-            # Dispatch the upload as a background task.
-            tasks.append(
-                asyncio.create_task(
-                    self._upload_part_task(
-                        uploader,
-                        self._message_api.private_file_channel,
-                        i + 1,
-                        part_size,
-                        self._upload_semaphore,
+
+        if total_size >= 0:
+            # Case 1: Total size is known upfront.
+            for i, part_size in enumerate(self._partition(total_size, max_part_size)):
+                await self._upload_semaphore.acquire()
+
+                file_msg.name = f"[part{i+1}]{original_name}"
+                file_msg.size = part_size
+
+                uploader = FileUploader(self._message_api.discord_api.next_bot, file_msg)
+                await uploader.upload()
+
+                tasks.append(
+                    asyncio.create_task(
+                        self._upload_part_task(
+                            uploader,
+                            self._message_api.private_file_channel,
+                            i + 1,
+                            part_size,
+                            self._upload_semaphore,
+                        )
                     )
                 )
-            )
+                await asyncio.sleep(0)
+                file_msg.next_part(part_size)
+        else:
+            # Case 2: Total size is unknown (e.g. streaming from SFTP).
+            # We read and upload parts until the stream is exhausted.
+            i = 0
+            while True:
+                await self._upload_semaphore.acquire()
+
+                file_msg.name = f"[part{i+1}]{original_name}"
+                # Set size to -1 for FileUploader to allow reading up to max_part_size
+                # without being capped by FileMessageFromStream.read().
+                file_msg.size = -1
+
+                uploader = FileUploader(self._message_api.discord_api.next_bot, file_msg)
+                part_size = await uploader.upload()
+
+                if part_size == 0 and i > 0:
+                    # EOF reached after some parts were uploaded.
+                    self._upload_semaphore.release()
+                    break
+
+                tasks.append(
+                    asyncio.create_task(
+                        self._upload_part_task(
+                            uploader,
+                            self._message_api.private_file_channel,
+                            i + 1,
+                            part_size,
+                            self._upload_semaphore,
+                        )
+                    )
+                )
+                await asyncio.sleep(0)
 
-            # Yield to the event loop between parts so the Discord
-            # gateway heartbeat and other async tasks can make progress
-            # during large uploads.
-            await asyncio.sleep(0)
+                if part_size < max_part_size:
+                    # Last part reached.
+                    break
 
-            # Advance the source message to the next part
-            file_msg.next_part(part_size)
+                file_msg.next_part(part_size)
+                i += 1
 
         # Wait for all uploads to complete.
         return list(await asyncio.gather(*tasks))

dcfs/core/repository/impl/file_content/file_uploader.py

@@ -66,7 +66,11 @@ async def upload(self) -> int:
         await self._file_msg.open()
         self._buffer = io.BytesIO()
         while True:
-            chunk = await self._file_msg.read(1024 * 1024)
+            remaining = limit - self._buffer.tell()
+            if remaining <= 0:
+                break
+
+            chunk = await self._file_msg.read(min(1024 * 1024, remaining))
             if not chunk:
                 break
             self._buffer.write(chunk)

dcfs/reqres.py

@@ -275,7 +275,12 @@ def new(
         )
 
     async def read(self, length: int) -> bytes:
-        size_to_return = min(length, self.get_size() - self._read_size)
+        total_size = self.get_size()
+        if total_size >= 0:
+            size_to_return = min(length, total_size - self._read_size)
+        else:
+            size_to_return = length
+
         if size_to_return <= 0:
             return b""