mirror of https://github.com/langgenius/dify.git
99 lines
3.3 KiB
Python
99 lines
3.3 KiB
Python
import json
|
|
import urllib.parse
|
|
|
|
import requests
|
|
from flask_login import current_user
|
|
|
|
from extensions.ext_database import db
|
|
from models.data_source import DataSourceBinding
|
|
|
|
|
|
class OAuthDataSource:
|
|
def __init__(self, client_id: str, client_secret: str, redirect_uri: str):
|
|
self.client_id = client_id
|
|
self.client_secret = client_secret
|
|
self.redirect_uri = redirect_uri
|
|
|
|
def get_authorization_url(self):
|
|
raise NotImplementedError()
|
|
|
|
def get_access_token(self, code: str):
|
|
raise NotImplementedError()
|
|
|
|
|
|
class NotionOAuth(OAuthDataSource):
|
|
_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"
|
|
|
|
def get_authorization_url(self):
|
|
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)}"
|
|
|
|
def get_access_token(self, code: str):
|
|
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')
|
|
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')
|
|
# 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
|
|
}
|
|
# save data source binding
|
|
data_source_binding = DataSourceBinding(
|
|
tenant_id=current_user.current_tenant_id,
|
|
access_token=access_token,
|
|
source_info=json.dumps(source_info)
|
|
)
|
|
db.session.add(data_source_binding)
|
|
db.session.commit()
|
|
|
|
def get_authorized_pages(self, access_token: str):
|
|
pages = []
|
|
data = {
|
|
'filter': {
|
|
"value": "page",
|
|
"property": "object"
|
|
}
|
|
}
|
|
headers = {
|
|
'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['results']
|
|
for result in results:
|
|
page_id = result['id']
|
|
page_name = result['properties']['title']['title'][0]['plain_text']
|
|
page_icon = result['icon']
|
|
page = {
|
|
'page_id': page_id,
|
|
'page_name': page_name,
|
|
'page_icon': page_icon
|
|
}
|
|
pages.append(page)
|
|
return pages
|