Google Drive Improvements (#3057)

* Google Drive Improvements * mypy * should work! * variable cleanup * final fixes
danswer-ai · Nov 7, 2024 · 2758ffd · 2758ffd
1 parent 07a1b49
commit 2758ffd
Show file tree

Hide file tree

Showing 22 changed files with 802 additions and 525 deletions.
diff --git a/backend/danswer/connectors/gmail/connector.py b/backend/danswer/connectors/gmail/connector.py
@@ -277,16 +277,16 @@ def _fetch_threads(
                 fields=THREAD_LIST_FIELDS,
                 q=query,
             ):
-                full_thread = add_retries(
-                    lambda: gmail_service.users()
-                    .threads()
-                    .get(
-                        userId=user_email,
-                        id=thread["id"],
-                        fields=THREAD_FIELDS,
-                    )
-                    .execute()
-                )()
+                full_threads = execute_paginated_retrieval(
+                    retrieval_function=gmail_service.users().threads().get,
+                    list_key=None,
+                    userId=user_email,
+                    fields=THREAD_FIELDS,
+                    id=thread["id"],
+                )
+                # full_threads is an iterator containing a single thread
+                # so we need to convert it to a list and grab the first element
+                full_thread = list(full_threads)[0]
                 doc = thread_to_document(full_thread)
                 if doc is None:
                     continue