Reorg public-private functions

This commit is contained in:
Pratik Kabra
2023-07-20 18:04:48 -05:00
parent 610fe6ebc4
commit b33c8b1d7c

View File

@@ -65,13 +65,6 @@ class NotionConnector(LoadConnector, PollConnector):
"Notion-Version": "2022-06-28",
}
def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:
"""Applies integration token to headers"""
self.headers[
"Authorization"
] = f'Bearer {credentials["notion_integration_token"]}'
return None
def _read_blocks(self, block_id: str, num_tabs: int = 0) -> str:
"""Reads blocks for a page"""
done = False
@@ -157,25 +150,6 @@ class NotionConnector(LoadConnector, PollConnector):
res.raise_for_status()
return NotionSearchResponse(**res.json())
def load_from_state(self) -> GenerateDocumentsOutput:
"""Loads all page data from a Notion workspace.
Returns:
List[Document]: List of documents.
"""
query_dict = {
"filter": {"property": "object", "value": "page"},
"page_size": self.batch_size,
}
while True:
db_res = self._search_notion(query_dict)
pages = [NotionPage(**page) for page in db_res.results]
yield self._read_pages(pages)
if db_res.has_more:
query_dict["start_cursor"] = db_res.next_cursor
else:
break
def _filter_pages_by_time(
self,
pages: List[Dict[str, Any]],
@@ -202,6 +176,32 @@ class NotionConnector(LoadConnector, PollConnector):
filtered_pages += [NotionPage(**page)]
return filtered_pages
def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:
"""Applies integration token to headers"""
self.headers[
"Authorization"
] = f'Bearer {credentials["notion_integration_token"]}'
return None
def load_from_state(self) -> GenerateDocumentsOutput:
"""Loads all page data from a Notion workspace.
Returns:
List[Document]: List of documents.
"""
query_dict = {
"filter": {"property": "object", "value": "page"},
"page_size": self.batch_size,
}
while True:
db_res = self._search_notion(query_dict)
pages = [NotionPage(**page) for page in db_res.results]
yield self._read_pages(pages)
if db_res.has_more:
query_dict["start_cursor"] = db_res.next_cursor
else:
break
def poll_source(
self, start: SecondsSinceUnixEpoch, end: SecondsSinceUnixEpoch
) -> GenerateDocumentsOutput: