cloudscope 0.1.0__py3-none-any.whl

cloudscope/backends/drive.py

@@ -0,0 +1,568 @@
+"""Google Drive backend implementation."""
+
+from __future__ import annotations
+
+import asyncio
+import io
+from collections import OrderedDict
+from collections.abc import AsyncIterator
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+from googleapiclient.discovery import build
+from googleapiclient.http import MediaFileUpload, MediaIoBaseDownload
+
+from cloudscope.auth.drive_oauth import get_drive_credentials
+from cloudscope.backends.base import (
+    AuthenticationError,
+    CloudScopeError,
+    NetworkError,
+    NotFoundError,
+    ProgressCallback,
+)
+from cloudscope.backends.registry import register_backend
+from cloudscope.models.cloud_file import CloudFile, CloudFileType
+
+# Google Workspace MIME types that need export instead of direct download
+WORKSPACE_MIME_TYPES: dict[str, dict[str, str]] = {
+    "application/vnd.google-apps.document": {
+        "default": "application/pdf",
+        "pdf": "application/pdf",
+        "docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
+        "txt": "text/plain",
+    },
+    "application/vnd.google-apps.spreadsheet": {
+        "default": "text/csv",
+        "csv": "text/csv",
+        "xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
+        "pdf": "application/pdf",
+    },
+    "application/vnd.google-apps.presentation": {
+        "default": "application/pdf",
+        "pdf": "application/pdf",
+        "pptx": "application/vnd.openxmlformats-officedocument.presentationml.presentation",
+    },
+    "application/vnd.google-apps.drawing": {
+        "default": "image/png",
+        "png": "image/png",
+        "pdf": "application/pdf",
+        "svg": "image/svg+xml",
+    },
+}
+
+# File fields to request from Drive API
+FILE_FIELDS = "id, name, mimeType, size, modifiedTime, md5Checksum, parents, trashed"
+LIST_FIELDS = f"nextPageToken, files({FILE_FIELDS})"
+
+# Max items in the path-to-ID cache
+PATH_CACHE_SIZE = 1024
+
+
+class GoogleDriveBackend:
+    """CloudBackend implementation for Google Drive."""
+
+    def __init__(
+        self,
+        client_secrets_path: str | None = None,
+    ) -> None:
+        self._client_secrets_path = client_secrets_path
+        self._service: Any = None
+        self._path_cache: OrderedDict[str, str] = OrderedDict()
+
+    @property
+    def backend_type(self) -> str:
+        return "drive"
+
+    @property
+    def display_name(self) -> str:
+        return "Google Drive"
+
+    # --- Connection ---
+
+    async def connect(self) -> None:
+        try:
+            credentials = await asyncio.to_thread(
+                get_drive_credentials, self._client_secrets_path
+            )
+            self._service = await asyncio.to_thread(
+                build, "drive", "v3", credentials=credentials
+            )
+            # Validate by fetching about info
+            await asyncio.to_thread(
+                lambda: self._service.about().get(fields="user").execute()
+            )
+        except FileNotFoundError:
+            raise
+        except Exception as e:
+            err_str = str(e).lower()
+            if "credential" in err_str or "token" in err_str or "unauthorized" in err_str:
+                raise AuthenticationError(f"Google Drive authentication failed: {e}") from e
+            raise NetworkError(f"Failed to connect to Google Drive: {e}") from e
+
+    async def disconnect(self) -> None:
+        self._service = None
+        self._path_cache.clear()
+
+    async def is_connected(self) -> bool:
+        if self._service is None:
+            return False
+        try:
+            await asyncio.to_thread(
+                lambda: self._service.about().get(fields="user").execute()
+            )
+            return True
+        except Exception:
+            return False
+
+    # --- Container listing ---
+
+    async def list_containers(self) -> list[str]:
+        """List available drives (My Drive + shared drives)."""
+        self._ensure_connected()
+
+        def _list() -> list[str]:
+            drives = ["My Drive"]
+            try:
+                result = self._service.drives().list(pageSize=100).execute()
+                for drive in result.get("drives", []):
+                    drives.append(drive["name"])
+            except Exception:
+                pass  # Shared drives may not be available
+            return drives
+
+        return await asyncio.to_thread(_list)
+
+    # --- Browsing ---
+
+    async def list_files(
+        self,
+        container: str,
+        prefix: str = "",
+        recursive: bool = False,
+    ) -> list[CloudFile]:
+        self._ensure_connected()
+
+        def _list() -> list[CloudFile]:
+            parent_id = self._resolve_parent_id(container, prefix)
+            query = f"'{parent_id}' in parents and trashed = false"
+
+            files: list[CloudFile] = []
+            page_token: str | None = None
+
+            while True:
+                result = (
+                    self._service.files()
+                    .list(
+                        q=query,
+                        fields=LIST_FIELDS,
+                        pageSize=1000,
+                        pageToken=page_token,
+                        orderBy="folder,name",
+                    )
+                    .execute()
+                )
+
+                for item in result.get("files", []):
+                    cloud_file = self._item_to_cloud_file(item, prefix)
+                    files.append(cloud_file)
+
+                    if recursive and cloud_file.is_folder:
+                        # For recursive listing, we'd need to recurse
+                        # but this is handled by list_files_recursive
+                        pass
+
+                page_token = result.get("nextPageToken")
+                if not page_token:
+                    break
+
+            return files
+
+        return await asyncio.to_thread(_list)
+
+    async def stat(self, container: str, path: str) -> CloudFile:
+        self._ensure_connected()
+
+        def _stat() -> CloudFile:
+            file_id = self._resolve_path_to_id(container, path)
+            if not file_id:
+                raise NotFoundError(f"drive://{container}/{path} not found")
+            result = (
+                self._service.files()
+                .get(fileId=file_id, fields=FILE_FIELDS)
+                .execute()
+            )
+            parent_prefix = path.rsplit("/", 1)[0] if "/" in path else ""
+            return self._item_to_cloud_file(result, parent_prefix)
+
+        return await asyncio.to_thread(_stat)
+
+    async def exists(self, container: str, path: str) -> bool:
+        try:
+            await self.stat(container, path)
+            return True
+        except NotFoundError:
+            return False
+
+    # --- Transfer ---
+
+    async def download(
+        self,
+        container: str,
+        remote_path: str,
+        local_path: str,
+        progress_callback: ProgressCallback | None = None,
+    ) -> None:
+        self._ensure_connected()
+
+        def _download() -> None:
+            file_id = self._resolve_path_to_id(container, remote_path)
+            if not file_id:
+                raise NotFoundError(f"drive://{container}/{remote_path} not found")
+
+            # Get file metadata to check MIME type
+            meta = self._service.files().get(fileId=file_id, fields="mimeType,size").execute()
+            mime_type = meta.get("mimeType", "")
+            total_size = int(meta.get("size", 0))
+
+            Path(local_path).parent.mkdir(parents=True, exist_ok=True)
+
+            if mime_type in WORKSPACE_MIME_TYPES:
+                # Export Workspace files
+                export_mime = WORKSPACE_MIME_TYPES[mime_type]["default"]
+                request = self._service.files().export_media(
+                    fileId=file_id, mimeType=export_mime
+                )
+            else:
+                request = self._service.files().get_media(fileId=file_id)
+
+            with open(local_path, "wb") as f:
+                downloader = MediaIoBaseDownload(f, request)
+                done = False
+                while not done:
+                    status, done = downloader.next_chunk()
+                    if progress_callback and status:
+                        transferred = int(status.progress() * total_size) if total_size else 0
+                        progress_callback(transferred, total_size)
+
+            if progress_callback:
+                final_size = Path(local_path).stat().st_size
+                progress_callback(final_size, final_size)
+
+        await asyncio.to_thread(_download)
+
+    async def upload(
+        self,
+        container: str,
+        local_path: str,
+        remote_path: str,
+        progress_callback: ProgressCallback | None = None,
+    ) -> CloudFile:
+        self._ensure_connected()
+
+        def _upload() -> dict:
+            # Resolve parent folder
+            parent_prefix = remote_path.rsplit("/", 1)[0] if "/" in remote_path else ""
+            file_name = remote_path.rsplit("/", 1)[-1]
+            parent_id = self._resolve_parent_id(container, parent_prefix)
+
+            file_metadata: dict[str, Any] = {
+                "name": file_name,
+                "parents": [parent_id],
+            }
+
+            media = MediaFileUpload(local_path, resumable=True)
+            total_size = Path(local_path).stat().st_size
+
+            request = self._service.files().create(
+                body=file_metadata,
+                media_body=media,
+                fields=FILE_FIELDS,
+            )
+
+            response = None
+            while response is None:
+                status, response = request.next_chunk()
+                if progress_callback and status:
+                    transferred = int(status.progress() * total_size)
+                    progress_callback(transferred, total_size)
+
+            if progress_callback:
+                progress_callback(total_size, total_size)
+
+            # Invalidate cache for the parent
+            self._invalidate_cache_prefix(parent_prefix)
+
+            return response
+
+        result = await asyncio.to_thread(_upload)
+        parent_prefix = remote_path.rsplit("/", 1)[0] if "/" in remote_path else ""
+        return self._item_to_cloud_file(result, parent_prefix)
+
+    # --- Mutation ---
+
+    async def delete(self, container: str, path: str) -> None:
+        self._ensure_connected()
+
+        def _delete() -> None:
+            file_id = self._resolve_path_to_id(container, path)
+            if not file_id:
+                raise NotFoundError(f"drive://{container}/{path} not found")
+            self._service.files().delete(fileId=file_id).execute()
+            self._invalidate_cache_prefix(path)
+
+        await asyncio.to_thread(_delete)
+
+    async def create_folder(self, container: str, path: str) -> CloudFile:
+        self._ensure_connected()
+
+        def _create() -> dict:
+            parent_prefix = path.rsplit("/", 1)[0] if "/" in path else ""
+            folder_name = path.rsplit("/", 1)[-1]
+            parent_id = self._resolve_parent_id(container, parent_prefix)
+
+            file_metadata = {
+                "name": folder_name,
+                "mimeType": "application/vnd.google-apps.folder",
+                "parents": [parent_id],
+            }
+            result = (
+                self._service.files()
+                .create(body=file_metadata, fields=FILE_FIELDS)
+                .execute()
+            )
+            self._invalidate_cache_prefix(parent_prefix)
+            return result
+
+        result = await asyncio.to_thread(_create)
+        parent_prefix = path.rsplit("/", 1)[0] if "/" in path else ""
+        return self._item_to_cloud_file(result, parent_prefix)
+
+    async def move(self, container: str, src: str, dst: str) -> CloudFile:
+        self._ensure_connected()
+
+        def _move() -> dict:
+            file_id = self._resolve_path_to_id(container, src)
+            if not file_id:
+                raise NotFoundError(f"drive://{container}/{src} not found")
+
+            # Get current parents
+            file_meta = self._service.files().get(fileId=file_id, fields="parents").execute()
+            previous_parents = ",".join(file_meta.get("parents", []))
+
+            dst_parent_prefix = dst.rsplit("/", 1)[0] if "/" in dst else ""
+            new_name = dst.rsplit("/", 1)[-1]
+            new_parent_id = self._resolve_parent_id(container, dst_parent_prefix)
+
+            result = (
+                self._service.files()
+                .update(
+                    fileId=file_id,
+                    addParents=new_parent_id,
+                    removeParents=previous_parents,
+                    body={"name": new_name},
+                    fields=FILE_FIELDS,
+                )
+                .execute()
+            )
+            self._invalidate_cache_prefix(src)
+            self._invalidate_cache_prefix(dst)
+            return result
+
+        result = await asyncio.to_thread(_move)
+        dst_parent = dst.rsplit("/", 1)[0] if "/" in dst else ""
+        return self._item_to_cloud_file(result, dst_parent)
+
+    async def copy(self, container: str, src: str, dst: str) -> CloudFile:
+        self._ensure_connected()
+
+        def _copy() -> dict:
+            file_id = self._resolve_path_to_id(container, src)
+            if not file_id:
+                raise NotFoundError(f"drive://{container}/{src} not found")
+
+            dst_parent_prefix = dst.rsplit("/", 1)[0] if "/" in dst else ""
+            new_name = dst.rsplit("/", 1)[-1]
+            new_parent_id = self._resolve_parent_id(container, dst_parent_prefix)
+
+            result = (
+                self._service.files()
+                .copy(
+                    fileId=file_id,
+                    body={"name": new_name, "parents": [new_parent_id]},
+                    fields=FILE_FIELDS,
+                )
+                .execute()
+            )
+            self._invalidate_cache_prefix(dst)
+            return result
+
+        result = await asyncio.to_thread(_copy)
+        dst_parent = dst.rsplit("/", 1)[0] if "/" in dst else ""
+        return self._item_to_cloud_file(result, dst_parent)
+
+    # --- Sync support ---
+
+    async def list_files_recursive(
+        self, container: str, prefix: str = ""
+    ) -> AsyncIterator[CloudFile]:
+        files = await self.list_files(container, prefix)
+        for f in files:
+            if f.is_folder:
+                async for child in self.list_files_recursive(container, f.path):
+                    yield child
+            else:
+                yield f
+
+    # --- Path-to-ID resolution ---
+
+    def _resolve_parent_id(self, container: str, prefix: str) -> str:
+        """Resolve a path prefix to a Drive folder ID."""
+        if not prefix:
+            return self._get_root_id(container)
+
+        cached = self._cache_get(f"{container}:{prefix}")
+        if cached:
+            return cached
+
+        # Walk each path component
+        parts = prefix.strip("/").split("/")
+        current_id = self._get_root_id(container)
+
+        for i, part in enumerate(parts):
+            partial_path = "/".join(parts[: i + 1])
+            cached = self._cache_get(f"{container}:{partial_path}")
+            if cached:
+                current_id = cached
+                continue
+
+            query = (
+                f"name = '{_escape_query(part)}' and "
+                f"'{current_id}' in parents and "
+                f"mimeType = 'application/vnd.google-apps.folder' and "
+                f"trashed = false"
+            )
+            result = (
+                self._service.files()
+                .list(q=query, fields="files(id, name)", pageSize=1)
+                .execute()
+            )
+            files = result.get("files", [])
+            if not files:
+                raise NotFoundError(f"Folder not found: {partial_path}")
+            current_id = files[0]["id"]
+            self._cache_set(f"{container}:{partial_path}", current_id)
+
+        return current_id
+
+    def _resolve_path_to_id(self, container: str, path: str) -> str | None:
+        """Resolve a full file path to its Drive file ID."""
+        cached = self._cache_get(f"{container}:{path}")
+        if cached:
+            return cached
+
+        parent_prefix = path.rsplit("/", 1)[0] if "/" in path else ""
+        file_name = path.rsplit("/", 1)[-1]
+
+        try:
+            parent_id = self._resolve_parent_id(container, parent_prefix)
+        except NotFoundError:
+            return None
+
+        query = (
+            f"name = '{_escape_query(file_name)}' and "
+            f"'{parent_id}' in parents and "
+            f"trashed = false"
+        )
+        result = (
+            self._service.files()
+            .list(q=query, fields="files(id)", pageSize=1)
+            .execute()
+        )
+        files = result.get("files", [])
+        if not files:
+            return None
+        file_id = files[0]["id"]
+        self._cache_set(f"{container}:{path}", file_id)
+        return file_id
+
+    def _get_root_id(self, container: str) -> str:
+        """Get the root folder ID for a container."""
+        if container == "My Drive":
+            return "root"
+        # Shared drives
+        cached = self._cache_get(f"__drive__:{container}")
+        if cached:
+            return cached
+        result = (
+            self._service.drives()
+            .list(q=f"name = '{_escape_query(container)}'", pageSize=1)
+            .execute()
+        )
+        drives = result.get("drives", [])
+        if not drives:
+            raise NotFoundError(f"Drive not found: {container}")
+        drive_id = drives[0]["id"]
+        self._cache_set(f"__drive__:{container}", drive_id)
+        return drive_id
+
+    # --- Cache management ---
+
+    def _cache_get(self, key: str) -> str | None:
+        if key in self._path_cache:
+            self._path_cache.move_to_end(key)
+            return self._path_cache[key]
+        return None
+
+    def _cache_set(self, key: str, value: str) -> None:
+        self._path_cache[key] = value
+        self._path_cache.move_to_end(key)
+        while len(self._path_cache) > PATH_CACHE_SIZE:
+            self._path_cache.popitem(last=False)
+
+    def _invalidate_cache_prefix(self, prefix: str) -> None:
+        """Remove all cache entries that start with the given prefix."""
+        to_remove = [k for k in self._path_cache if prefix in k]
+        for k in to_remove:
+            del self._path_cache[k]
+
+    # --- Helpers ---
+
+    def _item_to_cloud_file(self, item: dict, parent_prefix: str) -> CloudFile:
+        """Convert a Drive API file resource to a CloudFile."""
+        mime = item.get("mimeType", "")
+        is_folder = mime == "application/vnd.google-apps.folder"
+        name = item["name"]
+        path = f"{parent_prefix}/{name}".lstrip("/") if parent_prefix else name
+
+        modified = None
+        if "modifiedTime" in item:
+            modified = datetime.fromisoformat(
+                item["modifiedTime"].replace("Z", "+00:00")
+            )
+
+        # Cache the ID
+        self._cache_set(f":{path}", item["id"])
+
+        return CloudFile(
+            name=name,
+            path=path,
+            file_type=CloudFileType.FOLDER if is_folder else CloudFileType.FILE,
+            size=int(item.get("size", 0)),
+            last_modified=modified,
+            checksum=item.get("md5Checksum"),
+            content_type=mime if not is_folder else None,
+            native_id=item["id"],
+        )
+
+    def _ensure_connected(self) -> None:
+        if self._service is None:
+            raise CloudScopeError("Not connected. Call connect() first.")
+
+
+def _escape_query(value: str) -> str:
+    """Escape single quotes in Drive API query strings."""
+    return value.replace("\\", "\\\\").replace("'", "\\'")
+
+
+register_backend("drive", GoogleDriveBackend)