dify/api/libs/oauth_data_source.py

import datetime
import urllib.parse
from typing import Any

import requests
from flask_login import current_user  # type: ignore

from extensions.ext_database import db
from models.source import DataSourceOauthBinding


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"
    _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",
        }
        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,
            "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,
            )
        ).first()
        if data_source_binding:
            data_source_binding.source_info = source_info
            data_source_binding.disabled = False
            data_source_binding.updated_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
            db.session.commit()
        else:
            new_data_source_binding = DataSourceOauthBinding(
                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 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 = DataSourceOauthBinding.query.filter(
            db.and_(
                DataSourceOauthBinding.tenant_id == current_user.current_tenant_id,
                DataSourceOauthBinding.provider == "notion",
                DataSourceOauthBinding.access_token == access_token,
            )
        ).first()
        if data_source_binding:
            data_source_binding.source_info = source_info
            data_source_binding.disabled = False
            data_source_binding.updated_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
            db.session.commit()
        else:
            new_data_source_binding = DataSourceOauthBinding(
                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 = DataSourceOauthBinding.query.filter(
            db.and_(
                DataSourceOauthBinding.tenant_id == current_user.current_tenant_id,
                DataSourceOauthBinding.provider == "notion",
                DataSourceOauthBinding.id == binding_id,
                DataSourceOauthBinding.disabled == False,
            )
        ).first()
        if data_source_binding:
            # get all authorized pages
            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),
            }
            data_source_binding.source_info = new_source_info
            data_source_binding.disabled = False
            data_source_binding.updated_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
            db.session.commit()
        else:
            raise ValueError("Data source binding not found")

    def get_authorized_pages(self, access_token: str):
        pages = []
        page_results = self.notion_page_search(access_token)
        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"]
            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]
            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"]
            else:
                page_name = "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]
            page = {
                "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):
        results = []
        next_cursor = None
        has_more = True

        while has_more:
            data: dict[str, Any] = {
                "filter": {"value": "page", "property": "object"},
                **({"start_cursor": next_cursor} if next_cursor else {}),
            }
            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.extend(response_json.get("results", []))

            has_more = response_json.get("has_more", False)
            next_cursor = response_json.get("next_cursor", None)

        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",
        }
        response = requests.get(url=f"{self._NOTION_BLOCK_SEARCH}/{block_id}", headers=headers)
        response_json = response.json()
        if response.status_code != 200:
            message = response_json.get("message", "unknown error")
            raise ValueError(f"Error fetching block parent page ID: {message}")
        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",
        }
        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):
        results = []
        next_cursor = None
        has_more = True

        while has_more:
            data: dict[str, Any] = {
                "filter": {"value": "database", "property": "object"},
                **({"start_cursor": next_cursor} if next_cursor else {}),
            }
            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.extend(response_json.get("results", []))

            has_more = response_json.get("has_more", False)
            next_cursor = response_json.get("next_cursor", None)

        return results