From 24c101e0ff08ec1fa2f6a1a97726c34c4e8fa01d Mon Sep 17 00:00:00 2001 From: Guillaume Charest <1690085+gcharest@users.noreply.github.com> Date: Fri, 2 Feb 2024 09:43:40 -0500 Subject: [PATCH] fix: add loop to handle pagination Co-authored-by: Pat Heard --- .../google_workspace/google_drive.py | 34 ++++++++++++------- 1 file changed, 21 insertions(+), 13 deletions(-) diff --git a/app/integrations/google_workspace/google_drive.py b/app/integrations/google_workspace/google_drive.py index c1786c37..dfc4da78 100644 --- a/app/integrations/google_workspace/google_drive.py +++ b/app/integrations/google_workspace/google_drive.py @@ -260,20 +260,28 @@ def list_folders_in_folder(folder): list: A list of folders in the folder. """ service = get_google_service("drive", "v3") - results = ( - service.files() - .list( - pageSize=25, - supportsAllDrives=True, - includeItemsFromAllDrives=True, - corpora="user", - q="parents in '{}' and mimeType = 'application/vnd.google-apps.folder' and trashed=false".format( - folder - ), - fields="nextPageToken, files(id, name)", + page_token = None + all_files = [] + while True: + results = ( + service.files() + .list( + pageSize=25, + supportsAllDrives=True, + includeItemsFromAllDrives=True, + corpora="user", + q="parents in '{}' and mimeType = 'application/vnd.google-apps.folder' and trashed=false".format( + folder + ), + fields="nextPageToken, files(id, name)", + pageToken=page_token, + ) + .execute() ) - .execute() - ) + all_files.extend(results.get('files', [])) + page_token = results.get('nextPageToken') + if not page_token: + break return results.get("files", [])