drive/drive_client.py

"""Google Drive access: authentication, folders, identity lookup, resumable uploads.

Metadata calls use the official ``google-api-python-client``. File uploads use
Google's documented resumable-upload protocol directly over
``google.auth.transport.requests.AuthorizedSession`` so that the session URI can be
persisted in SQLite and an interrupted upload can continue after a process restart.
https://developers.google.com/drive/api/guides/manage-uploads#resumable
"""
from __future__ import annotations

import json
import logging
import os
import random
import re
import threading
import time
from pathlib import Path
from typing import Any, Callable
from urllib.parse import quote

from errors import (
    AuthError,
    BotError,
    DriveQuotaError,
    Interrupted,
    PermanentError,
    TransientError,
    UploadSessionExpired,
)
from models import DriveFile
from naming import normalize_key

log = logging.getLogger(__name__)

FOLDER_MIME = "application/vnd.google-apps.folder"
SCOPES = {
    "drive.file": "https://www.googleapis.com/auth/drive.file",
    "drive": "https://www.googleapis.com/auth/drive",
}
UPLOAD_URL = "https://www.googleapis.com/upload/drive/v3/files"
CHUNK_ALIGN = 256 * 1024
FILE_FIELDS = "id,name,size,md5Checksum,parents,mimeType,trashed,appProperties"
# appProperties key holding "<channel_id>:<message_id>"; used to recognise our own
# uploads after a crash, so a lost response never leads to a duplicate upload.
SOURCE_PROPERTY = "tg_source"

_RATE_LIMIT_REASONS = {"userRateLimitExceeded", "rateLimitExceeded", "sharingRateLimitExceeded"}
_QUOTA_REASONS = {"storageQuotaExceeded", "teamDriveFileLimitExceeded", "activeItemCreationLimitExceeded"}


def _network_errors() -> tuple[type[BaseException], ...]:
    errors: list[type[BaseException]] = [OSError, TimeoutError]  # requests' exceptions subclass OSError
    try:
        import httplib2

        errors.append(httplib2.HttpLib2Error)
    except ImportError:  # pragma: no cover
        pass
    try:
        from google.auth.exceptions import TransportError

        errors.append(TransportError)
    except ImportError:  # pragma: no cover
        pass
    return tuple(errors)


NETWORK_ERRORS = _network_errors()


def escape_query(value: str) -> str:
    """Escape a literal for the Drive ``q`` search syntax."""
    return value.replace("\\", "\\\\").replace("'", "\\'")


def backoff_delay(attempt: int, base: float, cap: float = 64.0) -> float:
    return min(base * (2 ** (attempt - 1)), cap) * random.uniform(0.8, 1.2)


def parse_error_payload(payload: Any) -> tuple[str | None, str | None]:
    """Extract (reason, message) from a Google JSON error body."""
    if isinstance(payload, (bytes, bytearray)):
        payload = payload.decode("utf-8", "replace")
    if isinstance(payload, str):
        try:
            payload = json.loads(payload)
        except ValueError:
            return None, payload[:200] or None
    if not isinstance(payload, dict):
        return None, None
    error = payload.get("error")
    if not isinstance(error, dict):
        return None, str(error) if error else None
    reason = None
    details = error.get("errors")
    if isinstance(details, list) and details and isinstance(details[0], dict):
        reason = details[0].get("reason")
    return reason or error.get("status"), error.get("message")


def classify_http_error(
    status: int, reason: str | None = None, message: str | None = None, retry_after: float | None = None
) -> BotError:
    text = f"Google Drive API error {status}" + (f" ({reason})" if reason else "") + (f": {message}" if message else "")
    if status == 401:
        return AuthError(text, public="Google authorization expired or was revoked; run `python main.py drive-auth`.")
    if status == 403 and reason in _QUOTA_REASONS:
        return DriveQuotaError(
            text,
            public="Google Drive storage quota exceeded (service accounts have no My Drive quota; use a Shared Drive).",
        )
    if status == 403 and reason == "dailyLimitExceeded":
        return TransientError(text, retry_after=retry_after or 3600, rate_limited=True, public="Drive daily API limit reached")
    if status == 429 or (status == 403 and reason in _RATE_LIMIT_REASONS):
        return TransientError(text, retry_after=retry_after, rate_limited=True, public="Google Drive rate limit")
    if status in (408, 500, 502, 503, 504) or status >= 500:
        return TransientError(text, retry_after=retry_after, public=f"Google Drive temporary error {status}")
    return PermanentError(text, public=f"Google Drive rejected the request ({status}{', ' + reason if reason else ''})")


# -- authentication --------------------------------------------------------------------


def load_credentials(config):
    """Build Google credentials without needing files on disk when env vars are set.

    Service account: GOOGLE_SERVICE_ACCOUNT_JSON (parsed in memory) or a key file.
    OAuth (personal Drive): GOOGLE_OAUTH_TOKEN_JSON (authorized-user JSON containing
    the refresh token) or the local token file written by `main.py drive-auth`.
    """
    from config import load_json_secret

    scopes = [SCOPES[config.drive_scope]]
    if config.google_auth_mode == "service_account":
        from google.oauth2 import service_account

        if config.google_service_account_json:
            info = load_json_secret(config.google_service_account_json)
            return service_account.Credentials.from_service_account_info(info, scopes=scopes)
        return service_account.Credentials.from_service_account_file(str(config.google_service_account_file), scopes=scopes)

    from google.auth.exceptions import RefreshError
    from google.auth.transport.requests import Request
    from google.oauth2.credentials import Credentials

    from_env = bool(config.google_oauth_token_json)
    if from_env:
        creds = Credentials.from_authorized_user_info(load_json_secret(config.google_oauth_token_json), scopes)
    else:
        token_file: Path = config.google_token_file
        if not token_file.exists():
            raise AuthError(
                "Google token not found", public="Google Drive is not authorized yet; run `python main.py drive-auth`."
            )
        creds = Credentials.from_authorized_user_file(str(token_file), scopes)
    if not creds.valid:
        if not creds.refresh_token:
            raise AuthError("Google token has no refresh token", public="Run `python main.py drive-auth` again.")
        try:
            creds.refresh(Request())
        except RefreshError as exc:
            raise AuthError(
                f"Google token refresh failed: {exc}",
                public="Google authorization expired or was revoked; run `python main.py drive-auth` and update the token.",
            ) from exc
        if not from_env:
            save_token(creds, config.google_token_file)
    return creds


def save_token(creds, token_file: Path) -> None:
    token_file.parent.mkdir(parents=True, exist_ok=True)
    fd = os.open(token_file, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
    with os.fdopen(fd, "w") as fh:
        fh.write(creds.to_json())


def run_oauth_flow(config, *, open_browser: bool = True) -> Path:
    from google_auth_oauthlib.flow import InstalledAppFlow

    flow = InstalledAppFlow.from_client_secrets_file(str(config.google_client_secrets_file), [SCOPES[config.drive_scope]])
    creds = flow.run_local_server(port=0, open_browser=open_browser)
    save_token(creds, config.google_token_file)
    return config.google_token_file


# -- client ------------------------------------------------------------------------------


class DriveClient:
    def __init__(
        self,
        service,
        session,
        *,
        chunk_size: int,
        shared_drive_id: str | None = None,
        http_retries: int = 5,
        base_delay: float = 2.0,
        sleep: Callable[[float], None] = time.sleep,
    ):
        self._service = service
        self._shared_drive_id = shared_drive_id
        self._retries = http_retries
        self._base_delay = base_delay
        self._sleep = sleep
        self._lock = threading.Lock()  # httplib2 (used by the discovery client) is not thread-safe
        self.uploader = ResumableUploader(
            session, chunk_size=chunk_size, retries=http_retries, base_delay=base_delay, sleep=sleep
        )

    @classmethod
    def from_config(cls, config) -> "DriveClient":
        from google.auth.transport.requests import AuthorizedSession
        from googleapiclient.discovery import build

        creds = load_credentials(config)
        service = build("drive", "v3", credentials=creds, cache_discovery=False)
        return cls(
            service,
            AuthorizedSession(creds),
            chunk_size=config.upload_chunk_size,
            shared_drive_id=config.shared_drive_id,
            http_retries=config.http_retries,
        )

    # -- low level ------------------------------------------------------------------------

    def _execute(self, build_request: Callable[[], Any], *, retries: int | None = None, not_found_ok: bool = False):
        from google.auth.exceptions import RefreshError
        from googleapiclient.errors import HttpError

        retries = self._retries if retries is None else retries
        attempt = 0
        while True:
            attempt += 1
            try:
                with self._lock:
                    return build_request().execute()
            except HttpError as exc:
                status = int(exc.resp.status)
                if status == 404 and not_found_ok:
                    return None
                retry_after = exc.resp.get("retry-after")
                error = classify_http_error(
                    status, *parse_error_payload(exc.content), retry_after=float(retry_after) if retry_after else None
                )
            except RefreshError as exc:
                raise AuthError(
                    f"Google token refresh failed: {exc}",
                    public="Google authorization expired or was revoked; run `python main.py drive-auth`.",
                ) from exc
            except NETWORK_ERRORS as exc:
                error = TransientError(f"network error talking to Google Drive: {exc}", public="Network error (Drive)")
            if not isinstance(error, TransientError) or attempt > retries:
                raise error
            delay = error.retry_after or backoff_delay(attempt, self._base_delay)
            log.warning("Drive request failed (%s); retry %d/%d in %.1fs", error, attempt, retries, delay)
            self._sleep(delay)

    def _list(self, query: str, *, order_by: str | None = None) -> list[DriveFile]:
        files: list[DriveFile] = []
        page_token = None
        while True:
            kwargs: dict[str, Any] = dict(
                q=query,
                fields=f"nextPageToken, files({FILE_FIELDS})",
                pageSize=100,
                supportsAllDrives=True,
                includeItemsFromAllDrives=True,
            )
            if order_by:
                kwargs["orderBy"] = order_by
            if self._shared_drive_id:
                kwargs.update(corpora="drive", driveId=self._shared_drive_id)
            if page_token:
                kwargs["pageToken"] = page_token
            response = self._execute(lambda: self._service.files().list(**kwargs))
            files.extend(DriveFile.from_api(f) for f in response.get("files", []))
            page_token = response.get("nextPageToken")
            if not page_token:
                return files

    # -- queries ------------------------------------------------------------------------

    def get_file(self, file_id: str) -> DriveFile | None:
        data = self._execute(
            lambda: self._service.files().get(fileId=file_id, fields=FILE_FIELDS, supportsAllDrives=True),
            not_found_ok=True,
        )
        return DriveFile.from_api(data) if data else None

    def find_folders(self, parent_id: str, name: str) -> list[DriveFile]:
        query = (
            f"'{escape_query(parent_id)}' in parents and name = '{escape_query(name)}' "
            f"and mimeType = '{FOLDER_MIME}' and trashed = false"
        )
        return self._list(query, order_by="createdTime")

    def find_by_name(self, parent_id: str, name: str) -> list[DriveFile]:
        query = (
            f"'{escape_query(parent_id)}' in parents and name = '{escape_query(name)}' "
            f"and mimeType != '{FOLDER_MIME}' and trashed = false"
        )
        return self._list(query)

    def find_by_source_key(self, key: str) -> list[DriveFile]:
        query = (
            f"appProperties has {{ key='{SOURCE_PROPERTY}' and value='{escape_query(key)}' }} and trashed = false"
        )
        return self._list(query)

    # -- mutations ------------------------------------------------------------------------

    def create_folder(self, parent_id: str, name: str) -> DriveFile:
        # Single attempt: folder creation is not idempotent. FolderResolver searches
        # again before retrying, so a lost response cannot create a duplicate.
        body = {"name": name, "mimeType": FOLDER_MIME, "parents": [parent_id]}
        data = self._execute(
            lambda: self._service.files().create(body=body, fields=FILE_FIELDS, supportsAllDrives=True), retries=0
        )
        return DriveFile.from_api(data)

    def share_with_link(self, file_id: str) -> None:
        body = {"type": "anyone", "role": "reader"}
        self._execute(
            lambda: self._service.permissions().create(fileId=file_id, body=body, fields="id", supportsAllDrives=True)
        )

    def upload_file(
        self,
        path: Path,
        *,
        name: str,
        parent_id: str,
        mime_type: str,
        size: int,
        app_properties: dict[str, str],
        description: str | None = None,
        session_uri: str | None,
        on_session: Callable[[str], None],
        on_progress: Callable[[int, int], None],
        should_stop: Callable[[], bool],
    ) -> DriveFile:
        metadata: dict[str, Any] = {"name": name, "parents": [parent_id], "appProperties": app_properties}
        if description:
            metadata["description"] = description
        data = self.uploader.upload(
            path,
            metadata=metadata,
            mime_type=mime_type,
            size=size,
            session_uri=session_uri,
            on_session=on_session,
            on_progress=on_progress,
            should_stop=should_stop,
        )
        return DriveFile.from_api(data)

    def check_access(self) -> None:
        """Cheap authenticated call used at startup to validate credentials."""
        self._execute(lambda: self._service.about().get(fields="user(displayName)"))

    def check_root_folder(self, folder_id: str) -> DriveFile:
        folder = self.get_file(folder_id)
        if folder is None:
            raise PermanentError(
                f"Drive root folder {folder_id} not found",
                public=(
                    "DRIVE_ROOT_FOLDER_ID is not visible to these credentials. With GOOGLE_DRIVE_SCOPE=drive.file the "
                    "bot only sees folders it created: run `python main.py drive-init`, or set GOOGLE_DRIVE_SCOPE=drive "
                    "and re-run `python main.py drive-auth`. For a service account, share the folder with it."
                ),
            )
        if folder.mime_type != FOLDER_MIME or folder.trashed:
            raise PermanentError(f"Drive root {folder_id} is not an active folder", public="DRIVE_ROOT_FOLDER_ID is not a usable folder")
        return folder


# -- resumable upload ---------------------------------------------------------------------


def parse_range(header: str | None) -> int:
    """'bytes=0-1048575' -> 1048576 bytes persisted; no header -> 0."""
    if not header:
        return 0
    match = re.search(r"(\d+)-(\d+)", header)
    return int(match.group(2)) + 1 if match else 0


def _response_error(resp) -> BotError:
    try:
        payload = resp.json()
    except ValueError:
        payload = getattr(resp, "text", None)
    retry_after = resp.headers.get("Retry-After") if resp.headers else None
    try:
        retry = float(retry_after) if retry_after else None
    except ValueError:
        retry = None
    return classify_http_error(resp.status_code, *parse_error_payload(payload), retry_after=retry)


class ResumableUploader:
    def __init__(
        self,
        session,
        *,
        chunk_size: int,
        retries: int = 5,
        base_delay: float = 2.0,
        sleep: Callable[[float], None] = time.sleep,
        upload_url: str = UPLOAD_URL,
    ):
        if chunk_size <= 0 or chunk_size % CHUNK_ALIGN:
            raise ValueError("chunk_size must be a positive multiple of 256 KiB")
        self._session = session
        self._chunk_size = chunk_size
        self._retries = retries
        self._base_delay = base_delay
        self._sleep = sleep
        self._upload_url = upload_url

    def upload(
        self,
        path: Path,
        *,
        metadata: dict[str, Any],
        mime_type: str,
        size: int,
        session_uri: str | None,
        on_session: Callable[[str], None],
        on_progress: Callable[[int, int], None],
        should_stop: Callable[[], bool],
    ) -> dict[str, Any]:
        if size <= 0:
            raise PermanentError("refusing to upload an empty file", public="Downloaded file is empty")
        offset = 0
        if session_uri:
            status, value = self.query_status(session_uri, size)
            if status == "done":
                on_progress(size, size)
                return value
            if status == "expired":
                log.info("Saved upload session expired; starting a new one")
                session_uri = None
            else:
                offset = value
                log.info("Resuming upload session at byte %d of %d", offset, size)
        if not session_uri:
            session_uri = self._initiate(metadata, mime_type, size)
            on_session(session_uri)
            offset = 0
        on_progress(offset, size)

        failures = 0
        with open(path, "rb") as fh:
            while True:
                if should_stop():
                    raise Interrupted(f"upload paused at byte {offset} of {size}")
                fh.seek(offset)
                chunk = fh.read(min(self._chunk_size, size - offset))
                if not chunk:
                    status, value = self.query_status(session_uri, size)
                    if status == "done":
                        return value
                    raise TransientError("upload finished locally but Drive did not finalize the file")
                headers = {
                    "Content-Length": str(len(chunk)),
                    "Content-Range": f"bytes {offset}-{offset + len(chunk) - 1}/{size}",
                }
                try:
                    resp = self._session.request("PUT", session_uri, data=chunk, headers=headers, timeout=(30, 600))
                except NETWORK_ERRORS as exc:
                    error: BotError = TransientError(f"network error during upload: {exc}", public="Network error during upload")
                else:
                    if resp.status_code in (200, 201):
                        on_progress(size, size)
                        return resp.json()
                    if resp.status_code == 308:
                        offset = parse_range(resp.headers.get("Range"))
                        failures = 0
                        on_progress(offset, size)
                        continue
                    if resp.status_code in (404, 410):
                        raise UploadSessionExpired("upload session expired", public="Upload session expired; will restart")
                    error = _response_error(resp)
                    if not isinstance(error, TransientError):
                        raise error
                failures += 1
                if failures > self._retries:
                    raise error
                delay = getattr(error, "retry_after", None) or backoff_delay(failures, self._base_delay)
                log.warning("Upload chunk failed (%s); retry %d/%d in %.1fs", error, failures, self._retries, delay)
                self._sleep(delay)
                status, value = self.query_status(session_uri, size)
                if status == "done":
                    on_progress(size, size)
                    return value
                if status == "expired":
                    raise UploadSessionExpired("upload session expired", public="Upload session expired; will restart")
                offset = value

    def query_status(self, session_uri: str, size: int) -> tuple[str, Any]:
        """Ask Drive how much of a resumable session it has: done / partial / expired."""
        attempt = 0
        while True:
            attempt += 1
            try:
                resp = self._session.request(
                    "PUT", session_uri, headers={"Content-Length": "0", "Content-Range": f"bytes */{size}"}, timeout=(30, 120)
                )
            except NETWORK_ERRORS as exc:
                error: BotError = TransientError(f"network error querying upload status: {exc}", public="Network error during upload")
            else:
                if resp.status_code in (200, 201):
                    return "done", resp.json()
                if resp.status_code == 308:
                    return "partial", parse_range(resp.headers.get("Range"))
                if resp.status_code in (404, 410):
                    return "expired", None
                error = _response_error(resp)
                if not isinstance(error, TransientError):
                    raise error
            if attempt > self._retries:
                raise error
            self._sleep(getattr(error, "retry_after", None) or backoff_delay(attempt, self._base_delay))

    def _initiate(self, metadata: dict[str, Any], mime_type: str, size: int) -> str:
        url = f"{self._upload_url}?uploadType=resumable&supportsAllDrives=true&fields={quote(FILE_FIELDS)}"
        headers = {
            "Content-Type": "application/json; charset=UTF-8",
            "X-Upload-Content-Type": mime_type,
            "X-Upload-Content-Length": str(size),
        }
        body = json.dumps(metadata).encode("utf-8")
        attempt = 0
        while True:
            attempt += 1
            try:
                resp = self._session.request("POST", url, data=body, headers=headers, timeout=(30, 120))
            except NETWORK_ERRORS as exc:
                error: BotError = TransientError(f"network error starting upload: {exc}", public="Network error during upload")
            else:
                if resp.status_code == 200:
                    location = resp.headers.get("Location") or resp.headers.get("location")
                    if location:
                        return location
                    error = TransientError("Drive did not return an upload session URI")
                else:
                    error = _response_error(resp)
                    if not isinstance(error, TransientError):
                        raise error
            if attempt > self._retries:
                raise error
            self._sleep(getattr(error, "retry_after", None) or backoff_delay(attempt, self._base_delay))


# -- folder resolution --------------------------------------------------------------------


class FolderResolver:
    """Map folder names to Drive folder IDs without ever creating duplicates.

    Lookup order: SQLite cache (verified once per process) -> Drive search (oldest
    match wins if duplicates already exist) -> create. After a failed create the
    search is repeated before trying again, so a create whose response was lost is
    found instead of duplicated.
    """

    def __init__(self, drive: DriveClient, db, root_id: str, *, attempts: int = 4, sleep=time.sleep):
        self._drive = drive
        self._db = db
        self.root_id = root_id
        self._attempts = attempts
        self._sleep = sleep
        self._verified: set[str] = set()

    def resolve(self, names: list[str]) -> str:
        parent = self.root_id
        for name in names:
            parent = self._ensure(parent, name)
        return parent

    def _ensure(self, parent_id: str, name: str) -> str:
        key = normalize_key(name)
        cached = self._db.get_folder(parent_id, key)
        if cached:
            if cached in self._verified:
                return cached
            folder = self._drive.get_file(cached)
            if folder and not folder.trashed and folder.mime_type == FOLDER_MIME:
                self._verified.add(cached)
                return cached
            log.warning("Cached folder %r (%s) is gone; looking it up again", name, cached)
            self._db.delete_folder(parent_id, key)

        attempt = 0
        while True:
            attempt += 1
            matches = [f for f in self._drive.find_folders(parent_id, name) if normalize_key(f.name) == key]
            if matches:
                if len(matches) > 1:
                    log.warning("%d folders named %r exist; reusing the oldest (%s)", len(matches), name, matches[0].id)
                return self._remember(parent_id, key, name, matches[0].id)
            try:
                created = self._drive.create_folder(parent_id, name)
            except TransientError:
                if attempt >= self._attempts:
                    raise
                self._sleep(backoff_delay(attempt, 2.0))
                continue
            log.info("Created Drive folder %r (%s)", name, created.id)
            return self._remember(parent_id, key, name, created.id)

    def _remember(self, parent_id: str, key: str, name: str, folder_id: str) -> str:
        self._db.put_folder(parent_id, key, name, folder_id)
        self._verified.add(folder_id)
        return folder_id