refactor: init hook useRouterTargetWorkspace (#1127)

apps/web/src/components/workspace-layout/index.tsx

@@ -1,18 +1,14 @@
+import { useGlobalState } from '@affine/store';
 import { useRouter } from 'next/router';
-import { PropsWithChildren } from 'react';
+import { PropsWithChildren, useEffect } from 'react';
 
 import HelpIsland from '@/components/help-island';
 import { WorkSpaceSliderBar } from '@/components/workspace-slider-bar';
-import useEnsureWorkspace from '@/hooks/use-ensure-workspace';
+import { useRouterTargetWorkspace } from '@/hooks/use-router-target-workspace';
 
 import { PageLoading } from '../loading';
 import { StyledPage, StyledToolWrapper, StyledWrapper } from './styles';
 
-export const WorkspaceDefender = ({ children }: PropsWithChildren) => {
-  const { workspaceLoaded } = useEnsureWorkspace();
-  return <>{workspaceLoaded ? children : <PageLoading />}</>;
-};
-
 export const WorkspaceLayout = ({ children }: PropsWithChildren) => {
   const router = useRouter();
 
@@ -35,10 +31,22 @@ export const WorkspaceLayout = ({ children }: PropsWithChildren) => {
 };
 
 export const Layout = ({ children }: PropsWithChildren) => {
-  return (
-    <WorkspaceDefender>
-      <WorkspaceLayout>{children}</WorkspaceLayout>
-    </WorkspaceDefender>
-  );
+  const { targetWorkspace, exist } = useRouterTargetWorkspace();
+  const router = useRouter();
+  const loadWorkspace = useGlobalState(store => store.loadWorkspace);
+  useEffect(() => {
+    if (!exist) {
+      router.replace('/404');
+    }
+  }, [exist, router]);
+  useEffect(() => {
+    if (exist && targetWorkspace) {
+      loadWorkspace(targetWorkspace.id);
+    }
+  }, [exist, loadWorkspace, targetWorkspace]);
+  if (!targetWorkspace) {
+    return <PageLoading />;
+  }
+  return <WorkspaceLayout>{children}</WorkspaceLayout>;
 };
 export default Layout;

apps/web/src/hooks/use-ensure-workspace.ts

@@ -1,73 +0,0 @@
-import { assertEquals } from '@blocksuite/global/utils';
-import { useRouter } from 'next/router';
-import { useCallback, useEffect, useState } from 'react';
-
-import { useDataCenter, useGlobalState } from '@/store/app';
-
-// todo: refactor with suspense mode
-// It is a fully effective hook
-// Cause it not just ensure workspace loaded, but also have router change.
-export const useEnsureWorkspace = () => {
-  const dataCenter = useDataCenter();
-  const currentWorkspace = useGlobalState(
-    useCallback(store => store.currentDataCenterWorkspace, [])
-  );
-  const loadWorkspace = useGlobalState(
-    useCallback(store => store.loadWorkspace, [])
-  );
-  const router = useRouter();
-  const [currentWorkspaceId, setCurrentWorkspaceId] = useState<string | null>(
-    typeof router.query.workspaceId === 'string'
-      ? router.query.workspaceId
-      : null
-  );
-
-  // const defaultOutLineWorkspaceId = '99ce7eb7';
-  // console.log(defaultOutLineWorkspaceId);
-  useEffect(() => {
-    const abortController = new AbortController();
-
-    const workspaceId =
-      (router.query.workspaceId as string) || dataCenter.workspaces[0]?.id;
-
-    // If router.query.workspaceId is not in workspace list, jump to 404 page
-    // If workspaceList is empty, we need to create a default workspace but not jump to 404
-    if (
-      workspaceId &&
-      dataCenter.workspaces.length &&
-      dataCenter.workspaces.findIndex(
-        meta => meta.id.toString() === workspaceId
-      ) === -1
-    ) {
-      router.push('/404');
-      return;
-    }
-    // If user is not login and input a custom workspaceId, jump to 404 page
-    // if (
-    //   !user &&
-    //   router.query.workspaceId &&
-    //   router.query.workspaceId !== defaultOutLineWorkspaceId
-    // ) {
-    //   router.push('/404');
-    //   return;
-    // }
-
-    loadWorkspace(workspaceId, abortController.signal).then(unit => {
-      if (!abortController.signal.aborted && unit) {
-        setCurrentWorkspaceId(unit.id);
-        assertEquals(unit.id, workspaceId);
-      }
-    });
-
-    return () => {
-      abortController.abort();
-    };
-  }, [dataCenter, loadWorkspace, router]);
-
-  return {
-    workspaceLoaded: currentWorkspace?.id === currentWorkspaceId,
-    activeWorkspaceId: currentWorkspace?.id ?? router.query.workspaceId,
-  };
-};
-
-export default useEnsureWorkspace;

apps/web/src/hooks/use-page-helper.ts

@@ -5,7 +5,7 @@ import { uuidv4, Workspace } from '@blocksuite/store';
 // eslint-disable-next-line @typescript-eslint/no-restricted-imports
 import type { QueryContent } from '@blocksuite/store/dist/workspace/search';
 import { useRouter } from 'next/router';
-import { useCallback } from 'react';
+import { useCallback, useMemo } from 'react';
 
 import { useChangePageMeta } from '@/hooks/use-change-page-meta';
 import { useGlobalState } from '@/store/app';
@@ -48,104 +48,107 @@ export const usePageHelper = (): EditorHandlers => {
     useCallback(store => store.currentDataCenterWorkspace, [])
   );
 
-  return {
-    createPage: ({
-      pageId = uuidv4().replaceAll('-', ''),
-      title = '',
-    } = {}) => {
-      return new Promise(resolve => {
-        if (!currentWorkspace) {
-          return resolve(null);
-        }
-        currentWorkspace.blocksuiteWorkspace?.createPage(pageId);
-        currentWorkspace.blocksuiteWorkspace?.signals.pageAdded.once(
-          addedPageId => {
-            currentWorkspace.blocksuiteWorkspace?.setPageMeta(addedPageId, {
-              title,
-            });
-            resolve(addedPageId);
+  return useMemo(
+    () => ({
+      createPage: ({
+        pageId = uuidv4().replaceAll('-', ''),
+        title = '',
+      } = {}) => {
+        return new Promise(resolve => {
+          if (!currentWorkspace) {
+            return resolve(null);
           }
-        );
-      });
-    },
-    toggleFavoritePage: async pageId => {
-      const pageMeta = getPageMeta(currentWorkspace, pageId);
-      if (!pageMeta) {
-        return Promise.reject('No page');
-      }
-      const favorite = !pageMeta.favorite;
-      changePageMeta(pageMeta.id, {
-        favorite,
-      });
-      return favorite;
-    },
-    toggleDeletePage: async pageId => {
-      const pageMeta = getPageMeta(currentWorkspace, pageId);
-
-      if (!pageMeta) {
-        return Promise.reject('No page');
-      }
-      const trash = !pageMeta.trash;
-
-      changePageMeta(pageMeta.id, {
-        trash,
-        trashDate: +new Date(),
-      });
-      return trash;
-    },
-    search: (query: QueryContent, workspace?: Workspace) => {
-      if (workspace) {
-        return workspace.search(query);
-      }
-      if (currentWorkspace) {
-        if (currentWorkspace.blocksuiteWorkspace) {
-          return currentWorkspace.blocksuiteWorkspace.search(query);
+          currentWorkspace.blocksuiteWorkspace?.createPage(pageId);
+          currentWorkspace.blocksuiteWorkspace?.signals.pageAdded.once(
+            addedPageId => {
+              currentWorkspace.blocksuiteWorkspace?.setPageMeta(addedPageId, {
+                title,
+              });
+              resolve(addedPageId);
+            }
+          );
+        });
+      },
+      toggleFavoritePage: async pageId => {
+        const pageMeta = getPageMeta(currentWorkspace, pageId);
+        if (!pageMeta) {
+          return Promise.reject('No page');
         }
-      }
-      return new Map();
-    },
-    changePageMode: async (pageId, mode) => {
-      const pageMeta = getPageMeta(currentWorkspace, pageId);
-      if (!pageMeta) {
-        return Promise.reject('No page');
-      }
+        const favorite = !pageMeta.favorite;
+        changePageMeta(pageMeta.id, {
+          favorite,
+        });
+        return favorite;
+      },
+      toggleDeletePage: async pageId => {
+        const pageMeta = getPageMeta(currentWorkspace, pageId);
 
-      editor?.setAttribute('mode', mode as string);
+        if (!pageMeta) {
+          return Promise.reject('No page');
+        }
+        const trash = !pageMeta.trash;
 
-      changePageMeta(pageMeta.id, {
-        mode,
-      });
-      return mode;
-    },
-    permanentlyDeletePage: pageId => {
-      // TODO:  workspace.meta.removePage or workspace.removePage?
+        changePageMeta(pageMeta.id, {
+          trash,
+          trashDate: +new Date(),
+        });
+        return trash;
+      },
+      search: (query: QueryContent, workspace?: Workspace) => {
+        if (workspace) {
+          return workspace.search(query);
+        }
+        if (currentWorkspace) {
+          if (currentWorkspace.blocksuiteWorkspace) {
+            return currentWorkspace.blocksuiteWorkspace.search(query);
+          }
+        }
+        return new Map();
+      },
+      changePageMode: async (pageId, mode) => {
+        const pageMeta = getPageMeta(currentWorkspace, pageId);
+        if (!pageMeta) {
+          return Promise.reject('No page');
+        }
 
-      currentWorkspace!.blocksuiteWorkspace?.meta.removePage(pageId);
-    },
-    openPage: (pageId, query = {}, newTab = false) => {
-      pageId = pageId.replace('space:', '');
+        editor?.setAttribute('mode', mode as string);
 
-      if (newTab) {
-        window.open(`/workspace/${currentWorkspace?.id}/${pageId}`, '_blank');
-        return Promise.resolve(true);
-      }
-      return router.push({
-        pathname: `/workspace/${currentWorkspace?.id}/${pageId}`,
-        query,
-      });
-    },
-    getPageMeta: pageId => {
-      if (!currentWorkspace) {
-        return null;
-      }
+        changePageMeta(pageMeta.id, {
+          mode,
+        });
+        return mode;
+      },
+      permanentlyDeletePage: pageId => {
+        // TODO:  workspace.meta.removePage or workspace.removePage?
 
-      return (
-        (currentWorkspace.blocksuiteWorkspace?.meta.pageMetas.find(
-          page => page.id === pageId
-        ) as PageMeta) || null
-      );
-    },
-  };
+        currentWorkspace!.blocksuiteWorkspace?.meta.removePage(pageId);
+      },
+      openPage: (pageId, query = {}, newTab = false) => {
+        pageId = pageId.replace('space:', '');
+
+        if (newTab) {
+          window.open(`/workspace/${currentWorkspace?.id}/${pageId}`, '_blank');
+          return Promise.resolve(true);
+        }
+        return router.push({
+          pathname: `/workspace/${currentWorkspace?.id}/${pageId}`,
+          query,
+        });
+      },
+      getPageMeta: pageId => {
+        if (!currentWorkspace) {
+          return null;
+        }
+
+        return (
+          (currentWorkspace.blocksuiteWorkspace?.meta.pageMetas.find(
+            page => page.id === pageId
+          ) as PageMeta) || null
+        );
+      },
+    }),
+    [changePageMeta, currentWorkspace, editor, router]
+  );
 };
 
 export default usePageHelper;

apps/web/src/hooks/use-router-target-workspace.ts

@@ -0,0 +1,26 @@
+import { useDataCenter, useDataCenterWorkspace } from '@affine/store';
+import { useRouter } from 'next/router';
+import { useMemo } from 'react';
+
+export function useRouterTargetWorkspace() {
+  const router = useRouter();
+  const dataCenter = useDataCenter();
+  const workspaceId =
+    typeof router.query.workspaceId === 'string'
+      ? router.query.workspaceId
+      : dataCenter.workspaces.at(0)?.id ?? null;
+  const targetWorkspace = useDataCenterWorkspace(workspaceId);
+  const notExist = useMemo(
+    () =>
+      workspaceId &&
+      dataCenter.workspaces.length &&
+      dataCenter.workspaces.findIndex(
+        meta => meta.id.toString() === workspaceId
+      ) === -1,
+    [dataCenter.workspaces, workspaceId]
+  );
+  return {
+    targetWorkspace,
+    exist: !notExist,
+  };
+}

apps/web/src/pages/workspace/[workspaceId]/index.tsx

@@ -1,44 +1,48 @@
 import { useRouter } from 'next/router';
-import { useCallback, useEffect } from 'react';
+import { useEffect } from 'react';
 
 import { PageLoading } from '@/components/loading';
-import useEnsureWorkspace from '@/hooks/use-ensure-workspace';
 import usePageHelper from '@/hooks/use-page-helper';
-import { useGlobalState } from '@/store/app';
+import { useRouterTargetWorkspace } from '@/hooks/use-router-target-workspace';
 
 const WorkspaceIndex = () => {
   const router = useRouter();
-  const currentWorkspace = useGlobalState(
-    useCallback(store => store.currentDataCenterWorkspace, [])
-  );
+  const { targetWorkspace, exist } = useRouterTargetWorkspace();
   const { createPage } = usePageHelper();
-  const { workspaceLoaded, activeWorkspaceId } = useEnsureWorkspace();
 
   useEffect(() => {
+    if (!exist) {
+      router.push('/404');
+      return;
+    }
+    const abortController = new AbortController();
     const initPage = async () => {
-      if (!workspaceLoaded) {
+      if (abortController.signal.aborted) {
+        return;
+      }
+      if (!targetWorkspace) {
         return;
       }
       const savedPageId =
-        currentWorkspace?.blocksuiteWorkspace?.meta.pageMetas.find(
+        targetWorkspace.blocksuiteWorkspace?.meta.pageMetas.find(
           meta => !meta.trash
         )?.id;
       if (savedPageId) {
-        router.replace(`/workspace/${activeWorkspaceId}/${savedPageId}`);
+        router.replace(`/workspace/${targetWorkspace.id}/${savedPageId}`);
         return;
+      } else {
+        const pageId = await createPage();
+        if (abortController.signal.aborted) {
+          return;
+        }
+        router.replace(`/workspace/${targetWorkspace.id}/${pageId}`);
       }
-
-      const pageId = await createPage();
-      router.replace(`/workspace/${activeWorkspaceId}/${pageId}`);
     };
     initPage();
-  }, [
-    currentWorkspace,
-    createPage,
-    router,
-    workspaceLoaded,
-    activeWorkspaceId,
-  ]);
+    return () => {
+      abortController.abort();
+    };
+  }, [targetWorkspace, createPage, router, exist]);
 
   return <PageLoading />;
 };

apps/web/src/pages/workspace/index.tsx

@@ -1,23 +1,19 @@
 import { useRouter } from 'next/router';
-import { useCallback, useEffect } from 'react';
+import { useEffect } from 'react';
 
 import { PageLoading } from '@/components/loading';
-import useEnsureWorkspace from '@/hooks/use-ensure-workspace';
-import { useGlobalState } from '@/store/app';
+import { useRouterTargetWorkspace } from '@/hooks/use-router-target-workspace';
 
 export const WorkspaceIndex = () => {
   const router = useRouter();
-  const currentWorkspace = useGlobalState(
-    useCallback(store => store.currentDataCenterWorkspace, [])
-  );
-  const { workspaceLoaded } = useEnsureWorkspace();
-
+  const { targetWorkspace, exist } = useRouterTargetWorkspace();
   useEffect(() => {
-    if (workspaceLoaded) {
-      router.push(`/workspace/${currentWorkspace?.id}`);
+    if (!exist) {
+      router.push('/404');
+    } else if (targetWorkspace) {
+      router.push(`/workspace/${targetWorkspace.id}`);
     }
-  }, [currentWorkspace, router, workspaceLoaded]);
-
+  }, [targetWorkspace, exist, router]);
   return <PageLoading />;
 };
 

packages/store/src/app/datacenter/index.tsx

@@ -116,6 +116,15 @@ export function useDataCenter() {
   return data as DataCenter;
 }
 
+export function useDataCenterWorkspace(
+  workspaceId: string | null
+): WorkspaceUnit | null {
+  const { data } = useSWR<WorkspaceUnit | null>(['datacenter', workspaceId], {
+    fallbackData: null,
+  });
+  return data ?? null;
+}
+
 export function DataCenterPreloader({ children }: React.PropsWithChildren) {
   const api = useGlobalStateApi();
   //# region effect for updating workspace page list

packages/store/src/app/index.tsx

@@ -84,7 +84,7 @@ export const useGlobalState: UseBoundStore<Store> = ((
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
 }) as any;
 
-export type DataKey = ['datacenter', string] | ['datacenter'];
+export type DataKey = ['datacenter', string | null] | ['datacenter'];
 
 const swrFetcher = async (keys: DataKey) => {
   assertEquals(keys[0], 'datacenter');
@@ -96,6 +96,9 @@ const swrFetcher = async (keys: DataKey) => {
       return dataCenter;
     });
   } else {
+    if (keys[1] === null) {
+      return null;
+    }
     const dataCenter = await dataCenterPromise;
     return dataCenter.loadWorkspace(keys[1]);
   }

packages/store/src/index.ts

@@ -1,4 +1,8 @@
 export * from './app';
 export type { PageMeta } from './app/datacenter';
 export { createDefaultWorkspace, DataCenterPreloader } from './app/datacenter';
-export { dataCenterPromise, useDataCenter } from './app/datacenter';
+export {
+  dataCenterPromise,
+  useDataCenter,
+  useDataCenterWorkspace,
+} from './app/datacenter';