Add timeout to all Notion calls

This commit is contained in:
Weves
2023-11-06 19:29:12 -08:00
committed by Chris Weaver
parent f5bf2e6374
commit abf9cc3248

View File

@@ -24,6 +24,8 @@ from danswer.utils.logger import setup_logger
logger = setup_logger() logger = setup_logger()
_NOTION_CALL_TIMEOUT = 30 # 30 seconds
@dataclass @dataclass
class NotionPage: class NotionPage:
@@ -96,7 +98,12 @@ class NotionConnector(LoadConnector, PollConnector):
logger.debug(f"Fetching children of block with ID '{block_id}'") logger.debug(f"Fetching children of block with ID '{block_id}'")
block_url = f"https://api.notion.com/v1/blocks/{block_id}/children" block_url = f"https://api.notion.com/v1/blocks/{block_id}/children"
query_params = None if not cursor else {"start_cursor": cursor} query_params = None if not cursor else {"start_cursor": cursor}
res = requests.get(block_url, headers=self.headers, params=query_params) res = requests.get(
block_url,
headers=self.headers,
params=query_params,
timeout=_NOTION_CALL_TIMEOUT,
)
try: try:
res.raise_for_status() res.raise_for_status()
except Exception as e: except Exception as e:
@@ -109,7 +116,11 @@ class NotionConnector(LoadConnector, PollConnector):
"""Fetch a page from it's ID via the Notion API.""" """Fetch a page from it's ID via the Notion API."""
logger.debug(f"Fetching page for ID '{page_id}'") logger.debug(f"Fetching page for ID '{page_id}'")
block_url = f"https://api.notion.com/v1/pages/{page_id}" block_url = f"https://api.notion.com/v1/pages/{page_id}"
res = requests.get(block_url, headers=self.headers) res = requests.get(
block_url,
headers=self.headers,
timeout=_NOTION_CALL_TIMEOUT,
)
try: try:
res.raise_for_status() res.raise_for_status()
except Exception as e: except Exception as e:
@@ -125,7 +136,12 @@ class NotionConnector(LoadConnector, PollConnector):
logger.debug(f"Fetching database for ID '{database_id}'") logger.debug(f"Fetching database for ID '{database_id}'")
block_url = f"https://api.notion.com/v1/databases/{database_id}/query" block_url = f"https://api.notion.com/v1/databases/{database_id}/query"
body = None if not cursor else {"start_cursor": cursor} body = None if not cursor else {"start_cursor": cursor}
res = requests.post(block_url, headers=self.headers, json=body) res = requests.post(
block_url,
headers=self.headers,
json=body,
timeout=_NOTION_CALL_TIMEOUT,
)
try: try:
res.raise_for_status() res.raise_for_status()
except Exception as e: except Exception as e:
@@ -286,6 +302,7 @@ class NotionConnector(LoadConnector, PollConnector):
"https://api.notion.com/v1/search", "https://api.notion.com/v1/search",
headers=self.headers, headers=self.headers,
json=query_dict, json=query_dict,
timeout=_NOTION_CALL_TIMEOUT,
) )
res.raise_for_status() res.raise_for_status()
return NotionSearchResponse(**res.json()) return NotionSearchResponse(**res.json())