estel/dify
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

chore(api/libs): Apply ruff format. (#7301)

Browse Source
This commit is contained in:
-LAN-
2024-08-15 17:53:12 +08:00
committed by GitHub
parent d07b2b9915
commit 9414143b5f
15 changed files with 270 additions and 320 deletions
Download Patch File Download Diff File Expand all files Collapse all files

230
api/libs/oauth_data_source.py
View File

@@ -21,53 +21,49 @@ class OAuthDataSource:
class NotionOAuth(OAuthDataSource):
_AUTH_URL = 'https://api.notion.com/v1/oauth/authorize'
_TOKEN_URL = 'https://api.notion.com/v1/oauth/token'
_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 get_authorization_url(self):
params = {
'client_id': self.client_id,
'response_type': 'code',
'redirect_uri': self.redirect_uri,
'owner': 'user'
"client_id": self.client_id,
"response_type": "code",
"redirect_uri": self.redirect_uri,
"owner": "user",
}
return f"{self._AUTH_URL}?{urllib.parse.urlencode(params)}"
def get_access_token(self, code: str):
data = {
'code': code,
'grant_type': 'authorization_code',
'redirect_uri': self.redirect_uri
}
headers = {'Accept': 'application/json'}
data = {"code": code, "grant_type": "authorization_code", "redirect_uri": self.redirect_uri}
headers = {"Accept": "application/json"}
auth = (self.client_id, self.client_secret)
response = requests.post(self._TOKEN_URL, data=data, auth=auth, headers=headers)
response_json = response.json()
access_token = response_json.get('access_token')
access_token = response_json.get("access_token")
if not access_token:
raise ValueError(f"Error in Notion OAuth: {response_json}")
workspace_name = response_json.get('workspace_name')
workspace_icon = response_json.get('workspace_icon')
workspace_id = response_json.get('workspace_id')
workspace_name = response_json.get("workspace_name")
workspace_icon = response_json.get("workspace_icon")
workspace_id = response_json.get("workspace_id")
# get all authorized pages
pages = self.get_authorized_pages(access_token)
source_info = {
'workspace_name': workspace_name,
'workspace_icon': workspace_icon,
'workspace_id': workspace_id,
'pages': pages,
'total': len(pages)
"workspace_name": workspace_name,
"workspace_icon": workspace_icon,
"workspace_id": workspace_id,
"pages": pages,
"total": len(pages),
}
# save data source binding
data_source_binding = DataSourceOauthBinding.query.filter(
db.and_(
DataSourceOauthBinding.tenant_id == current_user.current_tenant_id,
DataSourceOauthBinding.provider == 'notion',
DataSourceOauthBinding.access_token == access_token
DataSourceOauthBinding.provider == "notion",
DataSourceOauthBinding.access_token == access_token,
)
).first()
if data_source_binding:
@@ -79,7 +75,7 @@ class NotionOAuth(OAuthDataSource):
tenant_id=current_user.current_tenant_id,
access_token=access_token,
source_info=source_info,
provider='notion'
provider="notion",
)
db.session.add(new_data_source_binding)
db.session.commit()
@@ -91,18 +87,18 @@ class NotionOAuth(OAuthDataSource):
# get all authorized pages
pages = self.get_authorized_pages(access_token)
source_info = {
'workspace_name': workspace_name,
'workspace_icon': workspace_icon,
'workspace_id': workspace_id,
'pages': pages,
'total': len(pages)
"workspace_name": workspace_name,
"workspace_icon": workspace_icon,
"workspace_id": workspace_id,
"pages": pages,
"total": len(pages),
}
# save data source binding
data_source_binding = DataSourceOauthBinding.query.filter(
db.and_(
DataSourceOauthBinding.tenant_id == current_user.current_tenant_id,
DataSourceOauthBinding.provider == 'notion',
DataSourceOauthBinding.access_token == access_token
DataSourceOauthBinding.provider == "notion",
DataSourceOauthBinding.access_token == access_token,
)
).first()
if data_source_binding:
@@ -114,7 +110,7 @@ class NotionOAuth(OAuthDataSource):
tenant_id=current_user.current_tenant_id,
access_token=access_token,
source_info=source_info,
provider='notion'
provider="notion",
)
db.session.add(new_data_source_binding)
db.session.commit()
@@ -124,9 +120,9 @@ class NotionOAuth(OAuthDataSource):
data_source_binding = DataSourceOauthBinding.query.filter(
db.and_(
DataSourceOauthBinding.tenant_id == current_user.current_tenant_id,
DataSourceOauthBinding.provider == 'notion',
DataSourceOauthBinding.provider == "notion",
DataSourceOauthBinding.id == binding_id,
DataSourceOauthBinding.disabled == False
DataSourceOauthBinding.disabled == False,
)
).first()
if data_source_binding:
@@ -134,17 +130,17 @@ class NotionOAuth(OAuthDataSource):
pages = self.get_authorized_pages(data_source_binding.access_token)
source_info = data_source_binding.source_info
new_source_info = {
'workspace_name': source_info['workspace_name'],
'workspace_icon': source_info['workspace_icon'],
'workspace_id': source_info['workspace_id'],
'pages': pages,
'total': len(pages)
"workspace_name": source_info["workspace_name"],
"workspace_icon": source_info["workspace_icon"],
"workspace_id": source_info["workspace_id"],
"pages": pages,
"total": len(pages),
}
data_source_binding.source_info = new_source_info
data_source_binding.disabled = False
db.session.commit()
else:
raise ValueError('Data source binding not found')
raise ValueError("Data source binding not found")
def get_authorized_pages(self, access_token: str):
pages = []
@@ -152,143 +148,121 @@ class NotionOAuth(OAuthDataSource):
database_results = self.notion_database_search(access_token)
# get page detail
for page_result in page_results:
page_id = page_result['id']
page_name = 'Untitled'
for key in page_result['properties']:
if 'title' in page_result['properties'][key] and page_result['properties'][key]['title']:
title_list = page_result['properties'][key]['title']
if len(title_list) > 0 and 'plain_text' in title_list[0]:
page_name = title_list[0]['plain_text']
page_icon = page_result['icon']
page_id = page_result["id"]
page_name = "Untitled"
for key in page_result["properties"]:
if "title" in page_result["properties"][key] and page_result["properties"][key]["title"]:
title_list = page_result["properties"][key]["title"]
if len(title_list) > 0 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 == 'external' or icon_type == 'file':
url = page_icon[icon_type]['url']
icon = {
'type': 'url',
'url': url if url.startswith('http') else f'https://www.notion.so{url}'
}
icon_type = page_icon["type"]
if icon_type == "external" or icon_type == "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]
}
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 = 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'
elif parent_type == "workspace":
parent_id = "root"
else:
parent_id = parent[parent_type]
page = {
'page_id': page_id,
'page_name': page_name,
'page_icon': icon,
'parent_id': parent_id,
'type': 'page'
"page_id": page_id,
"page_name": page_name,
"page_icon": icon,
"parent_id": parent_id,
"type": "page",
}
pages.append(page)
# get database detail
for database_result in database_results:
page_id = database_result['id']
if len(database_result['title']) > 0:
page_name = database_result['title'][0]['plain_text']
page_id = database_result["id"]
if len(database_result["title"]) > 0:
page_name = database_result["title"][0]["plain_text"]
else:
page_name = 'Untitled'
page_icon = database_result['icon']
page_name = "Untitled"
page_icon = database_result["icon"]
if page_icon:
icon_type = page_icon['type']
if icon_type == 'external' or icon_type == 'file':
url = page_icon[icon_type]['url']
icon = {
'type': 'url',
'url': url if url.startswith('http') else f'https://www.notion.so{url}'
}
icon_type = page_icon["type"]
if icon_type == "external" or icon_type == "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]
}
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 = 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'
elif parent_type == "workspace":
parent_id = "root"
else:
parent_id = parent[parent_type]
page = {
'page_id': page_id,
'page_name': page_name,
'page_icon': icon,
'parent_id': parent_id,
'type': 'database'
"page_id": page_id,
"page_name": page_name,
"page_icon": icon,
"parent_id": parent_id,
"type": "database",
}
pages.append(page)
return pages
def notion_page_search(self, access_token: str):
data = {
'filter': {
"value": "page",
"property": "object"
}
}
data = {"filter": {"value": "page", "property": "object"}}
headers = {
'Content-Type': 'application/json',
'Authorization': f"Bearer {access_token}",
'Notion-Version': '2022-06-28',
"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 = response_json.get("results", [])
return results
def notion_block_parent_page_id(self, access_token: str, block_id: str):
headers = {
'Authorization': f"Bearer {access_token}",
'Notion-Version': '2022-06-28',
"Authorization": f"Bearer {access_token}",
"Notion-Version": "2022-06-28",
}
response = requests.get(url=f'{self._NOTION_BLOCK_SEARCH}/{block_id}', headers=headers)
response = requests.get(url=f"{self._NOTION_BLOCK_SEARCH}/{block_id}", headers=headers)
response_json = response.json()
parent = response_json['parent']
parent_type = parent['type']
if parent_type == 'block_id':
parent = response_json["parent"]
parent_type = parent["type"]
if parent_type == "block_id":
return self.notion_block_parent_page_id(access_token, parent[parent_type])
return parent[parent_type]
def notion_workspace_name(self, access_token: str):
headers = {
'Authorization': f"Bearer {access_token}",
'Notion-Version': '2022-06-28',
"Authorization": f"Bearer {access_token}",
"Notion-Version": "2022-06-28",
}
response = requests.get(url=self._NOTION_BOT_USER, headers=headers)
response_json = response.json()
if 'object' in response_json and response_json['object'] == 'user':
user_type = response_json['type']
if "object" in response_json and response_json["object"] == "user":
user_type = response_json["type"]
user_info = response_json[user_type]
if 'workspace_name' in user_info:
return user_info['workspace_name']
return 'workspace'
if "workspace_name" in user_info:
return user_info["workspace_name"]
return "workspace"
def notion_database_search(self, access_token: str):
data = {
'filter': {
"value": "database",
"property": "object"
}
}
data = {"filter": {"value": "database", "property": "object"}}
headers = {
'Content-Type': 'application/json',
'Authorization': f"Bearer {access_token}",
'Notion-Version': '2022-06-28',
"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 = response_json.get("results", [])
return results