google_manager/drive.py

from googleapiclient.discovery import build
from googleapiclient.errors import HttpError
import logging

logging.basicConfig(level=logging.INFO)


def create_folder(creds, folder_name, parent_folder_id=None) -> str:
    """ """
    drive_service = build("drive", "v3", credentials=creds)

    # Create a metadata dict with the folder name and parent folder ID (if any)
    folder_metadata = {
        "name": folder_name,
        "mimeType": "application/vnd.google-apps.folder",
    }
    if parent_folder_id:
        folder_metadata["parents"] = [parent_folder_id]

    try:
        folder = (
            drive_service.files().create(body=folder_metadata, fields="id").execute()
        )
        logging.info(
            f'Folder "{folder_name}" with ID "{folder["id"]}" created successfully'
        )
        return folder["id"]
    except HttpError as error:
        logging.info(
            f"An error occurred while creating the Google Drive folder: {error}"
        )
        return None


def search_folder(creds, folder_name):
    """ """
    try:
        # create drive api client
        service = build("drive", "v3", credentials=creds)
        files = []
        page_token = None

        while True:
            response = (
                service.files()
                .list(
                    q="mimeType='application/vnd.google-apps.folder' and trashed=false",
                    spaces="drive",
                    fields="nextPageToken, files(id, name)",
                    pageToken=page_token,
                )
                .execute()
            )
            for file in response.get("files", []):
                # Process change
                logging.info(f'Found file: {file.get("name")}, {file.get("id")}')
            files.extend(response.get("files", []))
            page_token = response.get("nextPageToken", None)
            if page_token is None:
                break

    except HttpError as error:
        logging.info(f"An error occurred: {error}")
        files = None

    return files