123 lines
4.4 KiB
Python
123 lines
4.4 KiB
Python
"""Google Drive client for uploading .docx files as Google Docs."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
from pathlib import Path
|
|
|
|
from googleapiclient.discovery import build
|
|
from googleapiclient.http import MediaFileUpload
|
|
|
|
from .google_auth import get_credentials
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
_FOLDER_CACHE_FILE = Path("data/drive_folder_cache.json")
|
|
_GOOGLE_DOC_MIME = "application/vnd.google-apps.document"
|
|
_FOLDER_MIME = "application/vnd.google-apps.folder"
|
|
|
|
|
|
class DriveClient:
|
|
"""Upload .docx files to Google Drive as Google Docs with shareable links."""
|
|
|
|
def __init__(self, root_folder_id: str, cache_file: Path | str = _FOLDER_CACHE_FILE):
|
|
self._root_folder_id = root_folder_id
|
|
self._cache_file = Path(cache_file)
|
|
self._folder_cache = self._load_cache()
|
|
creds = get_credentials()
|
|
self._service = build("drive", "v3", credentials=creds)
|
|
|
|
def _load_cache(self) -> dict[str, str]:
|
|
"""Load folder ID cache from disk."""
|
|
if self._cache_file.exists():
|
|
try:
|
|
return json.loads(self._cache_file.read_text())
|
|
except (json.JSONDecodeError, OSError):
|
|
log.warning("Corrupt drive folder cache, starting fresh")
|
|
return {}
|
|
|
|
def _save_cache(self) -> None:
|
|
"""Persist folder ID cache to disk."""
|
|
self._cache_file.parent.mkdir(parents=True, exist_ok=True)
|
|
self._cache_file.write_text(json.dumps(self._folder_cache, indent=2))
|
|
|
|
def ensure_client_folder(self, company_name: str) -> str:
|
|
"""Get or create a subfolder for the client under the root folder.
|
|
|
|
Returns the folder ID.
|
|
"""
|
|
cache_key = company_name.strip().lower()
|
|
if cache_key in self._folder_cache:
|
|
return self._folder_cache[cache_key]
|
|
|
|
# Check if folder already exists in Drive
|
|
query = (
|
|
f"name = '{company_name}' and "
|
|
f"'{self._root_folder_id}' in parents and "
|
|
f"mimeType = '{_FOLDER_MIME}' and trashed = false"
|
|
)
|
|
results = self._service.files().list(q=query, fields="files(id, name)").execute()
|
|
files = results.get("files", [])
|
|
|
|
if files:
|
|
folder_id = files[0]["id"]
|
|
log.info("Found existing Drive folder for '%s': %s", company_name, folder_id)
|
|
else:
|
|
# Create new folder
|
|
metadata = {
|
|
"name": company_name,
|
|
"mimeType": _FOLDER_MIME,
|
|
"parents": [self._root_folder_id],
|
|
}
|
|
folder = self._service.files().create(body=metadata, fields="id").execute()
|
|
folder_id = folder["id"]
|
|
log.info("Created Drive folder for '%s': %s", company_name, folder_id)
|
|
|
|
self._folder_cache[cache_key] = folder_id
|
|
self._save_cache()
|
|
return folder_id
|
|
|
|
def upload_as_google_doc(
|
|
self, local_path: Path | str, folder_id: str
|
|
) -> tuple[str, str]:
|
|
"""Upload a .docx file to Google Drive, converting to Google Doc.
|
|
|
|
Returns (file_id, web_view_link).
|
|
"""
|
|
local_path = Path(local_path)
|
|
if not local_path.exists():
|
|
raise FileNotFoundError("File not found: %s" % local_path)
|
|
|
|
metadata = {
|
|
"name": local_path.stem,
|
|
"parents": [folder_id],
|
|
"mimeType": _GOOGLE_DOC_MIME,
|
|
}
|
|
media = MediaFileUpload(
|
|
str(local_path),
|
|
mimetype="application/vnd.openxmlformats-officedocument.wordprocessingml.document",
|
|
)
|
|
result = (
|
|
self._service.files()
|
|
.create(body=metadata, media_body=media, fields="id, webViewLink")
|
|
.execute()
|
|
)
|
|
file_id = result["id"]
|
|
web_link = result["webViewLink"]
|
|
log.info("Uploaded '%s' as Google Doc: %s", local_path.name, web_link)
|
|
return file_id, web_link
|
|
|
|
def set_link_sharing(self, file_id: str, role: str = "commenter") -> None:
|
|
"""Make a file accessible to anyone with the link.
|
|
|
|
Args:
|
|
file_id: Google Drive file ID.
|
|
role: Permission role - 'reader', 'commenter', or 'writer'.
|
|
"""
|
|
permission = {"type": "anyone", "role": role}
|
|
self._service.permissions().create(
|
|
fileId=file_id, body=permission, fields="id"
|
|
).execute()
|
|
log.info("Set link sharing (role=%s) on file %s", role, file_id)
|