rclone-api 1.5.8__py3-none-any.whl

rclone_api/db/db.py

@@ -0,0 +1,277 @@
+"""
+Database module for rclone_api.
+"""
+
+import os
+from threading import Lock
+from typing import Optional
+
+from sqlmodel import Session, SQLModel, create_engine, select
+
+from rclone_api.db.models import RepositoryMeta, create_file_entry_model
+from rclone_api.file import FileItem
+
+
+def _to_table_name(remote_name: str) -> str:
+    return (
+        "files_"
+        + remote_name.replace(":", "_").replace(" ", "_").replace("/", "_").lower()
+    )
+
+
+class DB:
+    """Database class for rclone_api."""
+
+    def __init__(self, db_path_url: str):
+        """Initialize the database.
+
+        Args:
+            db_path: Path to the database file
+        """
+        self.db_path_url = db_path_url
+
+        # When running multiple commands in parallel, the database connection may fail once
+        # when the database is first populated.
+        retries = 2
+        for _ in range(retries):
+            try:
+                self.engine = create_engine(db_path_url)
+                SQLModel.metadata.create_all(self.engine)
+                break
+            except Exception as e:
+                print(f"Failed to connect to database. Retrying... {e}")
+        else:
+            raise Exception("Failed to connect to database.")
+        self._cache: dict[str, DBRepo] = {}
+        self._cache_lock = Lock()
+
+    def drop_all(self) -> None:
+        """Drop all tables in the database."""
+        SQLModel.metadata.drop_all(self.engine)
+
+    def close(self) -> None:
+        """Close the database connection and release resources."""
+        if hasattr(self, "engine") and self.engine is not None:
+            self.engine.dispose()
+
+    def add_files(self, files: list[FileItem]) -> None:
+        """Add files to the database.
+
+        Args:
+            remote_name: Name of the remote
+            files: List of file entries
+        """
+
+        partition: dict[str, list[FileItem]] = {}
+        for file in files:
+            partition.setdefault(file.remote, []).append(file)
+
+        for remote_name, files in partition.items():
+            repo = self.get_or_create_repo(remote_name)
+            repo.insert_files(files)
+
+    def query_all_files(self, remote_name: str) -> list[FileItem]:
+        """Query files from the database.
+
+        Args:
+            remote_name: Name of the remote
+        """
+        repo = self.get_or_create_repo(remote_name)
+        files = repo.get_all_files()
+        out: list[FileItem] = []
+        for file in files:
+            out.append(file)
+        return out
+
+    def get_or_create_repo(self, remote_name: str) -> "DBRepo":
+        """Get a table section for a remote.
+
+        Args:
+            remote_name: Name of the remote
+            table_name: Optional table name, will be derived from remote_name if not provided
+
+        Returns:
+            DBRepo: A table section for the remote
+        """
+        with self._cache_lock:
+            if remote_name in self._cache:
+                return self._cache[remote_name]
+            table_name = _to_table_name(remote_name)
+            out = DBRepo(self.engine, remote_name, table_name)
+            self._cache[remote_name] = out
+            return out
+
+
+class DBRepo:
+    """Table repo remote."""
+
+    def __init__(self, engine, remote_name: str, table_name: Optional[str] = None):
+        """Initialize a table section.
+
+        Args:
+            engine: SQLAlchemy engine
+            remote_name: Name of the remote
+            table_name: Optional table name, will be derived from remote_name if not provided
+        """
+        self.engine = engine
+        self.remote_name = remote_name
+
+        # If table_name is not provided, derive one from the remote name.
+        if table_name is None:
+            # table_name = (
+            #     "file_entries_"
+            #     + remote_name.replace(":", "_").replace(" ", "_").replace("/", "_").lower()
+            # )
+            table_name = _to_table_name(remote_name)
+        self.table_name = table_name
+
+        # Check if repository exists in RepositoryMeta; if not, create a new entry.
+        with Session(self.engine) as session:
+            existing_repo = session.exec(
+                select(RepositoryMeta).where(
+                    RepositoryMeta.repo_name == self.remote_name
+                )
+            ).first()
+            if not existing_repo:
+                repo_meta = RepositoryMeta(
+                    repo_name=self.remote_name, file_table_name=self.table_name
+                )
+                session.add(repo_meta)
+                session.commit()
+
+        # Dynamically create the file entry model and its table.
+        self.FileEntryModel = create_file_entry_model(self.table_name)
+        SQLModel.metadata.create_all(self.engine, tables=[self.FileEntryModel.__table__])  # type: ignore
+
+    def insert_file(self, file: FileItem) -> None:
+        """Insert a file entry into the table.
+
+        Args:
+            file: File entry
+        """
+        return self.insert_files([file])
+
+    def insert_files(self, files: list[FileItem]) -> None:
+        """
+        Insert multiple file entries into the table.
+
+        Three bulk operations are performed:
+        1. Select: Determine which files already exist.
+        2. Insert: Bulk-insert new file entries.
+        3. Update: Bulk-update existing file entries.
+
+        The FileEntryModel must define a unique constraint on (path, name) and have a primary key "id".
+        """
+        # Step 1: Bulk select existing records.
+        # get_exists() returns a set of FileItem objects (based on path_no_remote and name) that already exist.
+        existing_files = self.get_exists(files)
+
+        # Determine which files need to be updated vs. inserted.
+        needs_update = existing_files
+        is_new = set(files) - existing_files
+
+        # Step 2: Bulk insert new rows.
+        new_values = [
+            {
+                "path": file.path_no_remote,
+                "name": file.name,
+                "size": file.size,
+                "mime_type": file.mime_type,
+                "mod_time": file.mod_time,
+                "suffix": file.real_suffix,
+            }
+            for file in is_new
+        ]
+        with Session(self.engine) as session:
+            if new_values:
+                session.bulk_insert_mappings(self.FileEntryModel, new_values)  # type: ignore
+                session.commit()
+
+        # Step 3: Bulk update existing rows.
+        # First, query the database for the primary keys of rows that match the unique keys in needs_update.
+        with Session(self.engine) as session:
+            # Collect all unique paths from files needing update.
+            update_paths = [file.path_no_remote for file in needs_update]
+            # Query for existing rows matching any of these paths.
+            db_entries = session.exec(
+                select(self.FileEntryModel).where(
+                    self.FileEntryModel.path.in_(update_paths)  # type: ignore
+                )
+            ).all()
+
+            # Build a mapping from the unique key (path, name) to the primary key (id).
+            id_map = {(entry.path, entry.name): entry.id for entry in db_entries}
+
+            # Prepare bulk update mappings.
+            update_values = []
+            for file in needs_update:
+                key = (file.path_no_remote, file.name)
+                if key in id_map:
+                    update_values.append(
+                        {
+                            "id": id_map[key],
+                            "size": file.size,
+                            "mime_type": file.mime_type,
+                            "mod_time": file.mod_time,
+                            "suffix": file.real_suffix,
+                        }
+                    )
+            if update_values:
+                session.bulk_update_mappings(self.FileEntryModel, update_values)  # type: ignore
+                session.commit()
+
+    def get_exists(self, files: list[FileItem]) -> set[FileItem]:
+        """Get file entries from the table that exist among the given files.
+
+        Args:
+            files: List of file entries
+
+        Returns:
+            Set of FileItem instances whose 'path_no_remote' exists in the table.
+        """
+        # Extract unique paths from the input files.
+        paths = {file.path_no_remote for file in files}
+
+        with Session(self.engine) as session:
+            # Execute a single query to fetch all file paths in the table that match the input paths.
+            result = session.exec(
+                select(self.FileEntryModel.path).where(
+                    self.FileEntryModel.path.in_(paths)  # type: ignore
+                )
+            ).all()
+            # Convert the result to a set for fast membership tests.
+            existing_paths = set(result)
+
+        # Return the set of FileItem objects that have a path in the existing_paths.
+        return {file for file in files if file.path_no_remote in existing_paths}
+
+    def get_all_files(self) -> list[FileItem]:
+        """Get all files in the table.
+
+        Returns:
+            list: List of file entries
+        """
+        # with Session(self.engine) as session:
+        #     return session.exec(select(self.FileEntryModel)).all()
+        out: list[FileItem] = []
+        with Session(self.engine) as session:
+            query = session.exec(select(self.FileEntryModel)).all()
+            for item in query:
+                name = item.name  # type: ignore
+                size = item.size  # type: ignore
+                mime_type = item.mime_type  # type: ignore
+                mod_time = item.mod_time  # type: ignore
+                path = item.path  # type: ignore
+                parent = os.path.dirname(path)
+                if parent == "/" or parent == ".":
+                    parent = ""
+                o = FileItem(
+                    remote=self.remote_name,
+                    parent=parent,
+                    name=name,
+                    size=size,
+                    mime_type=mime_type,
+                    mod_time=mod_time,
+                )
+                out.append(o)
+        return out

rclone_api/db/models.py

@@ -0,0 +1,57 @@
+"""
+Database models for rclone_api.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Optional, Type
+
+from sqlalchemy import BigInteger, Column
+from sqlmodel import Field, SQLModel
+
+
+# Meta table that indexes all repositories
+class RepositoryMeta(SQLModel, table=True):
+    """Repository metadata table."""
+
+    id: Optional[int] = Field(default=None, primary_key=True)
+    repo_name: str
+    file_table_name: str  # The dedicated table name for file entries
+
+
+# Base FileEntry model that will be extended
+class FileEntry(SQLModel, ABC):
+    """Base file entry model with common fields."""
+
+    id: Optional[int] = Field(default=None, primary_key=True)
+    path: str = Field(index=True, unique=True)
+    suffix: str = Field(index=True)
+    name: str
+    size: int = Field(sa_column=Column(BigInteger))
+    mime_type: str
+    mod_time: str
+    hash: Optional[str] = Field(default=None)
+
+    @abstractmethod
+    def table_name(self) -> str:
+        """Return the table name for this file entry model."""
+        pass
+
+
+# Factory to dynamically create a FileEntry model with a given table name
+def create_file_entry_model(_table_name: str) -> Type[FileEntry]:
+    """Create a file entry model with a given table name.
+
+    Args:
+        table_name: Table name
+
+    Returns:
+        Type[FileEntryBase]: File entry model class with specified table name
+    """
+
+    class FileEntryConcrete(FileEntry, table=True):
+        __tablename__ = _table_name  # type: ignore # dynamically set table name
+
+        def table_name(self) -> str:
+            return _table_name
+
+    return FileEntryConcrete

rclone_api/deprecated.py

@@ -0,0 +1,24 @@
+import functools
+import warnings
+
+
+def deprecated(new_func_name: str):
+    """Decorator to mark functions as deprecated.
+
+    Args:
+        new_func_name: The name of the function that should be used instead.
+    """
+
+    def decorator(func):
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            warnings.warn(
+                f"{func.__name__}() is deprecated; use {new_func_name}() instead.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator

rclone_api/detail/copy_file_parts_resumable.py

@@ -0,0 +1,42 @@
+from rclone_api.rclone_impl import RcloneImpl
+from rclone_api.types import (
+    PartInfo,
+)
+
+
+def copy_file_parts_resumable(
+    self: RcloneImpl,
+    src: str,  # src:/Bucket/path/myfile.large.zst
+    dst_dir: str,  # dst:/Bucket/path/myfile.large.zst-parts/
+    part_infos: list[PartInfo] | None = None,
+    upload_threads: int = 10,
+    merge_threads: int = 5,
+    verbose: bool | None = None,
+) -> Exception | None:
+    # _upload_parts
+    from rclone_api.s3.multipart.upload_parts_resumable import upload_parts_resumable
+    from rclone_api.s3.multipart.upload_parts_server_side_merge import (
+        s3_server_side_multi_part_merge,
+    )
+
+    if verbose is None:
+        verbose = self.get_verbose()
+
+    err: Exception | None = upload_parts_resumable(
+        self=self,
+        src=src,
+        dst_dir=dst_dir,
+        part_infos=part_infos,
+        threads=upload_threads,
+    )
+    if isinstance(err, Exception):
+        return err
+    if dst_dir.endswith("/"):
+        dst_dir = dst_dir[:-1]
+    dst_info = f"{dst_dir}/info.json"
+    err = s3_server_side_multi_part_merge(
+        rclone=self, info_path=dst_info, max_workers=merge_threads, verbose=verbose
+    )
+    if isinstance(err, Exception):
+        return err
+    return None

rclone_api/detail/walk.py

@@ -0,0 +1,116 @@
+import random
+from queue import Queue
+from threading import Thread
+from typing import Generator
+
+from rclone_api import Dir
+from rclone_api.dir_listing import DirListing
+from rclone_api.remote import Remote
+from rclone_api.types import Order
+
+_MAX_OUT_QUEUE_SIZE = 50
+
+
+def walk_runner_breadth_first(
+    dir: Dir,
+    max_depth: int,
+    out_queue: Queue[DirListing | None],
+    order: Order = Order.NORMAL,
+) -> None:
+    queue: Queue[Dir] = Queue()
+    queue.put(dir)
+    try:
+        while not queue.empty():
+            current_dir = queue.get()
+            dirlisting = current_dir.ls(max_depth=0, order=order)
+            out_queue.put(dirlisting)
+            dirs = dirlisting.dirs
+
+            if max_depth != 0 and len(dirs) > 0:
+                for child in dirs:
+                    queue.put(child)
+            if max_depth < 0:
+                continue
+            if max_depth > 0:
+                max_depth -= 1
+        out_queue.put(None)
+    except KeyboardInterrupt:
+        import _thread
+
+        out_queue.put(None)
+
+        _thread.interrupt_main()
+
+
+def walk_runner_depth_first(
+    dir: Dir,
+    max_depth: int,
+    out_queue: Queue[DirListing | None],
+    order: Order = Order.NORMAL,
+) -> None:
+    try:
+        stack = [(dir, max_depth)]
+        while stack:
+            current_dir, depth = stack.pop()
+            dirlisting = current_dir.ls()
+            if order == Order.REVERSE:
+                dirlisting.dirs.reverse()
+            if order == Order.RANDOM:
+
+                random.shuffle(dirlisting.dirs)
+            if depth != 0:
+                for subdir in dirlisting.dirs:  # Process deeper directories first
+                    # stack.append((child, depth - 1 if depth > 0 else depth))
+                    next_depth = depth - 1 if depth > 0 else depth
+                    walk_runner_depth_first(subdir, next_depth, out_queue, order=order)
+            out_queue.put(dirlisting)
+        out_queue.put(None)
+    except KeyboardInterrupt:
+        import _thread
+
+        out_queue.put(None)
+        _thread.interrupt_main()
+
+
+def walk(
+    dir: Dir | Remote,
+    breadth_first: bool,
+    max_depth: int = -1,
+    order: Order = Order.NORMAL,
+) -> Generator[DirListing, None, None]:
+    """Walk through the given directory recursively.
+
+    Args:
+        dir: Directory or Remote to walk through
+        max_depth: Maximum depth to traverse (-1 for unlimited)
+
+    Yields:
+        DirListing: Directory listing for each directory encountered
+    """
+    try:
+        # Convert Remote to Dir if needed
+        if isinstance(dir, Remote):
+            dir = Dir(dir)
+        out_queue: Queue[DirListing | None] = Queue(maxsize=_MAX_OUT_QUEUE_SIZE)
+
+        def _task() -> None:
+            if breadth_first:
+                walk_runner_breadth_first(dir, max_depth, out_queue, order)
+            else:
+                walk_runner_depth_first(dir, max_depth, out_queue, order)
+
+        # Start worker thread
+        worker = Thread(
+            target=_task,
+            daemon=True,
+        )
+        worker.start()
+
+        while dirlisting := out_queue.get():
+            if dirlisting is None:
+                break
+            yield dirlisting
+
+        worker.join()
+    except KeyboardInterrupt:
+        pass

rclone_api/diff.py

@@ -0,0 +1,164 @@
+from dataclasses import dataclass
+from enum import Enum
+from queue import Queue
+from threading import Thread
+from typing import Generator
+
+from rclone_api.process import Process
+
+
+class DiffType(Enum):
+    EQUAL = "="
+    MISSING_ON_SRC = (
+        "-"  # means path was missing on the source, so only in the destination
+    )
+    MISSING_ON_DST = (
+        "+"  # means path was missing on the destination, so only in the source
+    )
+    DIFFERENT = "*"  # means path was present in source and destination but different.
+    ERROR = "!"  # means there was an error
+
+
+class DiffOption(Enum):
+    COMBINED = "combined"
+    MISSING_ON_SRC = "missing-on-src"
+    MISSING_ON_DST = "missing-on-dst"
+    DIFFER = "differ"
+    MATCH = "match"
+    ERROR = "error"
+
+
+@dataclass
+class DiffItem:
+    type: DiffType
+    path: str
+    src_prefix: str
+    dst_prefix: str
+
+    def __str__(self) -> str:
+        return f"{self.type.value} {self.path}"
+
+    def __repr__(self) -> str:
+        return f"{self.type.name} {self.path}"
+
+    def full_str(self) -> str:
+        return f"{self.type.name} {self.src_prefix}/{self.path} {self.dst_prefix}/{self.path}"
+
+    def dst_path(self) -> str:
+        return f"{self.dst_prefix}/{self.path}"
+
+    def src_path(self) -> str:
+        return f"{self.src_prefix}/{self.path}"
+
+
+def _parse_missing_on_src_dst(line: str) -> str | None:
+    if line.endswith("does-not-exist"):
+        # 2025/02/17 14:43:38 ERROR : zachs_video/breaking_ai_mind.mp4: file not in S3 bucket rclone-api-unit-test path does-not-exist
+        parts = line.split(" : ", 1)
+        if len(parts) < 1:
+            return None
+        right = parts[1]
+        file_path = right.split(":", 1)[0]
+        return file_path.strip()
+    return None
+
+
+def _classify_diff(
+    line: str, src_slug: str, dst_slug: str, diff_option: DiffOption
+) -> DiffItem | None:
+    def _new(type: DiffType, path: str) -> DiffItem:
+        return DiffItem(type, path, src_prefix=src_slug, dst_prefix=dst_slug)
+
+    if diff_option == DiffOption.COMBINED:
+        suffix = line[1:].strip() if len(line) > 0 else ""
+        if line.startswith(DiffType.EQUAL.value):
+            return _new(DiffType.EQUAL, suffix)
+        if line.startswith(DiffType.MISSING_ON_SRC.value):
+            return _new(DiffType.MISSING_ON_SRC, suffix)
+        if line.startswith(DiffType.MISSING_ON_DST.value):
+            return _new(DiffType.MISSING_ON_DST, suffix)
+        if line.startswith(DiffType.DIFFERENT.value):
+            return _new(DiffType.DIFFERENT, suffix)
+        if line.startswith(DiffType.ERROR.value):
+            return _new(DiffType.ERROR, suffix)
+        return None
+    if diff_option == DiffOption.MISSING_ON_SRC:
+        filename_src: str | None = _parse_missing_on_src_dst(line)
+        if filename_src is not None:
+            return _new(DiffType.MISSING_ON_SRC, filename_src)
+        return None
+    if diff_option == DiffOption.MISSING_ON_DST:
+        filename_dst: str | None = _parse_missing_on_src_dst(line)
+        if filename_dst is not None:
+            return _new(DiffType.MISSING_ON_DST, filename_dst)
+        return None
+    else:
+        raise ValueError(f"Unknown diff_option: {diff_option}")
+
+
+def _async_diff_stream_from_running_process(
+    running_process: Process,
+    src_slug: str,
+    dst_slug: str,
+    diff_option: DiffOption,
+    output: Queue[DiffItem | None],
+) -> None:
+    count = 0
+    first_few_lines: list[str] = []
+    try:
+        assert running_process.stdout is not None
+        n_max = 10
+        for line in iter(running_process.stdout.readline, b""):
+            try:
+                line_str = line.decode("utf-8").strip()
+                if len(first_few_lines) < n_max:
+                    first_few_lines.append(line_str)
+                # _classify_line_type
+                diff_item: DiffItem | None = _classify_diff(
+                    line_str, src_slug, dst_slug, diff_option
+                )
+                if diff_item is None:
+                    # Some other output that we don't care about, debug print etc.
+                    continue
+                output.put(diff_item)
+                count += 1
+                # print(f"unhandled: {line_str}")
+            except UnicodeDecodeError:
+                print("UnicodeDecodeError")
+                continue
+    except KeyboardInterrupt:
+        import _thread
+
+        print("KeyboardInterrupt")
+        _thread.interrupt_main()
+    except Exception as e:
+        import _thread
+
+        print(f"Error: {e}")
+        _thread.interrupt_main()
+    finally:
+        output.put(None)
+
+
+def diff_stream_from_running_process(
+    running_process: Process,
+    src_slug: str,
+    dst_slug: str,
+    diff_option: DiffOption,
+) -> Generator[DiffItem, None, None]:
+    output: Queue[DiffItem | None] = Queue()
+    # process_output_to_diff_stream(running_process, src_slug, dst_slug, output)
+
+    def _task() -> None:
+        _async_diff_stream_from_running_process(
+            running_process, src_slug, dst_slug, diff_option, output
+        )
+
+    thread = Thread(target=_task, daemon=True)
+    thread.start()
+    while True:
+        item = output.get()
+        if item is None:
+            break
+        yield item
+    thread.join(timeout=5)