From 0a6c2afb8a6d301fbe257818ad52ef1c4dcc1db5 Mon Sep 17 00:00:00 2001
From: Chris Weaver <25087905+Weves@users.noreply.github.com>
Date: Tue, 24 Oct 2023 15:00:50 -0700
Subject: [PATCH] Notion extra logging + small improvements (#621)

---
 backend/danswer/connectors/notion/connector.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/backend/danswer/connectors/notion/connector.py b/backend/danswer/connectors/notion/connector.py
index 2fa901fc6f8..5f854f44bdd 100644
--- a/backend/danswer/connectors/notion/connector.py
+++ b/backend/danswer/connectors/notion/connector.py
@@ -128,6 +128,7 @@ def _read_blocks(
             data = self._fetch_blocks(page_block_id, cursor)
 
             for result in data["results"]:
+                logger.debug(f"Found block for page '{page_block_id}': {result}")
                 result_block_id = result["id"]
                 result_type = result["type"]
                 result_obj = result[result_type]
@@ -280,6 +281,7 @@ def load_from_state(self) -> GenerateDocumentsOutput:
         # TODO: remove once Notion search issue is discovered
         if self.recursive_index_enabled and self.root_page_id:
             yield from self._recursive_load()
+            return
 
         query_dict = {
             "filter": {"property": "object", "value": "page"},
@@ -306,6 +308,7 @@ def poll_source(
         # TODO: remove once Notion search issue is discovered
         if self.recursive_index_enabled and self.root_page_id:
             yield from self._recursive_load()
+            return
 
         query_dict = {
             "page_size": self.batch_size,
@@ -328,7 +331,8 @@ def poll_source(
 if __name__ == "__main__":
     import os
 
-    connector = NotionConnector()
+    root_page_id = os.environ.get("NOTION_ROOT_PAGE_ID")
+    connector = NotionConnector(root_page_id=root_page_id)
     connector.load_credentials(
         {"notion_integration_token": os.environ.get("NOTION_INTEGRATION_TOKEN")}
     )