google_drive_client.py

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}')