s3ui 1.0.0__py3-none-any.whl

s3ui/core/transfers.py

@@ -0,0 +1,281 @@
+"""Transfer engine — manages queue of uploads and downloads."""
+
+from __future__ import annotations
+
+import datetime
+import logging
+import threading
+from typing import TYPE_CHECKING
+
+from PyQt6.QtCore import QObject, QThreadPool, pyqtSignal
+
+from s3ui.core.download_worker import DownloadWorker
+from s3ui.core.upload_worker import UploadWorker
+
+if TYPE_CHECKING:
+    from s3ui.core.s3_client import S3Client
+    from s3ui.db.database import Database
+
+logger = logging.getLogger("s3ui.transfers")
+
+
+class TransferEngine(QObject):
+    """Manages the transfer queue and worker pool."""
+
+    transfer_progress = pyqtSignal(int, int, int)  # transfer_id, bytes_done, total
+    transfer_speed = pyqtSignal(int, float)  # transfer_id, bytes_per_sec
+    transfer_status_changed = pyqtSignal(int, str)  # transfer_id, new_status
+    transfer_error = pyqtSignal(int, str, str)  # transfer_id, user_msg, detail
+    transfer_finished = pyqtSignal(int)  # transfer_id
+
+    def __init__(
+        self,
+        s3_client: S3Client,
+        db: Database,
+        bucket: str,
+        max_workers: int = 4,
+    ) -> None:
+        super().__init__()
+        self._s3 = s3_client
+        self._db = db
+        self._bucket = bucket
+        row = self._db.fetchone(
+            "SELECT id FROM buckets WHERE name = ? ORDER BY id DESC LIMIT 1",
+            (bucket,),
+        )
+        self._bucket_id = row["id"] if row else None
+        self._pool = QThreadPool()
+        self._pool.setMaxThreadCount(max_workers)
+
+        # Per-transfer control events
+        self._pause_events: dict[int, threading.Event] = {}
+        self._cancel_events: dict[int, threading.Event] = {}
+        self._active: set[int] = set()
+        self._paused_global = False
+
+    def enqueue(self, transfer_id: int) -> None:
+        """Submit a transfer to the worker pool."""
+        row = self._db.fetchone("SELECT * FROM transfers WHERE id = ?", (transfer_id,))
+        if not row:
+            logger.warning("Cannot enqueue transfer %d: not found", transfer_id)
+            return
+
+        pause_evt = threading.Event()
+        cancel_evt = threading.Event()
+        self._pause_events[transfer_id] = pause_evt
+        self._cancel_events[transfer_id] = cancel_evt
+        self._active.add(transfer_id)
+
+        if row["direction"] == "upload":
+            worker = UploadWorker(
+                transfer_id,
+                self._s3,
+                self._db,
+                self._bucket,
+                pause_evt,
+                cancel_evt,
+            )
+        else:
+            worker = DownloadWorker(
+                transfer_id,
+                self._s3,
+                self._db,
+                self._bucket,
+                pause_evt,
+                cancel_evt,
+            )
+
+        # Connect signals
+        worker.signals.progress.connect(self._on_progress)
+        worker.signals.speed.connect(self._on_speed)
+        worker.signals.finished.connect(self._on_finished)
+        worker.signals.failed.connect(self._on_failed)
+
+        self._pool.start(worker)
+        self.transfer_status_changed.emit(transfer_id, "in_progress")
+        logger.info("Enqueued transfer %d (%s)", transfer_id, row["direction"])
+
+    def pause(self, transfer_id: int) -> None:
+        """Pause a running transfer."""
+        evt = self._pause_events.get(transfer_id)
+        if evt:
+            evt.set()
+            self.transfer_status_changed.emit(transfer_id, "paused")
+
+    def resume(self, transfer_id: int) -> None:
+        """Resume a paused transfer by re-enqueuing it."""
+        self._db.execute(
+            "UPDATE transfers SET status = 'queued', updated_at = datetime('now') WHERE id = ?",
+            (transfer_id,),
+        )
+        self.enqueue(transfer_id)
+
+    def cancel(self, transfer_id: int) -> None:
+        """Cancel a transfer."""
+        evt = self._cancel_events.get(transfer_id)
+        if evt:
+            evt.set()
+        self.transfer_status_changed.emit(transfer_id, "cancelled")
+
+    def pause_all(self) -> None:
+        """Pause all active transfers."""
+        self._paused_global = True
+        for tid in list(self._active):
+            self.pause(tid)
+
+    def resume_all(self) -> None:
+        """Resume all paused transfers."""
+        self._paused_global = False
+        rows = self._db.fetchall(
+            "SELECT id FROM transfers WHERE status = 'paused' AND bucket_id = "
+            "(SELECT id FROM buckets WHERE name = ? LIMIT 1)",
+            (self._bucket,),
+        )
+        for row in rows:
+            self.resume(row["id"])
+
+    def retry(self, transfer_id: int) -> None:
+        """Retry a failed transfer."""
+        self._db.execute(
+            "UPDATE transfers SET status = 'queued', retry_count = 0, "
+            "error_message = NULL, updated_at = datetime('now') WHERE id = ?",
+            (transfer_id,),
+        )
+        self.enqueue(transfer_id)
+
+    def restore_pending(self) -> None:
+        """Restore transfers that were interrupted by an app shutdown."""
+        rows = self._db.fetchall(
+            "SELECT id, direction, local_path, status FROM transfers "
+            "WHERE status IN ('queued', 'in_progress', 'paused')"
+        )
+        for row in rows:
+            from pathlib import Path
+
+            local = Path(row["local_path"])
+
+            if row["direction"] == "upload" and not local.exists():
+                self._db.execute(
+                    "UPDATE transfers SET status = 'failed', "
+                    "error_message = 'Source file no longer exists.', "
+                    "updated_at = datetime('now') WHERE id = ?",
+                    (row["id"],),
+                )
+                logger.warning("Transfer %d: source file missing: %s", row["id"], local)
+                continue
+
+            if row["direction"] == "download" and not local.parent.exists():
+                self._db.execute(
+                    "UPDATE transfers SET status = 'failed', "
+                    "error_message = 'Destination directory no longer exists.', "
+                    "updated_at = datetime('now') WHERE id = ?",
+                    (row["id"],),
+                )
+                logger.warning("Transfer %d: dest dir missing: %s", row["id"], local.parent)
+                continue
+
+            # Reset in_progress to queued
+            if row["status"] == "in_progress":
+                self._db.execute(
+                    "UPDATE transfers SET status = 'queued', "
+                    "updated_at = datetime('now') WHERE id = ?",
+                    (row["id"],),
+                )
+
+            self.enqueue(row["id"])
+            logger.info("Restored transfer %d", row["id"])
+
+    # --- Signal handlers ---
+
+    def _on_progress(self, transfer_id: int, bytes_done: int, total: int) -> None:
+        self.transfer_progress.emit(transfer_id, bytes_done, total)
+
+    def _on_speed(self, transfer_id: int, bps: float) -> None:
+        self.transfer_speed.emit(transfer_id, bps)
+
+    def _on_finished(self, transfer_id: int) -> None:
+        self._cleanup(transfer_id)
+        self.transfer_finished.emit(transfer_id)
+        self._pick_next()
+
+    def _on_failed(self, transfer_id: int, user_msg: str, detail: str) -> None:
+        self._cleanup(transfer_id)
+        self.transfer_error.emit(transfer_id, user_msg, detail)
+        self._pick_next()
+
+    def _cleanup(self, transfer_id: int) -> None:
+        self._active.discard(transfer_id)
+        self._pause_events.pop(transfer_id, None)
+        self._cancel_events.pop(transfer_id, None)
+
+    def _pick_next(self) -> None:
+        """Start the next queued transfer if a slot is available."""
+        if self._paused_global:
+            return
+        if self._bucket_id is None:
+            row = self._db.fetchone(
+                "SELECT id FROM transfers WHERE status = 'queued' ORDER BY created_at ASC LIMIT 1"
+            )
+        else:
+            row = self._db.fetchone(
+                "SELECT id FROM transfers WHERE status = 'queued' AND bucket_id = ? "
+                "ORDER BY created_at ASC LIMIT 1",
+                (self._bucket_id,),
+            )
+        if row and row["id"] not in self._active:
+            self.enqueue(row["id"])
+
+    def cleanup_orphaned_uploads(self) -> int:
+        """Abort orphaned multipart uploads on S3 not tracked in the database.
+
+        Only aborts uploads older than 24 hours to avoid interfering with
+        uploads started by other tools. Returns the number of aborted uploads.
+        """
+        try:
+            s3_uploads = self._s3.list_multipart_uploads(self._bucket)
+        except Exception:
+            logger.warning("Failed to list multipart uploads for orphan cleanup")
+            return 0
+
+        # Collect known upload IDs from our database
+        rows = self._db.fetchall("SELECT upload_id FROM transfers WHERE upload_id IS NOT NULL")
+        known_ids = {r["upload_id"] for r in rows}
+
+        cutoff = datetime.datetime.now(datetime.UTC) - datetime.timedelta(hours=24)
+        aborted = 0
+
+        for upload in s3_uploads:
+            uid = upload["UploadId"]
+            if uid in known_ids:
+                continue  # We own this one
+
+            initiated = upload["Initiated"]
+            # Make sure initiated is offset-aware for comparison
+            if initiated.tzinfo is None:
+                initiated = initiated.replace(tzinfo=datetime.UTC)
+
+            if initiated < cutoff:
+                try:
+                    self._s3.abort_multipart_upload(self._bucket, upload["Key"], uid)
+                    aborted += 1
+                    logger.info(
+                        "Aborted orphaned multipart upload: key=%s upload_id=%s",
+                        upload["Key"],
+                        uid,
+                    )
+                except Exception:
+                    logger.warning(
+                        "Failed to abort orphaned upload: key=%s upload_id=%s",
+                        upload["Key"],
+                        uid,
+                    )
+            else:
+                logger.debug(
+                    "Skipping recent orphaned upload: key=%s upload_id=%s (< 24h old)",
+                    upload["Key"],
+                    uid,
+                )
+
+        if aborted:
+            logger.info("Orphan cleanup: aborted %d uploads", aborted)
+        return aborted

s3ui/core/upload_worker.py

@@ -0,0 +1,311 @@
+"""Upload worker — handles single and multipart uploads as a QRunnable."""
+
+from __future__ import annotations
+
+import logging
+import math
+import random
+import threading
+import time
+from typing import TYPE_CHECKING
+
+from PyQt6.QtCore import QObject, QRunnable, pyqtSignal
+
+from s3ui.constants import (
+    DEFAULT_PART_SIZE,
+    HUGE_PART_SIZE,
+    LARGE_PART_SIZE,
+    MAX_RETRY_ATTEMPTS,
+    MULTIPART_THRESHOLD,
+)
+
+if TYPE_CHECKING:
+    from s3ui.core.s3_client import S3Client
+    from s3ui.db.database import Database
+
+logger = logging.getLogger("s3ui.upload_worker")
+
+
+def select_part_size(file_size: int) -> int:
+    """Choose part size based on file size to stay under 10,000 parts."""
+    if file_size <= 50 * 1024**3:  # ≤50 GB
+        return DEFAULT_PART_SIZE
+    elif file_size <= 500 * 1024**3:  # ≤500 GB
+        return LARGE_PART_SIZE
+    else:
+        return HUGE_PART_SIZE
+
+
+class UploadWorkerSignals(QObject):
+    progress = pyqtSignal(int, int, int)  # transfer_id, bytes_done, total
+    speed = pyqtSignal(int, float)  # transfer_id, bytes_per_sec
+    finished = pyqtSignal(int)  # transfer_id
+    failed = pyqtSignal(int, str, str)  # transfer_id, user_msg, detail
+
+
+class UploadWorker(QRunnable):
+    """Uploads a file to S3, with multipart support and resume."""
+
+    def __init__(
+        self,
+        transfer_id: int,
+        s3_client: S3Client,
+        db: Database,
+        bucket: str,
+        pause_event: threading.Event,
+        cancel_event: threading.Event,
+    ) -> None:
+        super().__init__()
+        self.setAutoDelete(True)
+        self.signals = UploadWorkerSignals()
+        self.transfer_id = transfer_id
+        self._s3 = s3_client
+        self._db = db
+        self._bucket = bucket
+        self._pause = pause_event
+        self._cancel = cancel_event
+
+        # Speed tracking
+        self._speed_window: list[tuple[float, int]] = []
+        self._last_speed_emit = 0.0
+
+    def run(self) -> None:
+        try:
+            self._do_upload()
+        except Exception as e:
+            import traceback
+
+            logger.error("Upload %d failed: %s", self.transfer_id, e)
+            try:
+                self._mark_failed(str(e))
+            except Exception:
+                logger.exception("Failed to mark upload %d as failed", self.transfer_id)
+            self.signals.failed.emit(self.transfer_id, str(e), traceback.format_exc())
+
+    def _do_upload(self) -> None:
+        row = self._db.fetchone("SELECT * FROM transfers WHERE id = ?", (self.transfer_id,))
+        if not row:
+            self.signals.failed.emit(self.transfer_id, "Transfer record not found.", "")
+            return
+
+        from pathlib import Path
+
+        local_path = Path(row["local_path"])
+        if not local_path.exists():
+            self._mark_failed("Source file no longer exists.")
+            self.signals.failed.emit(
+                self.transfer_id,
+                "Source file no longer exists.",
+                str(local_path),
+            )
+            return
+
+        object_key = row["object_key"]
+        file_size = local_path.stat().st_size
+
+        # Update total_bytes if not set
+        if row["total_bytes"] is None or row["total_bytes"] != file_size:
+            self._db.execute(
+                "UPDATE transfers SET total_bytes = ? WHERE id = ?",
+                (file_size, self.transfer_id),
+            )
+
+        self._db.execute(
+            "UPDATE transfers SET status = 'in_progress', updated_at = datetime('now') "
+            "WHERE id = ?",
+            (self.transfer_id,),
+        )
+
+        if file_size < MULTIPART_THRESHOLD:
+            self._single_upload(local_path, object_key, file_size)
+        else:
+            self._multipart_upload(local_path, object_key, file_size, row)
+
+    def _single_upload(self, local_path, object_key: str, file_size: int) -> None:
+        data = local_path.read_bytes()
+        self._s3.put_object(self._bucket, object_key, data)
+        self._complete(file_size)
+
+    def _multipart_upload(self, local_path, object_key: str, file_size: int, row) -> None:
+        part_size = select_part_size(file_size)
+        num_parts = math.ceil(file_size / part_size)
+        upload_id = row["upload_id"]
+
+        # Initiate or resume
+        if not upload_id:
+            upload_id = self._s3.create_multipart_upload(self._bucket, object_key)
+            self._db.execute(
+                "UPDATE transfers SET upload_id = ? WHERE id = ?",
+                (upload_id, self.transfer_id),
+            )
+            # Create part records
+            for i in range(num_parts):
+                offset = i * part_size
+                size = min(part_size, file_size - offset)
+                self._db.execute(
+                    "INSERT OR IGNORE INTO transfer_parts "
+                    "(transfer_id, part_number, offset, size) VALUES (?, ?, ?, ?)",
+                    (self.transfer_id, i + 1, offset, size),
+                )
+        else:
+            # Resuming: reconcile with S3
+            s3_parts = self._s3.list_parts(self._bucket, object_key, upload_id)
+            s3_confirmed = {p["PartNumber"] for p in s3_parts}
+            for part_num in s3_confirmed:
+                self._db.execute(
+                    "UPDATE transfer_parts SET status = 'completed' "
+                    "WHERE transfer_id = ? AND part_number = ?",
+                    (self.transfer_id, part_num),
+                )
+
+        # Upload pending parts
+        pending = self._db.fetchall(
+            "SELECT * FROM transfer_parts WHERE transfer_id = ? AND status != 'completed' "
+            "ORDER BY part_number",
+            (self.transfer_id,),
+        )
+
+        bytes_done = self._get_transferred()
+        parts_for_complete = self._get_completed_parts()
+
+        with open(local_path, "rb") as f:
+            for part_row in pending:
+                if self._cancel.is_set():
+                    self._do_cancel(object_key, upload_id)
+                    return
+                if self._pause.is_set():
+                    self._do_pause()
+                    return
+
+                part_num = part_row["part_number"]
+                offset = part_row["offset"]
+                size = part_row["size"]
+
+                f.seek(offset)
+                data = f.read(size)
+                etag = self._upload_part_with_retry(object_key, upload_id, part_num, data)
+                if etag is None:
+                    return  # failed signal already emitted
+
+                self._db.execute(
+                    "UPDATE transfer_parts SET status = 'completed', etag = ? "
+                    "WHERE transfer_id = ? AND part_number = ?",
+                    (etag, self.transfer_id, part_num),
+                )
+
+                bytes_done += size
+                self._db.execute(
+                    "UPDATE transfers SET transferred = ?, updated_at = datetime('now') "
+                    "WHERE id = ?",
+                    (bytes_done, self.transfer_id),
+                )
+                self.signals.progress.emit(self.transfer_id, bytes_done, file_size)
+                self._update_speed(size)
+                parts_for_complete.append({"ETag": etag, "PartNumber": part_num})
+
+        # Complete
+        all_parts = sorted(self._get_all_completed_parts(), key=lambda p: p["PartNumber"])
+        self._s3.complete_multipart_upload(self._bucket, object_key, upload_id, all_parts)
+        self._complete(file_size)
+
+    def _upload_part_with_retry(
+        self, key: str, upload_id: str, part_num: int, data: bytes
+    ) -> str | None:
+        for attempt in range(MAX_RETRY_ATTEMPTS):
+            try:
+                return self._s3.upload_part(self._bucket, key, upload_id, part_num, data)
+            except Exception as e:
+                if attempt < MAX_RETRY_ATTEMPTS - 1:
+                    delay = _backoff_delay(attempt)
+                    logger.warning(
+                        "Upload part %d attempt %d failed, retrying in %.1fs: %s",
+                        part_num,
+                        attempt + 1,
+                        delay,
+                        e,
+                    )
+                    time.sleep(delay)
+                else:
+                    self._mark_failed(str(e))
+                    self.signals.failed.emit(
+                        self.transfer_id,
+                        f"Upload failed after {MAX_RETRY_ATTEMPTS} attempts.",
+                        str(e),
+                    )
+                    return None
+
+    def _complete(self, total: int) -> None:
+        self._db.execute(
+            "UPDATE transfers SET status = 'completed', transferred = ?, "
+            "updated_at = datetime('now') WHERE id = ?",
+            (total, self.transfer_id),
+        )
+        self.signals.progress.emit(self.transfer_id, total, total)
+        self.signals.finished.emit(self.transfer_id)
+        logger.info("Upload %d completed", self.transfer_id)
+
+    def _mark_failed(self, msg: str) -> None:
+        self._db.execute(
+            "UPDATE transfers SET status = 'failed', error_message = ?, "
+            "updated_at = datetime('now') WHERE id = ?",
+            (msg, self.transfer_id),
+        )
+
+    def _do_cancel(self, key: str, upload_id: str) -> None:
+        import contextlib
+
+        with contextlib.suppress(Exception):
+            self._s3.abort_multipart_upload(self._bucket, key, upload_id)
+        self._db.execute(
+            "UPDATE transfers SET status = 'cancelled', updated_at = datetime('now') WHERE id = ?",
+            (self.transfer_id,),
+        )
+        logger.info("Upload %d cancelled", self.transfer_id)
+
+    def _do_pause(self) -> None:
+        self._db.execute(
+            "UPDATE transfers SET status = 'paused', updated_at = datetime('now') WHERE id = ?",
+            (self.transfer_id,),
+        )
+        logger.info("Upload %d paused", self.transfer_id)
+
+    def _get_transferred(self) -> int:
+        row = self._db.fetchone(
+            "SELECT COALESCE(SUM(size), 0) as done FROM transfer_parts "
+            "WHERE transfer_id = ? AND status = 'completed'",
+            (self.transfer_id,),
+        )
+        return row["done"]
+
+    def _get_completed_parts(self) -> list[dict]:
+        rows = self._db.fetchall(
+            "SELECT part_number, etag FROM transfer_parts "
+            "WHERE transfer_id = ? AND status = 'completed'",
+            (self.transfer_id,),
+        )
+        return [{"ETag": r["etag"], "PartNumber": r["part_number"]} for r in rows]
+
+    def _get_all_completed_parts(self) -> list[dict]:
+        return self._get_completed_parts()
+
+    def _update_speed(self, chunk_bytes: int) -> None:
+        now = time.monotonic()
+        self._speed_window.append((now, chunk_bytes))
+        # Keep 3-second window
+        self._speed_window = [(t, b) for t, b in self._speed_window if now - t <= 3.0]
+        if now - self._last_speed_emit >= 0.5 and self._speed_window:
+            window_time = now - self._speed_window[0][0]
+            if window_time > 0:
+                total_bytes = sum(b for _, b in self._speed_window)
+                bps = total_bytes / window_time
+                self.signals.speed.emit(self.transfer_id, bps)
+            self._last_speed_emit = now
+
+
+def _backoff_delay(attempt: int) -> float:
+    """Exponential backoff with jitter: 0s, ~1s, ~4s."""
+    if attempt == 0:
+        return 0.0
+    base = 4 ** (attempt - 1)  # 1, 4
+    jitter_max = base * 0.5
+    return base + random.uniform(0, jitter_max)