Spaces:
Sleeping
Sleeping
File size: 6,698 Bytes
dae9b98 |
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 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 |
import os
import pickle
from google.auth.transport.requests import Request
from google_auth_oauthlib.flow import InstalledAppFlow
from googleapiclient.discovery import build
from googleapiclient.http import MediaFileUpload, MediaIoBaseUpload
import io
from pathlib import Path
import logging
logger = logging.getLogger(__name__)
# --- CONFIGURATION ---
# Get credentials from environment variables
CLIENT_ID = os.getenv("GOOGLE_CLIENT_ID")
CLIENT_SECRET = os.getenv("GOOGLE_CLIENT_SECRET")
# Google Drive folder IDs
EXPORTS_FOLDER_ID = "1k5iP4egzLrGJwnHkMhxt9bAkaCiieojO" # For Excel exports
IMAGES_FOLDER_ID = "1gd280IqcAzpAFTPeYsZjoBUOU9S7Zx3c" # For business card images
# Scopes define the level of access you are requesting.
SCOPES = ['https://www.googleapis.com/auth/drive.file']
TOKEN_PICKLE_FILE = 'token.pickle'
def get_drive_service():
"""Authenticates with Google and returns a Drive service object."""
creds = None
# The file token.pickle stores the user's access and refresh tokens.
if os.path.exists(TOKEN_PICKLE_FILE):
with open(TOKEN_PICKLE_FILE, 'rb') as token:
creds = pickle.load(token)
# If there are no (valid) credentials available, let the user log in.
if not creds or not creds.valid:
if creds and creds.expired and creds.refresh_token:
creds.refresh(Request())
else:
if not CLIENT_ID or not CLIENT_SECRET:
raise ValueError("GOOGLE_CLIENT_ID and GOOGLE_CLIENT_SECRET environment variables are required")
# Use client_config dictionary instead of a client_secret.json file
client_config = {
"installed": {
"client_id": CLIENT_ID,
"client_secret": CLIENT_SECRET,
"auth_uri": "https://accounts.google.com/o/oauth2/auth",
"token_uri": "https://oauth2.googleapis.com/token",
"redirect_uris": ["http://localhost"]
}
}
flow = InstalledAppFlow.from_client_config(client_config, SCOPES)
creds = flow.run_local_server(port=0)
# Save the credentials for the next run
with open(TOKEN_PICKLE_FILE, 'wb') as token:
pickle.dump(creds, token)
return build('drive', 'v3', credentials=creds)
def upload_file_to_drive(service, file_path=None, file_data=None, filename=None, folder_id=None, mimetype='application/octet-stream'):
"""
Uploads a file to a specific folder in Google Drive.
Args:
service: Google Drive service object
file_path: Path to local file (for file uploads)
file_data: Bytes data (for in-memory uploads)
filename: Name for the file in Drive
folder_id: ID of the target folder
mimetype: MIME type of the file
Returns:
dict: File information (id, webViewLink) or None if failed
"""
try:
if file_path and os.path.exists(file_path):
# Upload from local file
if not filename:
filename = os.path.basename(file_path)
media = MediaFileUpload(file_path, mimetype=mimetype, resumable=True)
logger.info(f"Uploading file from path: {file_path}")
elif file_data and filename:
# Upload from bytes data
file_io = io.BytesIO(file_data)
media = MediaIoBaseUpload(file_io, mimetype=mimetype, resumable=True)
logger.info(f"Uploading file from memory: {filename}")
else:
logger.error("Either file_path or (file_data + filename) must be provided")
return None
# Define the file's metadata
file_metadata = {
'name': filename,
'parents': [folder_id] if folder_id else []
}
logger.info(f"Uploading '{filename}' to Google Drive folder {folder_id}")
# Execute the upload request
file = service.files().create(
body=file_metadata,
media_body=media,
fields='id, webViewLink, name'
).execute()
logger.info(f"✅ File uploaded successfully!")
logger.info(f" File ID: {file.get('id')}")
logger.info(f" File Name: {file.get('name')}")
logger.info(f" View Link: {file.get('webViewLink')}")
return {
'id': file.get('id'),
'name': file.get('name'),
'webViewLink': file.get('webViewLink')
}
except Exception as e:
logger.error(f"Failed to upload file to Google Drive: {e}")
return None
def upload_excel_to_exports_folder(service, file_path=None, file_data=None, filename=None):
"""Upload Excel file to the exports folder."""
return upload_file_to_drive(
service,
file_path=file_path,
file_data=file_data,
filename=filename,
folder_id=EXPORTS_FOLDER_ID,
mimetype='application/vnd.openxmlformats-officedocument.spreadsheetml.sheet'
)
def upload_image_to_images_folder(service, file_path=None, file_data=None, filename=None, mimetype='image/png'):
"""Upload image file to the images folder."""
return upload_file_to_drive(
service,
file_path=file_path,
file_data=file_data,
filename=filename,
folder_id=IMAGES_FOLDER_ID,
mimetype=mimetype
)
def list_files_in_folder(service, folder_id, max_results=100):
"""List files in a specific Google Drive folder."""
try:
query = f"'{folder_id}' in parents"
results = service.files().list(
q=query,
maxResults=max_results,
fields="files(id, name, size, createdTime, webViewLink)"
).execute()
files = results.get('files', [])
logger.info(f"Found {len(files)} files in folder {folder_id}")
return files
except Exception as e:
logger.error(f"Failed to list files in folder {folder_id}: {e}")
return []
if __name__ == '__main__':
# Test the Google Drive connection
try:
drive_service = get_drive_service()
logger.info("Google Drive service initialized successfully")
# List files in both folders to verify access
exports_files = list_files_in_folder(drive_service, EXPORTS_FOLDER_ID)
images_files = list_files_in_folder(drive_service, IMAGES_FOLDER_ID)
print(f"Exports folder contains {len(exports_files)} files")
print(f"Images folder contains {len(images_files)} files")
except Exception as e:
logger.error(f"Failed to initialize Google Drive: {e}") |