fix: add loop to handle pagination

Co-authored-by: Pat Heard <[email protected]>
cds-snc · Feb 2, 2024 · 24c101e · 24c101e
1 parent aa88af7
commit 24c101e
Showing 1 changed file with 21 additions and 13 deletions.
diff --git 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", [])