""" accounts/tasks.py Celery tasks for the Not Cloud Storage system. Key tasks: - process_order_download: downloads all order files from Google Drive (or local storage), assembles a zip, and updates progress in DB. - cleanup_expired_downloads: removes old temp files. - cleanup_stale_chunks: removes incomplete chunked uploads. Google Workspace export limits ────────────────────────────── Workspace files (Docs, Sheets, Slides, etc.) must be *exported* (converted) by Google's servers before they can be downloaded. Google enforces hard size limits on these exports: • Google Docs → DOCX up to ~10 MB source size • Google Sheets → XLSX up to ~10 MB source size • Google Slides → PPTX up to ~100 MB source size • Other Workspace → PDF up to ~10 MB source size If a Workspace file exceeds its export limit Google returns HTTP 403 with reason "exportSizeLimitExceeded". The file is NOT a real binary — it is a JSON-based collaborative document that must be converted. There is no way to download it raw; the only options are: 1. The user manually exports it from drive.google.com as PDF / Office format. 2. Third-party conversion services (outside our scope). This module marks over-limit Workspace files as 'skipped' (not 'failed') and continues processing the rest of the order. The skip reason is stored in file_selection.download_error so the UI can surface it clearly. Regular binary files (videos, ZIPs, raw PDFs, images, etc.) can be up to 5 TB and are downloaded by streaming 10 MB chunks directly to disk — they are never loaded fully into memory. """ import os import shutil import zipfile import logging import requests as req from datetime import timedelta from celery import shared_task from celery.exceptions import SoftTimeLimitExceeded from django.utils import timezone from django.conf import settings from django.db import transaction logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- # Constants # --------------------------------------------------------------------------- DOWNLOAD_CHUNK_SIZE = 10 * 1024 * 1024 # 10 MB per Drive streaming chunk DOWNLOAD_TEMP_DIR = getattr(settings, 'DOWNLOAD_TEMP_DIR', '/tmp/notcloudstorage/downloads') CHUNK_UPLOAD_TEMP_DIR = getattr(settings, 'CHUNK_UPLOAD_TEMP_DIR', '/tmp/notcloudstorage/chunks') DOWNLOAD_FILE_TTL_HOURS = getattr(settings, 'DOWNLOAD_FILE_TTL_HOURS', 48) # Google Workspace MIME → (export MIME, file extension) WORKSPACE_EXPORT_MAP = { 'application/vnd.google-apps.document': { 'mime': 'application/vnd.openxmlformats-officedocument.wordprocessingml.document', 'ext': '.docx', 'label': 'Google Doc', }, 'application/vnd.google-apps.spreadsheet': { 'mime': 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet', 'ext': '.xlsx', 'label': 'Google Sheet', }, 'application/vnd.google-apps.presentation': { 'mime': 'application/vnd.openxmlformats-officedocument.presentationml.presentation', 'ext': '.pptx', 'label': 'Google Slides', }, 'application/vnd.google-apps.drawing': { 'mime': 'image/png', 'ext': '.png', 'label': 'Google Drawing', }, 'application/vnd.google-apps.form': { 'mime': 'application/pdf', 'ext': '.pdf', 'label': 'Google Form', }, 'application/vnd.google-apps.jam': { 'mime': 'application/pdf', 'ext': '.pdf', 'label': 'Jamboard', }, 'application/vnd.google-apps.script': { 'mime': 'application/vnd.google-apps.script+json', 'ext': '.json', 'label': 'Apps Script', }, } # Fallback export for unknown Workspace types WORKSPACE_FALLBACK = {'mime': 'application/pdf', 'ext': '.pdf', 'label': 'Workspace file'} # Already-compressed formats — use ZIP_STORED (no point recompressing) ALREADY_COMPRESSED_EXTS = { '.zip', '.gz', '.bz2', '.xz', '.7z', '.rar', '.mp4', '.mov', '.avi', '.mkv', '.webm', '.mp3', '.aac', '.ogg', '.flac', '.jpg', '.jpeg', '.png', '.gif', '.webp', '.heic', } # --------------------------------------------------------------------------- # Custom exceptions # --------------------------------------------------------------------------- class ExportTooLargeError(Exception): """ Raised when Google returns exportSizeLimitExceeded for a Workspace file. These files should be marked 'skipped', not 'failed' — the user must manually export them from Google Drive. """ class FileTooLargeForExportError(Exception): """Alias kept for compatibility.""" # --------------------------------------------------------------------------- # Main download task # --------------------------------------------------------------------------- @shared_task( bind=True, max_retries=3, default_retry_delay=90, soft_time_limit=3 * 3600, # 3 h soft limit → SoftTimeLimitExceeded time_limit=4 * 3600, # 4 h hard kill name='accounts.tasks.process_order_download', acks_late=True, # re-queue on worker crash ) def process_order_download(self, order_id: str, file_ids: list = None): """ Download all files for a TransferOrder from Google Drive (or local storage), assemble a zip, and update the DownloadTask record with per-file progress. Files that are too large to export from Google Workspace are marked 'skipped' and a human-readable explanation is stored; the task continues for all other files so the admin still receives a complete zip. """ from .models import TransferOrder, DownloadTask logger.info( "process_order_download started | order=%s | attempt=%d", order_id, self.request.retries + 1, ) # ── Fetch order ──────────────────────────────────────────────────────── try: order = ( TransferOrder.objects .select_related('user') .prefetch_related('file_selections__local_file') .get(id=order_id) ) except TransferOrder.DoesNotExist: logger.error("Order %s not found — aborting", order_id) return {'status': 'error', 'message': 'Order not found'} # ── Create / reset DownloadTask ──────────────────────────────────────── download_task, _ = DownloadTask.objects.update_or_create( order=order, defaults={ 'celery_task_id': self.request.id, 'status': 'in_progress', 'progress': 0, 'error_message': '', 'started_at': timezone.now(), 'completed_at': None, 'temp_zip_path': None, }, ) # ── Ensure temp directory ────────────────────────────────────────────── order_temp_dir = os.path.join(DOWNLOAD_TEMP_DIR, str(order_id)) os.makedirs(order_temp_dir, exist_ok=True) try: # ── OAuth token ──────────────────────────────────────────────────── access_token = _get_valid_access_token(order) if not access_token: # Instead of aborting, just log a warning and proceed without Drive service logging.warning("No valid Google OAuth access token available. Google Drive operations will be skipped.") service = None else: # ── Build Drive service only if token exists ─────────────────────────── from google.oauth2.credentials import Credentials from googleapiclient.discovery import build service = build('drive', 'v3', credentials=Credentials(token=access_token), cache_discovery=False) # ── Resolve file list ────────────────────────────────────────────── if file_ids: file_selections = list(order.file_selections.filter(id__in=file_ids)) else: file_selections = list(order.file_selections.all()) total = len(file_selections) downloaded = 0 failed = [] skipped = [] # Workspace files that exceeded the export size limit _update_order_status(order, 'download_in_progress', f'Starting download of {total} file(s)…', started=True) # ── Download loop ────────────────────────────────────────────────── for idx, fs in enumerate(file_selections, start=1): try: _update_file_status(fs, 'downloading', 0) final_name = _download_single_file(service, fs, order_temp_dir) if not fs.local_file and service is None: raise Exception("Google Drive access token not available. Cannot download this file.") _update_file_status( fs, 'downloaded', 100, filename=final_name, size=_file_size_on_disk(order_temp_dir, final_name), ) downloaded += 1 logger.info("Downloaded %d/%d: %s → %s", idx, total, fs.file_name, final_name) except SoftTimeLimitExceeded: raise # propagate — Celery will handle retry except ExportTooLargeError as exc: # ── Workspace file too large to export ───────────────────── # This is not a system error — mark as skipped, not failed, # so the admin understands what happened. error_msg = str(exc) logger.warning("Skipped (too large): %s — %s", fs.file_name, error_msg) _update_file_status(fs, 'skipped', 0, error=error_msg) skipped.append({'name': fs.file_name, 'reason': error_msg}) except Exception as exc: logger.warning("Failed: %s — %s", fs.file_name, exc, exc_info=True) _update_file_status(fs, 'failed', 0, error=str(exc)) failed.append({'name': fs.file_name, 'error': str(exc)}) # Progress: download phase occupies 0–85 % of overall progress progress = int((idx / total) * 85) _bump_task_progress(download_task, progress, f'Processed {idx}/{total} files…') _update_order_status(order, 'download_in_progress', f'Processed {idx}/{total} files…') # ── All files failed / skipped — abort before creating empty zip ── if downloaded == 0 and not skipped: raise RuntimeError( "Every file failed to download. Check individual file errors above." ) if downloaded == 0 and skipped: # Every file was a Workspace file that was too large — still informative raise RuntimeError( f"All {len(skipped)} file(s) are Google Workspace files that exceed the " "export size limit and cannot be downloaded automatically. " "Please ask the user to export them manually from Google Drive." ) # ── Build zip ────────────────────────────────────────────────────── _bump_task_progress(download_task, 88, 'Creating zip archive…') _update_order_status(order, 'download_in_progress', 'Creating zip archive…') zip_filename = f"Order_{order.order_number}_Files.zip" zip_path = os.path.join(order_temp_dir, zip_filename) _build_zip(order_temp_dir, zip_path, exclude=zip_filename) zip_size = os.path.getsize(zip_path) logger.info("Zip ready: %s (%.1f MB)", zip_path, zip_size / 1024 / 1024) # ── Determine final order status ─────────────────────────────────── if not failed and not skipped: final_status = 'download_completed' elif downloaded > 0: final_status = 'download_partial' # some succeeded else: final_status = 'download_failed' # Build human-readable status message parts = [f'{downloaded}/{total} file(s) downloaded'] if skipped: parts.append( f'{len(skipped)} Workspace file(s) skipped (too large to export — ' 'user must export manually from Google Drive)' ) if failed: parts.append(f'{len(failed)} file(s) failed') status_msg = '; '.join(parts) # ── Compute total downloaded bytes ───────────────────────────────── # Reload from DB to get sizes written by _update_file_status total_downloaded_bytes = sum( fs.download_size for fs in order.file_selections.filter(download_status='downloaded') if fs.download_size ) # ── Commit ───────────────────────────────────────────────────────── with transaction.atomic(): download_task.status = 'completed' download_task.progress = 100 download_task.temp_zip_path = zip_path download_task.total_size_downloaded = zip_size download_task.completed_at = timezone.now() download_task.save() order.status = final_status order.download_completed_at = timezone.now() order.download_status = status_msg order.download_log = { 'downloaded': downloaded, 'skipped': len(skipped), 'failed': len(failed), 'total_files': total, 'successful_downloads': downloaded, 'total_downloaded_size': total_downloaded_bytes, 'skipped_files': skipped, 'failed_files': failed, 'zip_path': zip_path, 'zip_size_bytes': zip_size, 'completed_at': timezone.now().isoformat(), } order.save(update_fields=[ 'status', 'download_completed_at', 'download_status', 'download_log', 'updated_at', ]) logger.info( "process_order_download complete | order=%s | downloaded=%d skipped=%d failed=%d | zip=%.1f MB", order_id, downloaded, len(skipped), len(failed), zip_size / 1024 / 1024, ) return { 'status': 'success', 'downloaded': downloaded, 'skipped': len(skipped), 'failed': len(failed), } except SoftTimeLimitExceeded: msg = "Task exceeded the time limit — will retry." logger.error("SoftTimeLimitExceeded | order=%s", order_id) _fail_task(download_task, order, msg) raise self.retry(countdown=120, max_retries=2) except Exception as exc: logger.error("process_order_download failed | order=%s | %s", order_id, exc, exc_info=True) _fail_task(download_task, order, str(exc)) if self.request.retries < self.max_retries: countdown = 60 * (self.request.retries + 1) raise self.retry(exc=exc, countdown=countdown) shutil.rmtree(order_temp_dir, ignore_errors=True) return {'status': 'failed', 'error': str(exc)} # --------------------------------------------------------------------------- # Scheduled cleanup tasks # --------------------------------------------------------------------------- @shared_task(name='accounts.tasks.cleanup_expired_downloads') def cleanup_expired_downloads(): """Remove temp download directories older than DOWNLOAD_FILE_TTL_HOURS.""" from .models import DownloadTask cutoff = timezone.now() - timedelta(hours=DOWNLOAD_FILE_TTL_HOURS) expired = DownloadTask.objects.filter(completed_at__lt=cutoff, status='completed') count = 0 for task in expired: if task.temp_zip_path: order_dir = os.path.dirname(task.temp_zip_path) if os.path.exists(order_dir): shutil.rmtree(order_dir, ignore_errors=True) logger.info("Cleaned up %s", order_dir) task.temp_zip_path = None task.save(update_fields=['temp_zip_path']) count += 1 logger.info("cleanup_expired_downloads: removed %d directories", count) return {'cleaned': count} @shared_task(name='accounts.tasks.cleanup_stale_chunks') def cleanup_stale_chunks(): """Remove chunked-upload temp directories that are more than 2 hours old.""" from .models import ChunkedUploadPart cutoff = timezone.now() - timedelta(hours=2) stale = ChunkedUploadPart.objects.filter(created_at__lt=cutoff) upload_ids = stale.values_list('upload_id', flat=True).distinct() count = 0 for uid in upload_ids: chunk_dir = os.path.join(CHUNK_UPLOAD_TEMP_DIR, uid) if os.path.exists(chunk_dir): shutil.rmtree(chunk_dir, ignore_errors=True) count += 1 stale.delete() logger.info("cleanup_stale_chunks: removed %d directories", count) return {'cleaned': count} # --------------------------------------------------------------------------- # Core file-download helper # --------------------------------------------------------------------------- def _download_single_file(service, file_selection, dest_dir: str) -> str: """ Download one FileSelection to dest_dir. Returns the final on-disk filename. Raises: ExportTooLargeError – Workspace file exceeds Google's export size limit. FileNotFoundError – File not found in Drive. PermissionError – OAuth token lacks sufficient scopes. ValueError – File is a folder (cannot be downloaded directly). Exception – Any other Drive / network error. """ from googleapiclient.http import MediaIoBaseDownload from googleapiclient.errors import HttpError # ── Local uploaded file ──────────────────────────────────────────────── if file_selection.local_file: local_path = file_selection.local_file.file.path if not os.path.exists(local_path): raise FileNotFoundError(f"Local file missing on server: {local_path}") dest_name = _unique_name(dest_dir, file_selection.file_name) shutil.copy2(local_path, os.path.join(dest_dir, dest_name)) return dest_name # ── Google Drive file ────────────────────────────────────────────────── file_id = file_selection.file_id # Fetch metadata try: metadata = service.files().get( fileId=file_id, fields='mimeType,name,size', ).execute() except HttpError as e: _raise_from_http_error(e, file_selection.file_name) mime_type = metadata.get('mimeType', '') drive_name = metadata.get('name', file_selection.file_name) # ── Folder guard ─────────────────────────────────────────────────────── if mime_type == 'application/vnd.google-apps.folder': raise ValueError( f"'{drive_name}' is a folder. " "Please select the individual files inside it instead." ) # ── Google Workspace file (must be exported / converted) ────────────── if mime_type.startswith('application/vnd.google-apps'): mapping = WORKSPACE_EXPORT_MAP.get(mime_type, WORKSPACE_FALLBACK) export_mime = mapping['mime'] label = mapping['label'] base = os.path.splitext(drive_name)[0] final_name = _unique_name(dest_dir, base + mapping['ext']) dest_path = os.path.join(dest_dir, final_name) logger.info("Exporting %s '%s' → %s", label, drive_name, mapping['ext']) try: request = service.files().export_media(fileId=file_id, mimeType=export_mime) _stream_to_disk(request, dest_path, file_selection) except HttpError as e: # ── exportSizeLimitExceeded — the only non-retryable Drive error ── if _is_export_too_large(e): raise ExportTooLargeError( f"'{drive_name}' is a {label} that is too large for Google to export " f"automatically (Google's limit is ~10 MB for Docs/Sheets, ~100 MB for " f"Slides). To get this file, open Google Drive, right-click the file, " f"choose Download or File → Download, and upload it here as a local file." ) _raise_from_http_error(e, drive_name) return final_name # ── Regular binary file (stream directly) ───────────────────────────── final_name = _unique_name(dest_dir, drive_name) dest_path = os.path.join(dest_dir, final_name) try: request = service.files().get_media(fileId=file_id) _stream_to_disk(request, dest_path, file_selection) except HttpError as e: _raise_from_http_error(e, drive_name) return final_name def _stream_to_disk(request, dest_path: str, file_selection): """ Stream a Drive MediaIoBaseDownload to dest_path, updating per-file progress in the DB on each chunk (best-effort, non-fatal on DB errors). Uses DOWNLOAD_CHUNK_SIZE (10 MB) chunks so memory footprint stays flat even for multi-GB files. """ from googleapiclient.http import MediaIoBaseDownload with open(dest_path, 'wb') as fh: downloader = MediaIoBaseDownload(fh, request, chunksize=DOWNLOAD_CHUNK_SIZE) done = False while not done: status, done = downloader.next_chunk() if status: pct = int(status.progress() * 100) try: type(file_selection).objects.filter(pk=file_selection.pk).update( progress_percentage=pct, download_status='downloading', ) except Exception: pass # progress update failure must never kill the download # --------------------------------------------------------------------------- # Zip builder # --------------------------------------------------------------------------- def _build_zip(source_dir: str, zip_path: str, exclude: str): """ Add every file in source_dir (except `exclude`) to a zip. • allowZip64=True — supports archives > 4 GB • Already-compressed formats use ZIP_STORED to avoid wasting CPU """ with zipfile.ZipFile(zip_path, 'w', allowZip64=True) as zf: for fname in sorted(os.listdir(source_dir)): if fname == exclude: continue fpath = os.path.join(source_dir, fname) if not os.path.isfile(fpath): continue ext = os.path.splitext(fname)[1].lower() compress = (zipfile.ZIP_STORED if ext in ALREADY_COMPRESSED_EXTS else zipfile.ZIP_DEFLATED) zf.write(fpath, fname, compress_type=compress) # --------------------------------------------------------------------------- # OAuth token helper # --------------------------------------------------------------------------- def _get_valid_access_token(order) -> str | None: """ Return a valid access token for the order's user. Priority: 1. Stored token on the order (if not expiring within 5 min) 2. Refresh via stored refresh_token 3. Current token from social_django UserSocialAuth """ from django.conf import settings as djsettings if order.access_token and order.token_expiry: if order.token_expiry > timezone.now() + timedelta(minutes=5): return order.access_token if order.refresh_token: try: resp = req.post( 'https://oauth2.googleapis.com/token', data={ 'client_id': djsettings.SOCIAL_AUTH_GOOGLE_OAUTH2_KEY, 'client_secret': djsettings.SOCIAL_AUTH_GOOGLE_OAUTH2_SECRET, 'refresh_token': order.refresh_token, 'grant_type': 'refresh_token', }, timeout=30, ) resp.raise_for_status() data = resp.json() new_token = data.get('access_token') expires_in = data.get('expires_in', 3600) if new_token: order.access_token = new_token order.token_expiry = timezone.now() + timedelta(seconds=expires_in) order.save(update_fields=['access_token', 'token_expiry']) logger.info("Token refreshed for order %s", order.id) return new_token except Exception as exc: logger.warning("Token refresh failed for order %s: %s", order.id, exc) try: from social_django.models import UserSocialAuth sa = UserSocialAuth.objects.filter(user=order.user, provider='google-oauth2').first() if sa: token = sa.extra_data.get('access_token') refresh = sa.extra_data.get('refresh_token') if token: order.access_token = token if refresh: order.refresh_token = refresh order.token_expiry = timezone.now() + timedelta(hours=1) order.save(update_fields=['access_token', 'refresh_token', 'token_expiry']) return token except Exception as exc: logger.warning("Social auth token retrieval failed for order %s: %s", order.id, exc) return None # --------------------------------------------------------------------------- # Small utilities # --------------------------------------------------------------------------- def _is_export_too_large(http_error) -> bool: """Return True if an HttpError indicates exportSizeLimitExceeded.""" if http_error.resp.status != 403: return False try: import json details = json.loads(http_error.content.decode()) errors = details.get('error', {}).get('errors', []) return any(e.get('reason') == 'exportSizeLimitExceeded' for e in errors) except Exception: # Fallback: check raw content string return b'exportSizeLimitExceeded' in (http_error.content or b'') def _raise_from_http_error(e, filename: str): """Convert a Drive HttpError to an appropriate Python exception.""" status = e.resp.status if status == 404: raise FileNotFoundError(f"'{filename}' not found in Google Drive") if status in (401, 403): raise PermissionError( f"Access denied for '{filename}'. " "The OAuth token may have expired or lack the required Drive scopes." ) raise RuntimeError(f"Google Drive API error {status} for '{filename}': {e}") def _update_order_status(order, status: str, msg: str, started: bool = False): fields = ['status', 'download_status', 'updated_at'] order.status = status order.download_status = msg if started: order.download_started_at = timezone.now() fields.append('download_started_at') order.save(update_fields=fields) def _update_file_status(fs, status: str, progress: int, filename: str = None, size: int = None, error: str = None): fs.download_status = status fs.progress_percentage = progress fields = ['download_status', 'progress_percentage'] if filename is not None: fs.downloaded_filename = filename fields.append('downloaded_filename') if size is not None: fs.download_size = size fields.append('download_size') if status == 'downloaded': fs.downloaded_at = timezone.now() fields.append('downloaded_at') if error is not None: fs.download_error = error fields.append('download_error') fs.save(update_fields=fields) def _bump_task_progress(task, progress: int, msg: str = ''): task.progress = progress task.save(update_fields=['progress', 'updated_at']) def _fail_task(task, order, error_msg: str): task.status = 'failed' task.error_message = error_msg task.save(update_fields=['status', 'error_message', 'updated_at']) order.status = 'download_failed' order.download_status = f'Download failed: {error_msg[:300]}' order.save(update_fields=['status', 'download_status', 'updated_at']) def _unique_name(directory: str, filename: str) -> str: """Return filename (or filename_N) that does not already exist in directory.""" base, ext = os.path.splitext(filename) candidate = filename counter = 1 while os.path.exists(os.path.join(directory, candidate)): candidate = f"{base}_{counter}{ext}" counter += 1 return candidate def _file_size_on_disk(directory: str, filename: str) -> int: path = os.path.join(directory, filename) return os.path.getsize(path) if os.path.exists(path) else 0