feat: add pagination support for Notion search (#11194)

langgenius · Dec 1, 2024 · 56c2d1c · 56c2d1c
1 parent a67b0d4
commit 56c2d1c
Showing 1 changed file with 45 additions and 18 deletions.
diff --git a/api/libs/oauth_data_source.py b/api/libs/oauth_data_source.py
@@ -221,15 +221,29 @@ def get_authorized_pages(self, access_token: str):
         return pages
 
     def notion_page_search(self, access_token: str):
-        data = {"filter": {"value": "page", "property": "object"}}
-        headers = {
-            "Content-Type": "application/json",
-            "Authorization": f"Bearer {access_token}",
-            "Notion-Version": "2022-06-28",
-        }
-        response = requests.post(url=self._NOTION_PAGE_SEARCH, json=data, headers=headers)
-        response_json = response.json()
-        results = response_json.get("results", [])
+        results = []
+        next_cursor = None
+        has_more = True
+
+        while has_more:
+            data = {
+                "filter": {"value": "page", "property": "object"},
+                **({"start_cursor": next_cursor} if next_cursor else {}),
+            }
+            headers = {
+                "Content-Type": "application/json",
+                "Authorization": f"Bearer {access_token}",
+                "Notion-Version": "2022-06-28",
+            }
+
+            response = requests.post(url=self._NOTION_PAGE_SEARCH, json=data, headers=headers)
+            response_json = response.json()
+
+            results.extend(response_json.get("results", []))
+
+            has_more = response_json.get("has_more", False)
+            next_cursor = response_json.get("next_cursor", None)
+
         return results
 
     def notion_block_parent_page_id(self, access_token: str, block_id: str):
@@ -260,13 +274,26 @@ def notion_workspace_name(self, access_token: str):
         return "workspace"
 
     def notion_database_search(self, access_token: str):
-        data = {"filter": {"value": "database", "property": "object"}}
-        headers = {
-            "Content-Type": "application/json",
-            "Authorization": f"Bearer {access_token}",
-            "Notion-Version": "2022-06-28",
-        }
-        response = requests.post(url=self._NOTION_PAGE_SEARCH, json=data, headers=headers)
-        response_json = response.json()
-        results = response_json.get("results", [])
+        results = []
+        next_cursor = None
+        has_more = True
+
+        while has_more:
+            data = {
+                "filter": {"value": "database", "property": "object"},
+                **({"start_cursor": next_cursor} if next_cursor else {}),
+            }
+            headers = {
+                "Content-Type": "application/json",
+                "Authorization": f"Bearer {access_token}",
+                "Notion-Version": "2022-06-28",
+            }
+            response = requests.post(url=self._NOTION_PAGE_SEARCH, json=data, headers=headers)
+            response_json = response.json()
+
+            results.extend(response_json.get("results", []))
+
+            has_more = response_json.get("has_more", False)
+            next_cursor = response_json.get("next_cursor", None)
+
         return results