import type { CollectionsAtom } from '@affine/component/page-list'; import type { Collection } from '@affine/env/filter'; import { DisposableGroup } from '@blocksuite/global/utils'; import { currentWorkspaceAtom } from '@toeverything/infra/atom'; import { type DBSchema, openDB } from 'idb'; import { atom } from 'jotai'; import { atomWithObservable } from 'jotai/utils'; import { Observable } from 'rxjs'; import type { Map as YMap } from 'yjs'; import { Doc as YDoc } from 'yjs'; import { sessionAtom } from '../atoms/cloud-user'; export interface PageCollectionDBV1 extends DBSchema { view: { key: Collection['id']; value: Collection; }; } export interface StorageCRUD { get: (key: string) => Promise; set: (key: string, value: Value) => Promise; delete: (key: string) => Promise; list: () => Promise; } type Subscribe = () => void; const collectionDBAtom = atom( openDB('page-view', 1, { upgrade(database) { database.createObjectStore('view', { keyPath: 'id', }); }, }) ); const callbackSet = new Set(); const localCollectionCRUDAtom = atom(get => ({ get: async (key: string) => { const db = await get(collectionDBAtom); const t = db.transaction('view').objectStore('view'); return (await t.get(key)) ?? null; }, set: async (key: string, value: Collection) => { const db = await get(collectionDBAtom); const t = db.transaction('view', 'readwrite').objectStore('view'); await t.put(value); callbackSet.forEach(cb => cb()); return key; }, delete: async (key: string) => { const db = await get(collectionDBAtom); const t = db.transaction('view', 'readwrite').objectStore('view'); callbackSet.forEach(cb => cb()); await t.delete(key); }, list: async () => { const db = await get(collectionDBAtom); const t = db.transaction('view').objectStore('view'); return t.getAllKeys(); }, })); const getCollections = async ( storage: StorageCRUD ): Promise => { return storage .list() .then(async keys => { return await Promise.all(keys.map(key => storage.get(key))).then(v => v.filter((v): v is Collection => v !== null) ); }) .catch(error => { console.error('Failed to load collections', error); return []; }); }; const pageCollectionBaseAtom = atomWithObservable(get => { const currentWorkspacePromise = get(currentWorkspaceAtom); const session = get(sessionAtom); const localCRUD = get(localCollectionCRUDAtom); const userId = session?.data?.user.id ?? null; const useLocalStorage = userId === null; return new Observable(subscriber => { // initial value subscriber.next([]); if (useLocalStorage) { getCollections(localCRUD).then(collections => { subscriber.next(collections); }); const fn = () => { getCollections(localCRUD).then(collections => { subscriber.next(collections); }); }; callbackSet.add(fn); return () => { callbackSet.delete(fn); }; } else { const group = new DisposableGroup(); currentWorkspacePromise.then(async currentWorkspace => { const collectionsFromLocal = await getCollections(localCRUD); const rootDoc = currentWorkspace.doc; const settingMap = rootDoc.getMap('settings') as YMap; if (!settingMap.has(userId)) { settingMap.set( userId, new YDoc({ guid: `${rootDoc.guid}:settings:${userId}`, }) ); } const settingDoc = settingMap.get(userId) as YDoc; if (!settingDoc.isLoaded) { settingDoc.load(); await settingDoc.whenLoaded; } const viewMap = settingDoc.getMap('view') as YMap; // sync local storage to doc collectionsFromLocal.map(v => viewMap.set(v.id, v)); // delete from indexeddb Promise.all( collectionsFromLocal.map(async v => { await localCRUD.delete(v.id); }) ).catch(error => { console.error('Failed to delete collections from indexeddb', error); }); const collectionsFromDoc: Collection[] = Array.from(viewMap.keys()) .map(key => viewMap.get(key)) .filter((v): v is Collection => !!v); const collections = [...collectionsFromDoc]; subscriber.next(collections); if (group.disposed) { return; } const fn = () => { const collectionsFromDoc: Collection[] = Array.from(viewMap.keys()) .map(key => viewMap.get(key)) .filter((v): v is Collection => !!v); const collections = [...collectionsFromLocal, ...collectionsFromDoc]; subscriber.next(collections); }; viewMap.observe(fn); group.add(() => { viewMap.unobserve(fn); }); }); return () => { group.dispose(); }; } }); }); export const currentCollectionsAtom: CollectionsAtom = atom( get => get(pageCollectionBaseAtom), async (get, _, apply) => { const collections = await get(pageCollectionBaseAtom); let newCollections: Collection[]; if (typeof apply === 'function') { newCollections = apply(collections); } else { newCollections = apply; } const session = get(sessionAtom); const userId = session?.data?.user.id ?? null; const useLocalStorage = userId === null; const added = newCollections.filter(v => !collections.includes(v)); const removed = collections.filter(v => !newCollections.includes(v)); if (useLocalStorage) { const localCRUD = get(localCollectionCRUDAtom); await Promise.all([ ...added.map(async v => { await localCRUD.set(v.id, v); }), ...removed.map(async v => { await localCRUD.delete(v.id); }), ]); } else { const currentWorkspace = await get(currentWorkspaceAtom); const rootDoc = currentWorkspace.doc; const settingMap = rootDoc.getMap('settings') as YMap; const settingDoc = settingMap.get(userId) as YDoc; const viewMap = settingDoc.getMap('view') as YMap; await Promise.all([ ...added.map(async v => { viewMap.set(v.id, v); }), ...removed.map(async v => { viewMap.delete(v.id); }), ]); } } );