Austin/dify
1
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2024-11-16 11:42:29 +08:00
Code Issues Actions 4 Packages Projects Releases Wiki Activity

Feature/self host notion import (#397)

Browse Source
This commit is contained in:
Jyong 2023-06-17 19:50:21 +08:00 committed by GitHub
parent 402b0b81d2
commit 226f28edcb
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 73 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
api/config.py
View File

@ -190,6 +190,9 @@ class Config:
# notion import setting
self.NOTION_CLIENT_ID = get_env('NOTION_CLIENT_ID')
self.NOTION_CLIENT_SECRET = get_env('NOTION_CLIENT_SECRET')
self.NOTION_INTEGRATION_TYPE = get_env('NOTION_INTEGRATION_TYPE')
self.NOTION_INTERNAL_SECRET = get_env('NOTION_INTERNAL_SECRET')
class CloudEditionConfig(Config):

8
api/controllers/console/auth/data_source_oauth.py
View File

@ -39,11 +39,17 @@ class OAuthDataSource(Resource):
print(vars(oauth_provider))
if not oauth_provider:
return {'error': 'Invalid provider'}, 400
if current_app.config.get('NOTION_INTEGRATION_TYPE') == 'internal':
internal_secret = current_app.config.get('NOTION_INTERNAL_SECRET')
oauth_provider.save_internal_access_token(internal_secret)
return redirect(f'{current_app.config.get("CONSOLE_URL")}?oauth_data_source=success')
else:
auth_url = oauth_provider.get_authorization_url()
return redirect(auth_url)
class OAuthDataSourceCallback(Resource):
def get(self, provider: str):
OAUTH_DATASOURCE_PROVIDERS = get_oauth_providers()

6
api/core/data_source/notion.py
View File

@ -84,7 +84,8 @@ class NotionPageReader(BaseReader):
heading = text
result_block_id = result["id"]
has_children = result["has_children"]
if has_children:
block_type = result["type"]
if has_children and block_type != 'child_page':
children_text = self._read_block(
result_block_id, num_tabs=num_tabs + 1
)
@ -184,7 +185,8 @@ class NotionPageReader(BaseReader):
result_block_id = result["id"]
has_children = result["has_children"]
if has_children:
block_type = result["type"]
if has_children and block_type != 'child_page':
children_text = self._read_block(
result_block_id, num_tabs=num_tabs + 1
)

56
api/libs/oauth_data_source.py
View File

@ -26,6 +26,7 @@ class NotionOAuth(OAuthDataSource):
_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 = {
@ -84,6 +85,41 @@ class NotionOAuth(OAuthDataSource):
db.session.add(new_data_source_binding)
db.session.commit()
def save_internal_access_token(self, access_token: str):
workspace_name = self.notion_workspace_name(access_token)
workspace_icon = None
workspace_id = current_user.current_tenant_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)
}
# save data source binding
data_source_binding = DataSourceBinding.query.filter(
db.and_(
DataSourceBinding.tenant_id == current_user.current_tenant_id,
DataSourceBinding.provider == 'notion',
DataSourceBinding.access_token == access_token
)
).first()
if data_source_binding:
data_source_binding.source_info = source_info
data_source_binding.disabled = False
db.session.commit()
else:
new_data_source_binding = DataSourceBinding(
tenant_id=current_user.current_tenant_id,
access_token=access_token,
source_info=source_info,
provider='notion'
)
db.session.add(new_data_source_binding)
db.session.commit()
def sync_data_source(self, binding_id: str):
# save data source binding
data_source_binding = DataSourceBinding.query.filter(
@ -222,7 +258,10 @@ class NotionOAuth(OAuthDataSource):
}
response = requests.post(url=self._NOTION_PAGE_SEARCH, json=data, headers=headers)
response_json = response.json()
if 'results' in response_json:
results = response_json['results']
else:
results = []
return results
def notion_block_parent_page_id(self, access_token: str, block_id: str):
@ -238,6 +277,20 @@ class NotionOAuth(OAuthDataSource):
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',
}
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']
user_info = response_json[user_type]
if 'workspace_name' in user_info:
return user_info['workspace_name']
return 'workspace'
def notion_database_search(self, access_token: str):
data = {
'filter': {
@ -252,5 +305,8 @@ class NotionOAuth(OAuthDataSource):
}
response = requests.post(url=self._NOTION_PAGE_SEARCH, json=data, headers=headers)
response_json = response.json()
if 'results' in response_json:
results = response_json['results']
else:
results = []
return results