Spaces:
Sleeping
Sleeping
File size: 3,656 Bytes
e69be74 4a10a29 e69be74 eba303d e69be74 eba303d e69be74 |
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 |
import os
import tempfile
from google.oauth2 import service_account
from googleapiclient.discovery import build
from googleapiclient.http import MediaFileUpload
from googleapiclient.errors import HttpError
class GoogleDriveClient:
def __init__(self, json_key_path=None, json_key_env_var='GOOGLE_SERVICE_ACCOUNT_KEY'):
"""
Initializes the Google Drive client.
:param json_key_path: Path to the JSON file (optional)
:param json_key_env_var: Environment variable name containing the JSON key (default 'GOOGLE_SERVICE_ACCOUNT_KEY')
"""
self.json_key_path = json_key_path
self.json_key_env_var = json_key_env_var
self.service = self._authenticate_google_drive()
def _authenticate_google_drive(self):
"""
Authenticates using the JSON key from a file or environment variable.
:return: Google Drive API service object
"""
if self.json_key_path and os.path.isfile(self.json_key_path):
# Authenticate using the file
creds = service_account.Credentials.from_service_account_file(
self.json_key_path,
scopes=['https://www.googleapis.com/auth/drive.file']
)
else:
# Get JSON key from environment variable
json_key = os.getenv(self.json_key_env_var)
if not json_key:
raise ValueError(f"Environment variable {self.json_key_env_var} is not set.")
# Save JSON key to a temporary file
with tempfile.NamedTemporaryFile(delete=False) as temp_file:
temp_file.write(json_key.encode('utf-8'))
temp_file_path = temp_file.name
# Create credentials object
creds = service_account.Credentials.from_service_account_file(
temp_file_path,
scopes=['https://www.googleapis.com/auth/drive.file']
)
# Remove the temporary file after creating the credentials object
os.remove(temp_file_path)
# Create and return the Google Drive API service object
return build('drive', 'v3', credentials=creds)
def upload_file(self, file_path, folder_id):
"""
Uploads a file to Google Drive.
:param file_path: Path to the local file
:param folder_id: ID of the folder on Google Drive
:return: URL of the uploaded file or error message
"""
try:
file_metadata = {
'name': os.path.basename(file_path),
'parents': [folder_id] # ID of the folder to upload the file to
}
media = MediaFileUpload(file_path, mimetype='application/octet-stream')
# Upload the file
file = self.service.files().create(body=file_metadata, media_body=media, fields='id').execute()
# Get the file ID and construct the URL
file_id = file.get('id')
file_link = f'https://drive.google.com/file/d/{file_id}/view?usp=drive_link'
return file_link
except HttpError as error:
return f"An error occurred: {error}"
# Example usage
if __name__ == '__main__':
# Path to the local file
file_path = './file.txt'
# ID of the folder on Google Drive
# folder_id = '19yiqYX_Z1rbHDxnFLJTvji0VcbhSO9cq'
folder_id = '10qtum6ykbGTyu7vvw59i3h1XSY3-lRpo'
# Initialize the client and upload the file
client = GoogleDriveClient(json_key_path='secrets/GOOGLE_SERVICE_DRIVE_KEY_435817.json')
file_link = client.upload_file(file_path, folder_id)
print(f'File upload result: {file_link}')
|