Fix/notion sync (#1258)

api/controllers/console/auth/data_source_oauth.py

@@ -45,10 +45,10 @@ class OAuthDataSource(Resource):
         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_WEB_URL")}?oauth_data_source=success')
+            return { 'data': '' }
         else:
             auth_url = oauth_provider.get_authorization_url()
-            return redirect(auth_url)
+            return { 'data': auth_url }, 200
 
 
 
@@ -62,6 +62,25 @@ class OAuthDataSourceCallback(Resource):
             return {'error': 'Invalid provider'}, 400
         if 'code' in request.args:
             code = request.args.get('code')
+
+            return redirect(f'{current_app.config.get("CONSOLE_WEB_URL")}?type=notion&code={code}')
+        elif 'error' in request.args:
+            error = request.args.get('error')
+
+            return redirect(f'{current_app.config.get("CONSOLE_WEB_URL")}?type=notion&error={error}')
+        else:
+            return redirect(f'{current_app.config.get("CONSOLE_WEB_URL")}?type=notion&error=Access denied')
+        
+
+class OAuthDataSourceBinding(Resource):
+    def get(self, provider: str):
+        OAUTH_DATASOURCE_PROVIDERS = get_oauth_providers()
+        with current_app.app_context():
+            oauth_provider = OAUTH_DATASOURCE_PROVIDERS.get(provider)
+        if not oauth_provider:
+            return {'error': 'Invalid provider'}, 400
+        if 'code' in request.args:
+            code = request.args.get('code')
             try:
                 oauth_provider.get_access_token(code)
             except requests.exceptions.HTTPError as e:
@@ -69,12 +88,7 @@ class OAuthDataSourceCallback(Resource):
                     f"An error occurred during the OAuthCallback process with {provider}: {e.response.text}")
                 return {'error': 'OAuth data source process failed'}, 400
 
-            return redirect(f'{current_app.config.get("CONSOLE_WEB_URL")}?oauth_data_source=success')
-        elif 'error' in request.args:
-            error = request.args.get('error')
-            return redirect(f'{current_app.config.get("CONSOLE_WEB_URL")}?oauth_data_source={error}')
-        else:
-            return redirect(f'{current_app.config.get("CONSOLE_WEB_URL")}?oauth_data_source=access_denied')
+            return {'result': 'success'}, 200
 
 
 class OAuthDataSourceSync(Resource):
@@ -101,4 +115,5 @@ class OAuthDataSourceSync(Resource):
 
 api.add_resource(OAuthDataSource, '/oauth/data-source/<string:provider>')
 api.add_resource(OAuthDataSourceCallback, '/oauth/data-source/callback/<string:provider>')
+api.add_resource(OAuthDataSourceBinding, '/oauth/data-source/binding/<string:provider>')
 api.add_resource(OAuthDataSourceSync, '/oauth/data-source/<string:provider>/<uuid:binding_id>/sync')

web/app/components/header/account-setting/data-source-page/data-source-notion/index.tsx

@@ -1,14 +1,15 @@
+import { useEffect, useState } from 'react'
+import useSWR from 'swr'
 import { useTranslation } from 'react-i18next'
-import Link from 'next/link'
 import { PlusIcon } from '@heroicons/react/24/solid'
 import cn from 'classnames'
 import Indicator from '../../../indicator'
 import Operate from './operate'
 import s from './style.module.css'
 import NotionIcon from '@/app/components/base/notion-icon'
-import { apiPrefix } from '@/config'
 import type { DataSourceNotion as TDataSourceNotion } from '@/models/common'
 import { useAppContext } from '@/context/app-context'
+import { fetchNotionConnection } from '@/service/common'
 
 type DataSourceNotionProps = {
   workspaces: TDataSourceNotion[]
@@ -18,9 +19,30 @@ const DataSourceNotion = ({
 }: DataSourceNotionProps) => {
   const { t } = useTranslation()
   const { isCurrentWorkspaceManager } = useAppContext()
+  const [canConnectNotion, setCanConnectNotion] = useState(false)
+  const { data } = useSWR(canConnectNotion ? '/oauth/data-source/notion' : null, fetchNotionConnection)
 
   const connected = !!workspaces.length
 
+  const handleConnectNotion = () => {
+    if (!isCurrentWorkspaceManager)
+      return
+
+    setCanConnectNotion(true)
+  }
+
+  const handleAuthAgain = () => {
+    if (data?.data)
+      window.location.href = data.data
+    else
+      setCanConnectNotion(true)
+  }
+
+  useEffect(() => {
+    if (data?.data)
+      window.location.href = data.data
+  }, [data])
+
   return (
     <div className='mb-2 border-[0.5px] border-gray-200 bg-gray-50 rounded-xl'>
       <div className='flex items-center px-3 py-[9px]'>
@@ -40,26 +62,28 @@ const DataSourceNotion = ({
         {
           connected
             ? (
-              <Link
+              <div
                 className={
                   `flex items-center ml-3 px-3 h-7 bg-white border border-gray-200
                   rounded-md text-xs font-medium text-gray-700
                   ${isCurrentWorkspaceManager ? 'cursor-pointer' : 'grayscale opacity-50 cursor-default'}`
                 }
-                href={isCurrentWorkspaceManager ? `${apiPrefix}/oauth/data-source/notion` : '/'}>
+                onClick={handleConnectNotion}
+              >
                 {t('common.dataSource.connect')}
-              </Link>
+              </div>
             )
             : (
-              <Link
-                href={isCurrentWorkspaceManager ? `${apiPrefix}/oauth/data-source/notion` : '/' }
+              <div
                 className={
                   `flex items-center px-3 h-7 bg-white border-[0.5px] border-gray-200 text-xs font-medium text-primary-600 rounded-md
                   ${isCurrentWorkspaceManager ? 'cursor-pointer' : 'grayscale opacity-50 cursor-default'}`
-                }>
+                }
+                onClick={handleConnectNotion}
+              >
                 <PlusIcon className='w-[14px] h-[14px] mr-[5px]' />
                 {t('common.dataSource.notion.addWorkspace')}
-              </Link>
+              </div>
             )
         }
       </div>
@@ -98,7 +122,7 @@ const DataSourceNotion = ({
                     }
                   </div>
                   <div className='mr-2 w-[1px] h-3 bg-gray-100' />
-                  <Operate workspace={workspace} />
+                  <Operate workspace={workspace} onAuthAgain={handleAuthAgain} />
                 </div>
               ))
             }

web/app/components/header/account-setting/data-source-page/data-source-notion/operate/index.tsx

@@ -1,11 +1,9 @@
 'use client'
 import { useTranslation } from 'react-i18next'
 import { Fragment } from 'react'
-import Link from 'next/link'
 import { useSWRConfig } from 'swr'
 import { EllipsisHorizontalIcon } from '@heroicons/react/24/solid'
 import { Menu, Transition } from '@headlessui/react'
-import { apiPrefix } from '@/config'
 import { syncDataSourceNotion, updateDataSourceNotionAction } from '@/service/common'
 import Toast from '@/app/components/base/toast'
 import type { DataSourceNotion } from '@/models/common'
@@ -15,9 +13,11 @@ import { Trash03 } from '@/app/components/base/icons/src/vender/line/general'
 
 type OperateProps = {
   workspace: DataSourceNotion
+  onAuthAgain: () => void
 }
 export default function Operate({
   workspace,
+  onAuthAgain,
 }: OperateProps) {
   const itemClassName = `
     flex px-3 py-2 hover:bg-gray-50 text-sm text-gray-700
@@ -71,9 +71,10 @@ export default function Operate({
               >
                 <div className="px-1 py-1">
                   <Menu.Item>
-                    <Link
+                    <div
                       className={itemClassName}
-                      href={`${apiPrefix}/oauth/data-source/notion`}>
+                      onClick={onAuthAgain}
+                    >
                       <FilePlus02 className={itemIconClassName} />
                       <div>
                         <div className='leading-5'>{t('common.dataSource.notion.changeAuthorizedPages')}</div>
@@ -81,7 +82,7 @@ export default function Operate({
                           {workspace.source_info.total} {t('common.dataSource.notion.pagesAuthorized')}
                         </div>
                       </div>
-                    </Link>
+                    </div>
                   </Menu.Item>
                   <Menu.Item>
                     <div className={itemClassName} onClick={handleSync}>

web/hooks/use-pay.tsx

@@ -7,7 +7,10 @@ import useSWR from 'swr'
 import { useContext } from 'use-context-selector'
 import I18n from '@/context/i18n'
 import { ProviderEnum } from '@/app/components/header/account-setting/model-page/declarations'
-import { fetchFreeQuotaVerify } from '@/service/common'
+import {
+  fetchDataSourceNotionBinding,
+  fetchFreeQuotaVerify,
+} from '@/service/common'
 import type { ConfirmCommonProps } from '@/app/components/base/confirm/common'
 import Confirm from '@/app/components/base/confirm/common'
 
@@ -92,19 +95,56 @@ export const useCheckFreeQuota = () => {
     : null
 }
 
+export const useCheckNotion = () => {
+  const router = useRouter()
+  const [confirm, setConfirm] = useState<ConfirmType | null>(null)
+  const [canBinding, setCanBinding] = useState(false)
+  const searchParams = useSearchParams()
+  const type = searchParams.get('type')
+  const notionCode = searchParams.get('code')
+  const notionError = searchParams.get('error')
+  const { data } = useSWR(
+    canBinding
+      ? `/oauth/data-source/binding/notion?code=${notionCode}`
+      : null,
+    fetchDataSourceNotionBinding,
+  )
+
+  useEffect(() => {
+    if (data)
+      router.replace('/', { forceOptimisticNavigation: false })
+  }, [data, router])
+  useEffect(() => {
+    if (type === 'notion') {
+      if (notionError) {
+        setConfirm({
+          type: 'danger',
+          title: notionError,
+        })
+      }
+      else if (notionCode) {
+        setCanBinding(true)
+      }
+    }
+  }, [type, notionCode, notionError])
+
+  return confirm
+}
+
 export const CheckModal = () => {
   const router = useRouter()
   const { t } = useTranslation()
   const [showPayStatusModal, setShowPayStatusModal] = useState(true)
   const anthropicConfirmInfo = useAnthropicCheckPay()
   const freeQuotaConfirmInfo = useCheckFreeQuota()
+  const notionConfirmInfo = useCheckNotion()
 
   const handleCancelShowPayStatusModal = useCallback(() => {
     setShowPayStatusModal(false)
     router.replace('/', { forceOptimisticNavigation: false })
   }, [router])
 
-  const confirmInfo = anthropicConfirmInfo || freeQuotaConfirmInfo
+  const confirmInfo = anthropicConfirmInfo || freeQuotaConfirmInfo || notionConfirmInfo
 
   if (!confirmInfo || !showPayStatusModal)
     return null

web/service/common.ts

@@ -188,3 +188,11 @@ export const fetchDocumentsLimit: Fetcher<DocumentsLimitResponse, string> = (url
 export const fetchFreeQuotaVerify: Fetcher<{ result: string; flag: boolean; reason: string }, string> = (url) => {
   return get(url) as Promise<{ result: string; flag: boolean; reason: string }>
 }
+
+export const fetchNotionConnection: Fetcher<{ data: string }, string> = (url) => {
+  return get(url) as Promise<{ data: string }>
+}
+
+export const fetchDataSourceNotionBinding: Fetcher<{ result: string }, string> = (url) => {
+  return get(url) as Promise<{ result: string }>
+}