breakpointsoftware/document-parser
0
1from __future__ import annotations2 3import json4import logging5import os6import re7import tempfile8from dataclasses import dataclass9from pathlib import Path10from typing import Any11 12from dotenv import load_dotenv13from google.oauth2.service_account import Credentials14 15from document_processing import SUPPORTED_EXTENSIONS16 17 18load_dotenv()19 20 21logger = logging.getLogger(__name__)22if not logger.handlers:23 handler = logging.StreamHandler()24 handler.setFormatter(logging.Formatter("%(asctime)s %(levelname)s %(name)s: %(message)s"))25 logger.addHandler(handler)26logger.setLevel(logging.INFO)27logger.propagate = False28 29 30DRIVE_READONLY_SCOPE = "https://www.googleapis.com/auth/drive.readonly"31DRIVE_FULL_SCOPE = "https://www.googleapis.com/auth/drive"32FOLDER_MIME_TYPE = "application/vnd.google-apps.folder"33 34# Google Workspace files require export, while binary files can be downloaded directly.35SUPPORTED_EXPORT_MIME_MAP = {36 "application/vnd.google-apps.document": (37 "application/vnd.openxmlformats-officedocument.wordprocessingml.document",38 ".docx",39 ),40}41 42SUPPORTED_BINARY_MIME_TO_SUFFIX = {43 "application/pdf": ".pdf",44 "text/plain": ".txt",45 "application/vnd.openxmlformats-officedocument.wordprocessingml.document": ".docx",46 "image/jpeg": ".jpg",47 "image/png": ".png",48}49 50 51@dataclass52class DriveDocument:53 document_id: str54 source_file: str55 mime_type: str56 modificationDate: str | None57 local_path: str58 59 60@dataclass61class DriveDownloadResult:62 documents: list[DriveDocument]63 discovered_count: int64 temp_dir: str65 66 67class GoogleDriveConfigError(RuntimeError):68 """Raised when required Google Drive environment configuration is missing or invalid."""69 70 71def load_drive_folder_ids_from_env(env_var: str = "GOOGLE_DRIVE_FOLDER_IDS") -> list[str]:72 raw = (os.getenv(env_var) or "").strip()73 if not raw:74 raise GoogleDriveConfigError(f"Missing {env_var}. Set one or more folder IDs separated by commas.")75 76 folder_ids = [item.strip() for item in raw.split(",") if item.strip()]77 if not folder_ids:78 raise GoogleDriveConfigError(f"{env_var} is empty after parsing.")79 80 logger.info("Loaded %s folder id(s) from env var %s", len(folder_ids), env_var)81 82 return folder_ids83 84 85def build_drive_credentials(scope: str = DRIVE_READONLY_SCOPE, service_account_json: str | None = None) -> Credentials:86 """Build Google Drive credentials.87 88 Args:89 scope: OAuth scope for the credentials90 service_account_json: Optional JSON string with service account credentials.91 If not provided, uses GOOGLE_SERVICE_ACCOUNT_JSON env var.92 """93 json_to_use = (service_account_json or "").strip() or (os.getenv("GOOGLE_SERVICE_ACCOUNT_JSON") or "").strip()94 service_account_file = (os.getenv("GOOGLE_SERVICE_ACCOUNT_FILE") or "").strip()95 96 if json_to_use:97 try:98 info: dict[str, Any] = json.loads(json_to_use)99 except json.JSONDecodeError as exc:100 raise GoogleDriveConfigError("GOOGLE_SERVICE_ACCOUNT_JSON is not valid JSON.") from exc101 logger.info("Using GOOGLE_SERVICE_ACCOUNT_JSON for Drive credentials")102 return Credentials.from_service_account_info(info, scopes=[scope])103 104 if service_account_file:105 logger.info("Using GOOGLE_SERVICE_ACCOUNT_FILE for Drive credentials")106 return Credentials.from_service_account_file(service_account_file, scopes=[scope])107 108 raise GoogleDriveConfigError(109 "Missing Google credentials. Set GOOGLE_SERVICE_ACCOUNT_JSON or GOOGLE_SERVICE_ACCOUNT_FILE."110 )111 112 113def build_drive_service(scope: str = DRIVE_READONLY_SCOPE, service_account_json: str | None = None):114 """Build Google Drive service.115 116 Args:117 scope: OAuth scope for the credentials118 service_account_json: Optional JSON string with service account credentials for tenant.119 """120 try:121 from googleapiclient.discovery import build122 except ImportError as exc:123 raise GoogleDriveConfigError(124 "google-api-python-client is required for Drive scanning. Add it to requirements and install dependencies."125 ) from exc126 127 credentials = build_drive_credentials(scope=scope, service_account_json=service_account_json)128 logger.info("Building Drive service with scope=%s", scope)129 return build("drive", "v3", credentials=credentials, cache_discovery=False)130 131 132def _escape_drive_query_value(value: str) -> str:133 return value.replace("'", "\\'")134 135 136def _sanitize_file_name(name: str) -> str:137 cleaned = re.sub(r'[<>:"/\\|?*\x00-\x1F]', "_", name).strip().rstrip(".")138 return cleaned or "document"139 140 141def _ensure_unique_destination(path: Path) -> Path:142 if not path.exists():143 return path144 145 counter = 1146 while True:147 candidate = path.with_name(f"{path.stem}_{counter}{path.suffix}")148 if not candidate.exists():149 return candidate150 counter += 1151 152 153def _list_files_from_folder(service, folder_id: str, include_subfolders: bool) -> list[dict[str, Any]]:154 logger.info("Listing Drive folder id=%s include_subfolders=%s", folder_id, include_subfolders)155 folders_to_scan = [folder_id]156 discovered: list[dict[str, Any]] = []157 158 while folders_to_scan:159 current_folder = folders_to_scan.pop(0)160 page_token = None161 162 while True:163 response = (164 service.files()165 .list(166 q=f"'{current_folder}' in parents and trashed = false",167 fields="nextPageToken, files(id, name, mimeType, modifiedTime)",168 supportsAllDrives=True,169 includeItemsFromAllDrives=True,170 pageSize=1000,171 pageToken=page_token,172 )173 .execute()174 )175 176 for file_meta in response.get("files", []):177 mime_type = file_meta.get("mimeType")178 if mime_type == FOLDER_MIME_TYPE:179 if include_subfolders:180 child_folder_id = file_meta.get("id")181 if child_folder_id:182 folders_to_scan.append(child_folder_id)183 continue184 185 discovered.append(file_meta)186 187 page_token = response.get("nextPageToken")188 if not page_token:189 break190 191 return discovered192 193 194def _is_supported_drive_file(file_meta: dict[str, Any]) -> bool:195 mime_type = (file_meta.get("mimeType") or "").strip().lower()196 if mime_type in SUPPORTED_BINARY_MIME_TO_SUFFIX:197 return True198 if mime_type in SUPPORTED_EXPORT_MIME_MAP:199 return True200 201 name = (file_meta.get("name") or "").strip()202 suffix = Path(name).suffix.lower()203 return suffix in SUPPORTED_EXTENSIONS204 205 206def _download_drive_file(service, file_meta: dict[str, Any], destination_dir: Path) -> str:207 try:208 from googleapiclient.http import MediaIoBaseDownload209 except ImportError as exc:210 raise GoogleDriveConfigError(211 "google-api-python-client is required for Drive scanning. Add it to requirements and install dependencies."212 ) from exc213 214 file_id = (file_meta.get("id") or "").strip()215 file_name = _sanitize_file_name((file_meta.get("name") or "document").strip())216 mime_type = (file_meta.get("mimeType") or "").strip().lower()217 218 suffix = Path(file_name).suffix.lower()219 request = None220 221 if mime_type in SUPPORTED_EXPORT_MIME_MAP:222 export_mime, export_suffix = SUPPORTED_EXPORT_MIME_MAP[mime_type]223 if suffix != export_suffix:224 file_name = f"{Path(file_name).stem}{export_suffix}"225 suffix = export_suffix226 request = service.files().export_media(fileId=file_id, mimeType=export_mime)227 else:228 if not suffix and mime_type in SUPPORTED_BINARY_MIME_TO_SUFFIX:229 file_name = f"{file_name}{SUPPORTED_BINARY_MIME_TO_SUFFIX[mime_type]}"230 suffix = Path(file_name).suffix.lower()231 request = service.files().get_media(fileId=file_id)232 233 if suffix not in SUPPORTED_EXTENSIONS:234 raise ValueError(f"Unsupported downloaded file type: {suffix or '[no extension]'}")235 236 destination = _ensure_unique_destination(destination_dir / file_name)237 238 with destination.open("wb") as stream:239 downloader = MediaIoBaseDownload(stream, request)240 done = False241 while not done:242 _, done = downloader.next_chunk()243 244 return str(destination)245 246 247def scan_drive_supported_documents(folder_ids: list[str], include_subfolders: bool = True) -> DriveDownloadResult:248 service = build_drive_service()249 logger.info("Starting Drive scan for %s folder(s)", len(folder_ids))250 251 discovered_files: list[dict[str, Any]] = []252 for folder_id in folder_ids:253 discovered_files.extend(_list_files_from_folder(service, folder_id, include_subfolders))254 255 supported_files = [file_meta for file_meta in discovered_files if _is_supported_drive_file(file_meta)]256 logger.info("Drive scan discovered=%s supported=%s", len(discovered_files), len(supported_files))257 temp_dir = tempfile.mkdtemp(prefix="drive_docs_")258 destination_dir = Path(temp_dir)259 260 documents: list[DriveDocument] = []261 for file_meta in supported_files:262 try:263 local_path = _download_drive_file(service, file_meta, destination_dir)264 except Exception:265 # Keep scan resilient if one file is malformed or inaccessible.266 logger.exception("Skipping file during Drive download due to error file_id=%s", file_meta.get("id"))267 continue268 269 documents.append(270 DriveDocument(271 document_id=str(file_meta.get("id") or "").strip(),272 source_file=str(file_meta.get("name") or "").strip(),273 mime_type=str(file_meta.get("mimeType") or "").strip(),274 modificationDate=str(file_meta.get("modifiedTime") or "").strip() or None,275 local_path=local_path,276 )277 )278 279 return DriveDownloadResult(280 documents=documents,281 discovered_count=len(supported_files),282 temp_dir=temp_dir,283 )284 285 286def _find_folder_in_parent(service, folder_name: str, parent_folder_id: str | None) -> str | None:287 query_parts = [288 "trashed = false",289 f"mimeType = '{FOLDER_MIME_TYPE}'",290 f"name = '{_escape_drive_query_value(folder_name)}'",291 ]292 if parent_folder_id:293 query_parts.append(f"'{parent_folder_id}' in parents")294 295 try:296 response = (297 service.files()298 .list(299 q=" and ".join(query_parts),300 fields="files(id,name)",301 supportsAllDrives=True,302 includeItemsFromAllDrives=True,303 pageSize=10,304 )305 .execute()306 )307 except Exception as exc:308 raise GoogleDriveConfigError(f"Failed to query folder '{folder_name}': {exc}") from exc309 310 files = response.get("files", [])311 if not files:312 logger.info("Drive folder not found name=%s parent=%s", folder_name, parent_folder_id or "root")313 return None314 315 logger.info("Found Drive folder name=%s parent=%s", folder_name, parent_folder_id or "root")316 return str(files[0].get("id") or "").strip() or None317 318 319def _create_folder_in_parent(service, folder_name: str, parent_folder_id: str | None) -> str:320 body: dict[str, Any] = {"name": folder_name, "mimeType": FOLDER_MIME_TYPE}321 if parent_folder_id:322 body["parents"] = [parent_folder_id]323 324 try:325 created = (326 service.files()327 .create(body=body, fields="id", supportsAllDrives=True)328 .execute()329 )330 except Exception as exc:331 raise GoogleDriveConfigError(332 f"Failed to create folder '{folder_name}' under parent '{parent_folder_id or 'root'}': {exc}"333 ) from exc334 335 folder_id = str(created.get("id") or "").strip()336 if not folder_id:337 raise GoogleDriveConfigError(f"Folder '{folder_name}' was created but no id was returned.")338 339 logger.info("Created Drive folder name=%s parent=%s id=%s", folder_name, parent_folder_id or "root", folder_id)340 341 return folder_id342 343 344def ensure_drive_folder_path(service, folder_path: str, root_folder_id: str | None = None) -> str:345 normalized = (folder_path or "").strip().strip("/")346 if not normalized:347 raise GoogleDriveConfigError("folder_path is required to resolve a Drive destination.")348 349 current_parent = (root_folder_id or "").strip() or None350 for part in [segment.strip() for segment in normalized.split("/") if segment.strip()]:351 folder_id = _find_folder_in_parent(service, part, current_parent)352 if folder_id is None:353 folder_id = _create_folder_in_parent(service, part, current_parent)354 current_parent = folder_id355 356 if not current_parent:357 raise GoogleDriveConfigError(f"Could not resolve Drive path '{folder_path}'.")358 359 logger.info("Resolved Drive path=%s to folder_id=%s", folder_path, current_parent)360 361 return current_parent362 363 364def move_file_to_path(service, file_id: str, destination_path: str, root_folder_id: str | None = None) -> str:365 normalized_file_id = (file_id or "").strip()366 if not normalized_file_id:367 raise GoogleDriveConfigError("file_id is required to move a Drive file.")368 369 destination_folder_id = ensure_drive_folder_path(service, destination_path, root_folder_id=root_folder_id)370 logger.info("Moving Drive file_id=%s to destination_path=%s", normalized_file_id, destination_path)371 372 try:373 metadata = (374 service.files()375 .get(fileId=normalized_file_id, fields="parents", supportsAllDrives=True)376 .execute()377 )378 except Exception as exc:379 raise GoogleDriveConfigError(f"Failed to read current parents for file '{normalized_file_id}': {exc}") from exc380 381 current_parents = [str(parent).strip() for parent in metadata.get("parents", []) if str(parent).strip()]382 383 if destination_folder_id in current_parents and len(current_parents) == 1:384 logger.info("Drive file already in destination file_id=%s folder_id=%s", normalized_file_id, destination_folder_id)385 return destination_folder_id386 387 request_kwargs: dict[str, Any] = {388 "fileId": normalized_file_id,389 "addParents": destination_folder_id,390 "supportsAllDrives": True,391 "fields": "id,parents",392 }393 remove_parents = ",".join(parent for parent in current_parents if parent != destination_folder_id)394 if remove_parents:395 request_kwargs["removeParents"] = remove_parents396 397 try:398 service.files().update(**request_kwargs).execute()399 except Exception as exc:400 raise GoogleDriveConfigError(401 f"Failed to move file '{normalized_file_id}' to Drive path '{destination_path}': {exc}"402 ) from exc403 404 logger.info("Moved Drive file_id=%s to folder_id=%s", normalized_file_id, destination_folder_id)405 406 return destination_folder_id407