import { DatabaseBlockModel, EmbedLinkedDocModel, EmbedSyncedDocModel, ListBlockModel, ParagraphBlockModel, SurfaceRefBlockModel, } from '@blocksuite/affine-model'; import { BlockSuiteError } from '@blocksuite/global/exceptions'; import type { AfterImportBlockPayload, BeforeImportBlockPayload, DeltaOperation, TransformerMiddleware, } from '@blocksuite/store'; import { filter, map } from 'rxjs'; import { matchModels } from '../../utils'; export const replaceIdMiddleware = (idGenerator: () => string): TransformerMiddleware => ({ slots, docCRUD }) => { const idMap = new Map(); // After Import const afterImportBlock$ = slots.afterImport.pipe( filter( (payload): payload is AfterImportBlockPayload => payload.type === 'block' ), map(({ model }) => model) ); afterImportBlock$ .pipe(filter(model => matchModels(model, [DatabaseBlockModel]))) .subscribe(model => { Object.keys(model.props.cells).forEach(cellId => { if (idMap.has(cellId)) { model.props.cells[idMap.get(cellId)!] = model.props.cells[cellId]; delete model.props.cells[cellId]; } }); }); // replace LinkedPage pageId with new id in paragraph blocks afterImportBlock$ .pipe( filter(model => matchModels(model, [ParagraphBlockModel, ListBlockModel]) ) ) .subscribe(model => { let prev = 0; const delta: DeltaOperation[] = []; for (const d of model.props.text.toDelta()) { if (d.attributes?.reference?.pageId) { const newId = idMap.get(d.attributes.reference.pageId); if (!newId) { prev += d.insert?.length ?? 0; continue; } if (prev > 0) { delta.push({ retain: prev }); } delta.push({ retain: d.insert?.length ?? 0, attributes: { reference: { ...d.attributes.reference, pageId: newId, }, }, }); prev = 0; } else { prev += d.insert?.length ?? 0; } } if (delta.length > 0) { model.props.text.applyDelta(delta); } }); afterImportBlock$ .pipe(filter(model => matchModels(model, [SurfaceRefBlockModel]))) .subscribe(model => { const original = model.props.reference; // If there exists a replacement, replace the reference with the new id. // Otherwise, // 1. If the reference is an affine:frame not in doc, generate a new id. // 2. If the reference is graph, keep the original id. if (idMap.has(original)) { model.props.reference = idMap.get(original)!; } else if ( model.props.refFlavour === 'affine:frame' && !model.store.hasBlock(original) ) { const newId = idGenerator(); idMap.set(original, newId); model.props.reference = newId; } }); // TODO(@fundon): process linked block/element afterImportBlock$ .pipe( filter(model => matchModels(model, [EmbedLinkedDocModel, EmbedSyncedDocModel]) ) ) .subscribe(model => { const original = model.props.pageId; // If the pageId is not in the doc, generate a new id. // If we already have a replacement, use it. if (!docCRUD.get(original)) { if (idMap.has(original)) { model.props.pageId = idMap.get(original)!; } else { const newId = idGenerator(); idMap.set(original, newId); model.props.pageId = newId; } } }); // Before Import slots.beforeImport .pipe(filter(payload => payload.type === 'page')) .subscribe(payload => { if (idMap.has(payload.snapshot.meta.id)) { payload.snapshot.meta.id = idMap.get(payload.snapshot.meta.id)!; return; } const newId = idGenerator(); idMap.set(payload.snapshot.meta.id, newId); payload.snapshot.meta.id = newId; }); slots.beforeImport .pipe( filter( (payload): payload is BeforeImportBlockPayload => payload.type === 'block' ) ) .subscribe(payload => { const { snapshot } = payload; if (snapshot.flavour === 'affine:page') { const index = snapshot.children.findIndex( c => c.flavour === 'affine:surface' ); if (index !== -1) { const [surface] = snapshot.children.splice(index, 1); snapshot.children.push(surface); } } const original = snapshot.id; let newId: string; if (idMap.has(original)) { newId = idMap.get(original)!; } else { newId = idGenerator(); idMap.set(original, newId); } snapshot.id = newId; if (snapshot.flavour === 'affine:surface') { // Generate new IDs for images and frames in advance. snapshot.children.forEach(child => { const original = child.id; if (idMap.has(original)) { newId = idMap.get(original)!; } else { newId = idGenerator(); idMap.set(original, newId); } }); Object.entries( snapshot.props.elements as Record> ).forEach(([_, value]) => { switch (value.type) { case 'connector': { let connection = value.source as Record; if (idMap.has(connection.id)) { const newId = idMap.get(connection.id); if (!newId) { throw new BlockSuiteError( BlockSuiteError.ErrorCode.TransformerError, `reference id must exist: ${connection.id}` ); } connection.id = newId; } connection = value.target as Record; if (idMap.has(connection.id)) { const newId = idMap.get(connection.id); if (!newId) { throw new BlockSuiteError( BlockSuiteError.ErrorCode.TransformerError, `reference id must exist: ${connection.id}` ); } connection.id = newId; } break; } case 'group': { const json = (value.children as Record) .json as Record; Object.entries(json).forEach(([key, value]) => { if (idMap.has(key)) { delete json[key]; const newKey = idMap.get(key); if (!newKey) { throw new BlockSuiteError( BlockSuiteError.ErrorCode.TransformerError, `reference id must exist: ${key}` ); } json[newKey] = value; } }); break; } default: break; } }); } }); };