first commit
This commit is contained in:
gosangam
+620
@@ -0,0 +1,620 @@
|
||||
"""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
|
||||
Reference in New Issue
Block a user