Underground-Digital/Workflow-Engine

1import datetime2import urllib.parse3 4import requests5from flask_login import current_user6 7from extensions.ext_database import db8from models.source import DataSourceOauthBinding9 10 11class OAuthDataSource:12    def __init__(self, client_id: str, client_secret: str, redirect_uri: str):13        self.client_id = client_id14        self.client_secret = client_secret15        self.redirect_uri = redirect_uri16 17    def get_authorization_url(self):18        raise NotImplementedError()19 20    def get_access_token(self, code: str):21        raise NotImplementedError()22 23 24class NotionOAuth(OAuthDataSource):25    _AUTH_URL = "https://api.notion.com/v1/oauth/authorize"26    _TOKEN_URL = "https://api.notion.com/v1/oauth/token"27    _NOTION_PAGE_SEARCH = "https://api.notion.com/v1/search"28    _NOTION_BLOCK_SEARCH = "https://api.notion.com/v1/blocks"29    _NOTION_BOT_USER = "https://api.notion.com/v1/users/me"30 31    def get_authorization_url(self):32        params = {33            "client_id": self.client_id,34            "response_type": "code",35            "redirect_uri": self.redirect_uri,36            "owner": "user",37        }38        return f"{self._AUTH_URL}?{urllib.parse.urlencode(params)}"39 40    def get_access_token(self, code: str):41        data = {"code": code, "grant_type": "authorization_code", "redirect_uri": self.redirect_uri}42        headers = {"Accept": "application/json"}43        auth = (self.client_id, self.client_secret)44        response = requests.post(self._TOKEN_URL, data=data, auth=auth, headers=headers)45 46        response_json = response.json()47        access_token = response_json.get("access_token")48        if not access_token:49            raise ValueError(f"Error in Notion OAuth: {response_json}")50        workspace_name = response_json.get("workspace_name")51        workspace_icon = response_json.get("workspace_icon")52        workspace_id = response_json.get("workspace_id")53        # get all authorized pages54        pages = self.get_authorized_pages(access_token)55        source_info = {56            "workspace_name": workspace_name,57            "workspace_icon": workspace_icon,58            "workspace_id": workspace_id,59            "pages": pages,60            "total": len(pages),61        }62        # save data source binding63        data_source_binding = DataSourceOauthBinding.query.filter(64            db.and_(65                DataSourceOauthBinding.tenant_id == current_user.current_tenant_id,66                DataSourceOauthBinding.provider == "notion",67                DataSourceOauthBinding.access_token == access_token,68            )69        ).first()70        if data_source_binding:71            data_source_binding.source_info = source_info72            data_source_binding.disabled = False73            data_source_binding.updated_at = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)74            db.session.commit()75        else:76            new_data_source_binding = DataSourceOauthBinding(77                tenant_id=current_user.current_tenant_id,78                access_token=access_token,79                source_info=source_info,80                provider="notion",81            )82            db.session.add(new_data_source_binding)83            db.session.commit()84 85    def save_internal_access_token(self, access_token: str):86        workspace_name = self.notion_workspace_name(access_token)87        workspace_icon = None88        workspace_id = current_user.current_tenant_id89        # get all authorized pages90        pages = self.get_authorized_pages(access_token)91        source_info = {92            "workspace_name": workspace_name,93            "workspace_icon": workspace_icon,94            "workspace_id": workspace_id,95            "pages": pages,96            "total": len(pages),97        }98        # save data source binding99        data_source_binding = DataSourceOauthBinding.query.filter(100            db.and_(101                DataSourceOauthBinding.tenant_id == current_user.current_tenant_id,102                DataSourceOauthBinding.provider == "notion",103                DataSourceOauthBinding.access_token == access_token,104            )105        ).first()106        if data_source_binding:107            data_source_binding.source_info = source_info108            data_source_binding.disabled = False109            data_source_binding.updated_at = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)110            db.session.commit()111        else:112            new_data_source_binding = DataSourceOauthBinding(113                tenant_id=current_user.current_tenant_id,114                access_token=access_token,115                source_info=source_info,116                provider="notion",117            )118            db.session.add(new_data_source_binding)119            db.session.commit()120 121    def sync_data_source(self, binding_id: str):122        # save data source binding123        data_source_binding = DataSourceOauthBinding.query.filter(124            db.and_(125                DataSourceOauthBinding.tenant_id == current_user.current_tenant_id,126                DataSourceOauthBinding.provider == "notion",127                DataSourceOauthBinding.id == binding_id,128                DataSourceOauthBinding.disabled == False,129            )130        ).first()131        if data_source_binding:132            # get all authorized pages133            pages = self.get_authorized_pages(data_source_binding.access_token)134            source_info = data_source_binding.source_info135            new_source_info = {136                "workspace_name": source_info["workspace_name"],137                "workspace_icon": source_info["workspace_icon"],138                "workspace_id": source_info["workspace_id"],139                "pages": pages,140                "total": len(pages),141            }142            data_source_binding.source_info = new_source_info143            data_source_binding.disabled = False144            data_source_binding.updated_at = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)145            db.session.commit()146        else:147            raise ValueError("Data source binding not found")148 149    def get_authorized_pages(self, access_token: str):150        pages = []151        page_results = self.notion_page_search(access_token)152        database_results = self.notion_database_search(access_token)153        # get page detail154        for page_result in page_results:155            page_id = page_result["id"]156            page_name = "Untitled"157            for key in page_result["properties"]:158                if "title" in page_result["properties"][key] and page_result["properties"][key]["title"]:159                    title_list = page_result["properties"][key]["title"]160                    if len(title_list) > 0 and "plain_text" in title_list[0]:161                        page_name = title_list[0]["plain_text"]162            page_icon = page_result["icon"]163            if page_icon:164                icon_type = page_icon["type"]165                if icon_type in {"external", "file"}:166                    url = page_icon[icon_type]["url"]167                    icon = {"type": "url", "url": url if url.startswith("http") else f"https://www.notion.so{url}"}168                else:169                    icon = {"type": "emoji", "emoji": page_icon[icon_type]}170            else:171                icon = None172            parent = page_result["parent"]173            parent_type = parent["type"]174            if parent_type == "block_id":175                parent_id = self.notion_block_parent_page_id(access_token, parent[parent_type])176            elif parent_type == "workspace":177                parent_id = "root"178            else:179                parent_id = parent[parent_type]180            page = {181                "page_id": page_id,182                "page_name": page_name,183                "page_icon": icon,184                "parent_id": parent_id,185                "type": "page",186            }187            pages.append(page)188            # get database detail189        for database_result in database_results:190            page_id = database_result["id"]191            if len(database_result["title"]) > 0:192                page_name = database_result["title"][0]["plain_text"]193            else:194                page_name = "Untitled"195            page_icon = database_result["icon"]196            if page_icon:197                icon_type = page_icon["type"]198                if icon_type in {"external", "file"}:199                    url = page_icon[icon_type]["url"]200                    icon = {"type": "url", "url": url if url.startswith("http") else f"https://www.notion.so{url}"}201                else:202                    icon = {"type": icon_type, icon_type: page_icon[icon_type]}203            else:204                icon = None205            parent = database_result["parent"]206            parent_type = parent["type"]207            if parent_type == "block_id":208                parent_id = self.notion_block_parent_page_id(access_token, parent[parent_type])209            elif parent_type == "workspace":210                parent_id = "root"211            else:212                parent_id = parent[parent_type]213            page = {214                "page_id": page_id,215                "page_name": page_name,216                "page_icon": icon,217                "parent_id": parent_id,218                "type": "database",219            }220            pages.append(page)221        return pages222 223    def notion_page_search(self, access_token: str):224        data = {"filter": {"value": "page", "property": "object"}}225        headers = {226            "Content-Type": "application/json",227            "Authorization": f"Bearer {access_token}",228            "Notion-Version": "2022-06-28",229        }230        response = requests.post(url=self._NOTION_PAGE_SEARCH, json=data, headers=headers)231        response_json = response.json()232        results = response_json.get("results", [])233        return results234 235    def notion_block_parent_page_id(self, access_token: str, block_id: str):236        headers = {237            "Authorization": f"Bearer {access_token}",238            "Notion-Version": "2022-06-28",239        }240        response = requests.get(url=f"{self._NOTION_BLOCK_SEARCH}/{block_id}", headers=headers)241        response_json = response.json()242        parent = response_json["parent"]243        parent_type = parent["type"]244        if parent_type == "block_id":245            return self.notion_block_parent_page_id(access_token, parent[parent_type])246        return parent[parent_type]247 248    def notion_workspace_name(self, access_token: str):249        headers = {250            "Authorization": f"Bearer {access_token}",251            "Notion-Version": "2022-06-28",252        }253        response = requests.get(url=self._NOTION_BOT_USER, headers=headers)254        response_json = response.json()255        if "object" in response_json and response_json["object"] == "user":256            user_type = response_json["type"]257            user_info = response_json[user_type]258            if "workspace_name" in user_info:259                return user_info["workspace_name"]260        return "workspace"261 262    def notion_database_search(self, access_token: str):263        data = {"filter": {"value": "database", "property": "object"}}264        headers = {265            "Content-Type": "application/json",266            "Authorization": f"Bearer {access_token}",267            "Notion-Version": "2022-06-28",268        }269        response = requests.post(url=self._NOTION_PAGE_SEARCH, json=data, headers=headers)270        response_json = response.json()271        results = response_json.get("results", [])272        return results273