mirror of
https://github.com/langgenius/dify.git
synced 2026-08-29 11:49:06 +08:00
303 lines
13 KiB
Python
303 lines
13 KiB
Python
"""Notion adapter for the OAuth data-source application port."""
|
|
|
|
import urllib.parse
|
|
from collections.abc import Mapping
|
|
from typing import Any, Literal, TypedDict, override
|
|
|
|
import httpx
|
|
from pydantic import TypeAdapter, ValidationError
|
|
|
|
from core.helper.ssrf_proxy import MaxRetriesExceededError, SSRFProxy
|
|
from core.tools.errors import ToolSSRFError
|
|
from services.data_source_oauth_service import DataSourceOAuthError, DataSourceProviderGateway
|
|
from services.entities.data_source_oauth_entities import DataSourceOAuthAuthorization
|
|
|
|
|
|
class _NotionPageSummary(TypedDict):
|
|
page_id: str
|
|
page_name: str
|
|
page_icon: dict[str, str] | None
|
|
parent_id: str
|
|
type: Literal["page", "database"]
|
|
|
|
|
|
class _NotionSourceInfo(TypedDict):
|
|
workspace_name: str | None
|
|
workspace_icon: str | None
|
|
workspace_id: str | None
|
|
pages: list[_NotionPageSummary]
|
|
total: int
|
|
|
|
|
|
_NOTION_SOURCE_INFO_ADAPTER = TypeAdapter(_NotionSourceInfo)
|
|
_NOTION_PAGE_SUMMARY_ADAPTER = TypeAdapter(_NotionPageSummary)
|
|
|
|
|
|
class NotionDataSourceGateway(DataSourceProviderGateway):
|
|
_AUTH_URL = "https://api.notion.com/v1/oauth/authorize"
|
|
_TOKEN_URL = "https://api.notion.com/v1/oauth/token"
|
|
_NOTION_PAGE_SEARCH = "https://api.notion.com/v1/search"
|
|
_NOTION_BLOCK_SEARCH = "https://api.notion.com/v1/blocks"
|
|
_NOTION_BOT_USER = "https://api.notion.com/v1/users/me"
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
client_id: str,
|
|
client_secret: str,
|
|
redirect_uri: str,
|
|
http_client: httpx.Client | SSRFProxy,
|
|
) -> None:
|
|
self._client_id = client_id
|
|
self._client_secret = client_secret
|
|
self._redirect_uri = redirect_uri
|
|
self._http_client = http_client
|
|
|
|
@override
|
|
def get_authorization_url(self) -> str:
|
|
params = {
|
|
"client_id": self._client_id,
|
|
"response_type": "code",
|
|
"redirect_uri": self._redirect_uri,
|
|
"owner": "user",
|
|
}
|
|
return f"{self._AUTH_URL}?{urllib.parse.urlencode(params)}"
|
|
|
|
@override
|
|
def authorize(self, code: str) -> DataSourceOAuthAuthorization:
|
|
response_json = self._post_json(
|
|
self._TOKEN_URL,
|
|
data={"code": code, "grant_type": "authorization_code", "redirect_uri": self._redirect_uri},
|
|
auth=(self._client_id, self._client_secret),
|
|
headers={"Accept": "application/json"},
|
|
)
|
|
access_token = response_json.get("access_token")
|
|
if not isinstance(access_token, str) or not access_token:
|
|
raise DataSourceOAuthError("Notion OAuth response did not include an access token")
|
|
|
|
try:
|
|
source_info = self._build_source_info(
|
|
workspace_name=self._optional_string(response_json.get("workspace_name")),
|
|
workspace_icon=self._optional_string(response_json.get("workspace_icon")),
|
|
workspace_id=self._optional_string(response_json.get("workspace_id")),
|
|
pages=self.get_authorized_pages(access_token),
|
|
)
|
|
except (KeyError, TypeError, ValidationError) as exc:
|
|
raise DataSourceOAuthError("Notion OAuth response was invalid") from exc
|
|
return DataSourceOAuthAuthorization(access_token=access_token, source_info=source_info)
|
|
|
|
@override
|
|
def authorize_internal(self, access_token: str, workspace_id: str) -> DataSourceOAuthAuthorization:
|
|
try:
|
|
source_info = self._build_source_info(
|
|
workspace_name=self.notion_workspace_name(access_token),
|
|
workspace_icon=None,
|
|
workspace_id=workspace_id,
|
|
pages=self.get_authorized_pages(access_token),
|
|
)
|
|
except (KeyError, TypeError, ValidationError) as exc:
|
|
raise DataSourceOAuthError("Notion integration response was invalid") from exc
|
|
return DataSourceOAuthAuthorization(access_token=access_token, source_info=source_info)
|
|
|
|
@override
|
|
def refresh(self, access_token: str, source_info: Mapping[str, object]) -> Mapping[str, object]:
|
|
try:
|
|
return self._build_source_info(
|
|
workspace_name=self._optional_string(source_info.get("workspace_name")),
|
|
workspace_icon=self._optional_string(source_info.get("workspace_icon")),
|
|
workspace_id=self._optional_string(source_info.get("workspace_id")),
|
|
pages=self.get_authorized_pages(access_token),
|
|
)
|
|
except (KeyError, TypeError, ValidationError) as exc:
|
|
raise DataSourceOAuthError("Notion integration response was invalid") from exc
|
|
|
|
def get_authorized_pages(self, access_token: str) -> list[_NotionPageSummary]:
|
|
pages: list[_NotionPageSummary] = []
|
|
page_results = self.notion_page_search(access_token)
|
|
database_results = self.notion_database_search(access_token)
|
|
for page_result in page_results:
|
|
page_id = page_result["id"]
|
|
page_name = "Untitled"
|
|
for prop in page_result["properties"].values():
|
|
if title_list := prop.get("title"):
|
|
if title_list and "plain_text" in title_list[0]:
|
|
page_name = title_list[0]["plain_text"]
|
|
page_icon = page_result["icon"]
|
|
if page_icon:
|
|
icon_type = page_icon["type"]
|
|
if icon_type in {"external", "file"}:
|
|
url = page_icon[icon_type]["url"]
|
|
icon = {"type": "url", "url": url if url.startswith("http") else f"https://www.notion.so{url}"}
|
|
else:
|
|
icon = {"type": "emoji", "emoji": page_icon[icon_type]}
|
|
else:
|
|
icon = None
|
|
parent = page_result["parent"]
|
|
parent_type = parent["type"]
|
|
if parent_type == "block_id":
|
|
parent_id = self.notion_block_parent_page_id(access_token, parent[parent_type])
|
|
elif parent_type == "workspace":
|
|
parent_id = "root"
|
|
else:
|
|
parent_id = parent[parent_type]
|
|
pages.append(
|
|
_NOTION_PAGE_SUMMARY_ADAPTER.validate_python(
|
|
{
|
|
"page_id": page_id,
|
|
"page_name": page_name,
|
|
"page_icon": icon,
|
|
"parent_id": parent_id,
|
|
"type": "page",
|
|
}
|
|
)
|
|
)
|
|
|
|
for database_result in database_results:
|
|
page_id = database_result["id"]
|
|
page_name = database_result["title"][0]["plain_text"] if database_result["title"] else "Untitled"
|
|
page_icon = database_result["icon"]
|
|
if page_icon:
|
|
icon_type = page_icon["type"]
|
|
if icon_type in {"external", "file"}:
|
|
url = page_icon[icon_type]["url"]
|
|
icon = {"type": "url", "url": url if url.startswith("http") else f"https://www.notion.so{url}"}
|
|
else:
|
|
icon = {"type": icon_type, icon_type: page_icon[icon_type]}
|
|
else:
|
|
icon = None
|
|
parent = database_result["parent"]
|
|
parent_type = parent["type"]
|
|
if parent_type == "block_id":
|
|
parent_id = self.notion_block_parent_page_id(access_token, parent[parent_type])
|
|
elif parent_type == "workspace":
|
|
parent_id = "root"
|
|
else:
|
|
parent_id = parent[parent_type]
|
|
pages.append(
|
|
_NOTION_PAGE_SUMMARY_ADAPTER.validate_python(
|
|
{
|
|
"page_id": page_id,
|
|
"page_name": page_name,
|
|
"page_icon": icon,
|
|
"parent_id": parent_id,
|
|
"type": "database",
|
|
}
|
|
)
|
|
)
|
|
return pages
|
|
|
|
def notion_page_search(self, access_token: str) -> list[dict[str, Any]]:
|
|
return self._search(access_token, object_type="page")
|
|
|
|
def notion_database_search(self, access_token: str) -> list[dict[str, Any]]:
|
|
return self._search(access_token, object_type="database")
|
|
|
|
def _search(self, access_token: str, *, object_type: str) -> list[dict[str, Any]]:
|
|
results: list[dict[str, Any]] = []
|
|
next_cursor: str | None = None
|
|
has_more = True
|
|
while has_more:
|
|
data: dict[str, Any] = {
|
|
"filter": {"value": object_type, "property": "object"},
|
|
**({"start_cursor": next_cursor} if next_cursor else {}),
|
|
}
|
|
response_json = self._post_json(
|
|
self._NOTION_PAGE_SEARCH,
|
|
json=data,
|
|
headers=self._notion_headers(access_token, content_type=True),
|
|
)
|
|
raw_results = response_json.get("results", [])
|
|
if not isinstance(raw_results, list):
|
|
raise DataSourceOAuthError("Notion search response was invalid")
|
|
results.extend(raw_results)
|
|
has_more = response_json.get("has_more", False) is True
|
|
raw_cursor = response_json.get("next_cursor")
|
|
if has_more and (not isinstance(raw_cursor, str) or not raw_cursor):
|
|
raise DataSourceOAuthError("Notion search response did not include a valid next_cursor")
|
|
next_cursor = raw_cursor if isinstance(raw_cursor, str) else None
|
|
return results
|
|
|
|
def notion_block_parent_page_id(self, access_token: str, block_id: str) -> str:
|
|
response_json = self._get_json(
|
|
f"{self._NOTION_BLOCK_SEARCH}/{block_id}",
|
|
headers=self._notion_headers(access_token),
|
|
)
|
|
parent = response_json.get("parent")
|
|
if not isinstance(parent, dict):
|
|
raise DataSourceOAuthError("Notion block response was invalid")
|
|
parent_type = parent.get("type")
|
|
if not isinstance(parent_type, str) or parent_type not in parent:
|
|
raise DataSourceOAuthError("Notion block parent was invalid")
|
|
parent_id = parent[parent_type]
|
|
if not isinstance(parent_id, str):
|
|
raise DataSourceOAuthError("Notion block parent identifier was invalid")
|
|
if parent_type == "block_id":
|
|
return self.notion_block_parent_page_id(access_token, parent_id)
|
|
return parent_id
|
|
|
|
def notion_workspace_name(self, access_token: str) -> str:
|
|
response_json = self._get_json(self._NOTION_BOT_USER, headers=self._notion_headers(access_token))
|
|
if response_json.get("object") == "user":
|
|
user_type = response_json.get("type")
|
|
user_info = response_json.get(user_type) if isinstance(user_type, str) else None
|
|
if isinstance(user_info, dict) and isinstance(user_info.get("workspace_name"), str):
|
|
return user_info["workspace_name"]
|
|
return "workspace"
|
|
|
|
def _post_json(self, url: str, **kwargs: Any) -> dict[str, Any]:
|
|
try:
|
|
response = self._http_client.post(url, **kwargs)
|
|
except (httpx.HTTPError, MaxRetriesExceededError, ToolSSRFError) as exc:
|
|
raise DataSourceOAuthError("Notion request failed") from exc
|
|
return self._response_json(response)
|
|
|
|
def _get_json(self, url: str, **kwargs: Any) -> dict[str, Any]:
|
|
try:
|
|
response = self._http_client.get(url, **kwargs)
|
|
except (httpx.HTTPError, MaxRetriesExceededError, ToolSSRFError) as exc:
|
|
raise DataSourceOAuthError("Notion request failed") from exc
|
|
return self._response_json(response)
|
|
|
|
@staticmethod
|
|
def _response_json(response: httpx.Response) -> dict[str, Any]:
|
|
try:
|
|
response.raise_for_status()
|
|
payload = response.json()
|
|
except (httpx.HTTPError, ValueError) as exc:
|
|
raise DataSourceOAuthError("Notion returned an unsuccessful response") from exc
|
|
if not isinstance(payload, dict):
|
|
raise DataSourceOAuthError("Notion response was invalid")
|
|
return payload
|
|
|
|
@staticmethod
|
|
def _notion_headers(access_token: str, *, content_type: bool = False) -> dict[str, str]:
|
|
headers = {
|
|
"Authorization": f"Bearer {access_token}",
|
|
"Notion-Version": "2022-06-28",
|
|
}
|
|
if content_type:
|
|
headers["Content-Type"] = "application/json"
|
|
return headers
|
|
|
|
@staticmethod
|
|
def _optional_string(value: object) -> str | None:
|
|
return value if isinstance(value, str) else None
|
|
|
|
@staticmethod
|
|
def _build_source_info(
|
|
*,
|
|
workspace_name: str | None,
|
|
workspace_icon: str | None,
|
|
workspace_id: str | None,
|
|
pages: list[_NotionPageSummary],
|
|
) -> _NotionSourceInfo:
|
|
return _NOTION_SOURCE_INFO_ADAPTER.validate_python(
|
|
{
|
|
"workspace_name": workspace_name,
|
|
"workspace_icon": workspace_icon,
|
|
"workspace_id": workspace_id,
|
|
"pages": pages,
|
|
"total": len(pages),
|
|
}
|
|
)
|