Spaces:
Sleeping
Sleeping
File size: 6,826 Bytes
80cb919 1c8c19d e3a165a 1c8c19d e3a165a 1c8c19d e3a165a 1c8c19d e3a165a 1c8c19d e3a165a 80cb919 1c8c19d e3a165a 1c8c19d 80cb919 e3a165a 1c8c19d 80cb919 e3a165a 1c8c19d 80cb919 e3a165a 80cb919 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 |
# Save final post-process to Google Drive
import os, json, logging
from typing import Optional
# Check if we're in local mode
IS_LOCAL = os.getenv("IS_LOCAL", "false").lower() == "true"
# Conditional imports for Google Drive (only when not in local mode)
if not IS_LOCAL:
try:
from google.oauth2 import service_account
from googleapiclient.discovery import build
from googleapiclient.http import MediaFileUpload
GOOGLE_DRIVE_AVAILABLE = True
except ImportError:
GOOGLE_DRIVE_AVAILABLE = False
# Create dummy classes for when Google Drive is not available
class service_account:
class Credentials:
@staticmethod
def from_service_account_info(*args, **kwargs):
raise ImportError("Google Drive dependencies not available")
class build:
@staticmethod
def build(*args, **kwargs):
raise ImportError("Google Drive dependencies not available")
class MediaFileUpload:
def __init__(self, *args, **kwargs):
raise ImportError("Google Drive dependencies not available")
else:
# Local mode: Google Drive is not available
GOOGLE_DRIVE_AVAILABLE = False
# Create dummy classes for local mode
class service_account:
class Credentials:
@staticmethod
def from_service_account_info(*args, **kwargs):
raise ImportError("Google Drive not available in local mode")
class build:
@staticmethod
def build(*args, **kwargs):
raise ImportError("Google Drive not available in local mode")
class MediaFileUpload:
def __init__(self, *args, **kwargs):
raise ImportError("Google Drive not available in local mode")
from utils.token import get_credentials
logger = logging.getLogger("dsaver")
if not logger.handlers:
logger.setLevel(logging.INFO)
fmt = logging.Formatter("[%(levelname)s] %(asctime)s - %(message)s")
handler = logging.StreamHandler()
handler.setFormatter(fmt)
logger.addHandler(handler)
class DriveSaver:
"""Google Drive uploader. Prefers OAuth; optional SA fallback (Shared Drive only)."""
def __init__(self, default_folder_id: Optional[str] = None):
self.service = None
self.folder_id = default_folder_id or os.getenv("GDRIVE_FOLDER_ID")
self.supports_all_drives = os.getenv("GDRIVE_FOLDER_IS_SHARED", "false").lower() in ("1","true","yes")
self.allow_sa_fallback = os.getenv("GDRIVE_ALLOW_SA_FALLBACK", "false").lower() in ("1","true","yes")
# Check if we're in local mode
if IS_LOCAL:
logger.info("🏠 Local mode: Google Drive integration disabled")
return
# Check if Google Drive is available
if not GOOGLE_DRIVE_AVAILABLE:
logger.warning("⚠️ Google Drive dependencies not available - DriveSaver will be disabled")
return
if not self.folder_id:
logger.warning("📁 No GDRIVE_FOLDER_ID set; uploads must provide folder_id explicitly")
self._initialize_service()
def _initialize_service(self):
if IS_LOCAL:
logger.info("🏠 Local mode: Skipping Google Drive service initialization")
return
if not GOOGLE_DRIVE_AVAILABLE:
logger.warning("⚠️ Google Drive dependencies not available - skipping service initialization")
return
creds = get_credentials()
if creds:
logger.info("✅ Using OAuth credentials")
else:
# Optional SA fallback — ONLY valid for Shared Drives where SA is a member
if self.allow_sa_fallback:
creds_env = os.getenv("GDRIVE_CREDENTIALS_JSON")
if creds_env:
try:
info = json.loads(creds_env)
if info.get("type") == "service_account":
creds = service_account.Credentials.from_service_account_info(
info, scopes=["https://www.googleapis.com/auth/drive"]
)
logger.info("✅ Using Service Account credentials (fallback)")
if not self.supports_all_drives:
logger.warning("⚠️ SA fallback without Shared Drive mode will likely fail (no quota). "
"Set GDRIVE_FOLDER_IS_SHARED=true and use a Shared Drive folder ID.")
else:
logger.error("❌ GDRIVE_CREDENTIALS_JSON is not a service account JSON")
except Exception as e:
logger.error(f"❌ Failed to init Service Account: {e}")
if not creds:
logger.error("❌ No valid Google credentials available (OAuth or SA).")
self.service = None
return
# Build Drive service
self.service = build("drive", "v3", credentials=creds)
logger.info("✅ Google Drive service initialized")
def upload_file_to_drive(self, file_path: str, folder_id: Optional[str] = None, mimetype: Optional[str] = None) -> bool:
if IS_LOCAL:
logger.info("🏠 Local mode: File upload to Google Drive skipped")
return False
if not GOOGLE_DRIVE_AVAILABLE:
logger.warning("⚠️ Google Drive dependencies not available - upload skipped")
return False
if not self.service:
logger.error("❌ Drive service not initialized")
return False
try:
target_folder = folder_id or self.folder_id
name = os.path.basename(file_path)
media = MediaFileUpload(file_path, mimetype=mimetype or "application/octet-stream")
metadata = {"name": name, "parents": [target_folder]}
req = self.service.files().create(
body=metadata,
media_body=media,
fields="id",
supportsAllDrives=self.supports_all_drives
)
req.execute()
logger.info(f"✅ Uploaded '{name}' to Drive (folder: {target_folder})")
return True
except Exception as e:
logger.error(f"❌ Drive upload failed: {e}")
return False
def is_service_available(self) -> bool:
return not IS_LOCAL and GOOGLE_DRIVE_AVAILABLE and self.service is not None
def set_folder_id(self, folder_id: str):
self.folder_id = folder_id
logger.info(f"📁 Default folder ID updated: {folder_id}")
|