fix(schema-migration): more robust migration of background-color & text-color attributes

packages/core/src/extensions/Collaboration/schemaMigration/SchemaMigrationPlugin.ts

@@ -1,5 +1,4 @@
 import { Plugin, PluginKey } from "@tiptap/pm/state";
-import { ySyncPluginKey } from "y-prosemirror";
 import * as Y from "yjs";
 
 import { BlockNoteExtension } from "../../../editor/BlockNoteExtension.js";
@@ -31,8 +30,12 @@ export class SchemaMigrationPlugin extends BlockNoteExtension {
           }
 
           if (
-            transactions.length !== 1 ||
-            !transactions[0].getMeta(ySyncPluginKey)
+            // If any of the transactions are not due to a yjs sync, we don't need to run the migration
+            !transactions.some((tr) => tr.getMeta("y-sync$")) ||
+            // If none of the transactions result in a document change, we don't need to run the migration
+            transactions.every((tr) => !tr.docChanged) ||
+            // If the fragment is still empty, we can't run the migration (since it has not yet been applied to the Y.Doc)
+            !fragment.firstChild
           ) {
             return undefined;
           }
@@ -44,6 +47,10 @@ export class SchemaMigrationPlugin extends BlockNoteExtension {
 
           this.migrationDone = true;
 
+          if (!tr.docChanged) {
+            return undefined;
+          }
+
           return tr;
         },
       }),

packages/core/src/extensions/Collaboration/schemaMigration/migrationRules/moveColorAttributes.test.ts

@@ -0,0 +1,130 @@
+import { expect, it } from "vitest";
+import * as Y from "yjs";
+import { BlockNoteEditor } from "../../../../editor/BlockNoteEditor.js";
+import { moveColorAttributes } from "./moveColorAttributes.js";
+import { prosemirrorJSONToYXmlFragment } from "y-prosemirror";
+
+it("can move color attributes on older documents", async () => {
+  const doc = new Y.Doc();
+  const fragment = doc.getXmlFragment("doc");
+  const editor = BlockNoteEditor.create({
+    initialContent: [
+      {
+        type: "paragraph",
+        content: "Welcome to this demo!",
+      },
+    ],
+  });
+
+  // Because this was a previous schema, we are creating the YFragment manually
+  const blockGroup = new Y.XmlElement("blockGroup");
+  const el = new Y.XmlElement("blockContainer");
+  el.setAttribute("id", "0");
+  el.setAttribute("backgroundColor", "red");
+  el.setAttribute("textColor", "blue");
+  const para = new Y.XmlElement("paragraph");
+  para.setAttribute("textAlignment", "left");
+  para.insert(0, [new Y.XmlText("Welcome to this demo!")]);
+  el.insert(0, [para]);
+  blockGroup.insert(0, [el]);
+  fragment.insert(0, [blockGroup]);
+
+  // Note that the blockContainer has the color attributes, but the paragraph does not.
+  expect(fragment.toJSON()).toMatchInlineSnapshot(
+    `"<blockgroup><blockcontainer backgroundColor="red" id="0" textColor="blue"><paragraph textAlignment="left">Welcome to this demo!</paragraph></blockcontainer></blockgroup>"`,
+  );
+
+  const tr = editor.prosemirrorState.tr;
+  moveColorAttributes(fragment, tr);
+  // Note that the color attributes have been moved to the paragraph.
+  expect(JSON.stringify(tr.doc.toJSON())).toMatchInlineSnapshot(
+    `"{"type":"doc","content":[{"type":"blockGroup","content":[{"type":"blockContainer","attrs":{"id":"0"},"content":[{"type":"paragraph","attrs":{"backgroundColor":"red","textColor":"blue","textAlignment":"left"},"content":[{"type":"text","text":"Welcome to this demo!"}]}]}]}]}"`,
+  );
+});
+
+it("does not move color attributes on newer documents", async () => {
+  const doc = new Y.Doc();
+  const fragment = doc.getXmlFragment("doc");
+  const editor = BlockNoteEditor.create({
+    initialContent: [
+      {
+        type: "paragraph",
+        content: "Welcome to this demo!",
+        props: {
+          backgroundColor: "red",
+          textColor: "blue",
+          // Set to non-default value to ensure it is not overridden by the migration rule.
+          textAlignment: "right",
+        },
+      },
+    ],
+  });
+
+  prosemirrorJSONToYXmlFragment(
+    editor.pmSchema,
+    JSON.parse(JSON.stringify(editor.prosemirrorState.doc.toJSON())),
+    fragment,
+  );
+
+  expect(fragment.toJSON()).toMatchInlineSnapshot(
+    // The color attributes are on the paragraph, not the blockContainer.
+    `"<blockgroup><blockcontainer id="0"><paragraph backgroundColor="red" textAlignment="right" textColor="blue">Welcome to this demo!</paragraph></blockcontainer></blockgroup>"`,
+  );
+
+  const tr = editor.prosemirrorState.tr;
+  moveColorAttributes(fragment, tr);
+  // The document will be unchanged because the color attributes are already on the paragraph.
+  expect(tr.docChanged).toBe(false);
+});
+
+it("can move color attributes on older documents multiple times", async () => {
+  const doc = new Y.Doc();
+  const fragment = doc.getXmlFragment("doc");
+  const editor = BlockNoteEditor.create({
+    initialContent: [
+      {
+        type: "paragraph",
+        content: "Welcome to this demo!",
+      },
+    ],
+  });
+
+  // Because this was a previous schema, we are creating the YFragment manually
+  const blockGroup = new Y.XmlElement("blockGroup");
+  const el = new Y.XmlElement("blockContainer");
+  el.setAttribute("id", "0");
+  el.setAttribute("backgroundColor", "red");
+  el.setAttribute("textColor", "blue");
+  const para = new Y.XmlElement("paragraph");
+  para.setAttribute("textAlignment", "left");
+  para.insert(0, [new Y.XmlText("Welcome to this demo!")]);
+  el.insert(0, [para]);
+  blockGroup.insert(0, [el]);
+  fragment.insert(0, [blockGroup]);
+
+  // Note that the blockContainer has the color attributes, but the paragraph does not.
+  expect(fragment.toJSON()).toMatchInlineSnapshot(
+    `"<blockgroup><blockcontainer backgroundColor="red" id="0" textColor="blue"><paragraph textAlignment="left">Welcome to this demo!</paragraph></blockcontainer></blockgroup>"`,
+  );
+
+  const tr = editor.prosemirrorState.tr;
+  moveColorAttributes(fragment, tr);
+  // Note that the color attributes have been moved to the paragraph.
+  expect(JSON.stringify(tr.doc.toJSON())).toMatchInlineSnapshot(
+    `"{"type":"doc","content":[{"type":"blockGroup","content":[{"type":"blockContainer","attrs":{"id":"0"},"content":[{"type":"paragraph","attrs":{"backgroundColor":"red","textColor":"blue","textAlignment":"left"},"content":[{"type":"text","text":"Welcome to this demo!"}]}]}]}]}"`,
+  );
+
+  el.setAttribute("backgroundColor", "green");
+  el.setAttribute("textColor", "yellow");
+
+  expect(fragment.toJSON()).toMatchInlineSnapshot(
+    `"<blockgroup><blockcontainer backgroundColor="green" id="0" textColor="yellow"><paragraph textAlignment="left">Welcome to this demo!</paragraph></blockcontainer></blockgroup>"`,
+  );
+
+  const nextTr = editor.prosemirrorState.tr;
+  moveColorAttributes(fragment, nextTr);
+  // Note that the color attributes have been moved to the paragraph.
+  expect(JSON.stringify(nextTr.doc.toJSON())).toMatchInlineSnapshot(
+    `"{"type":"doc","content":[{"type":"blockGroup","content":[{"type":"blockContainer","attrs":{"id":"0"},"content":[{"type":"paragraph","attrs":{"backgroundColor":"green","textColor":"yellow","textAlignment":"left"},"content":[{"type":"text","text":"Welcome to this demo!"}]}]}]}]}"`,
+  );
+});

packages/core/src/extensions/Collaboration/schemaMigration/migrationRules/moveColorAttributes.ts

@@ -23,14 +23,13 @@ const traverseElement = (
 export const moveColorAttributes: MigrationRule = (fragment, tr) => {
   // Stores necessary info for all `blockContainer` nodes which still have
   // `textColor` or `backgroundColor` attributes that need to be moved.
-  const targetBlockContainers: Record<
+  const targetBlockContainers: Map<
     string,
     {
-      textColor?: string;
-      backgroundColor?: string;
+      textColor: string | undefined;
+      backgroundColor: string | undefined;
     }
-  > = {};
-
+  > = new Map();
   // Finds all elements which still have `textColor` or `backgroundColor`
   // attributes in the current Yjs fragment.
   fragment.forEach((element) => {
@@ -40,39 +39,53 @@ export const moveColorAttributes: MigrationRule = (fragment, tr) => {
           element.nodeName === "blockContainer" &&
           element.hasAttribute("id")
         ) {
+          const textColor = element.getAttribute("textColor");
+          const backgroundColor = element.getAttribute("backgroundColor");
+
           const colors = {
-            textColor: element.getAttribute("textColor"),
-            backgroundColor: element.getAttribute("backgroundColor"),
+            textColor:
+              textColor === defaultProps.textColor.default
+                ? undefined
+                : textColor,
+            backgroundColor:
+              backgroundColor === defaultProps.backgroundColor.default
+                ? undefined
+                : backgroundColor,
           };
 
-          if (colors.textColor === defaultProps.textColor.default) {
-            colors.textColor = undefined;
-          }
-          if (colors.backgroundColor === defaultProps.backgroundColor.default) {
-            colors.backgroundColor = undefined;
-          }
-
           if (colors.textColor || colors.backgroundColor) {
-            targetBlockContainers[element.getAttribute("id")!] = colors;
+            targetBlockContainers.set(element.getAttribute("id")!, colors);
           }
         }
       });
     }
   });
 
+  if (targetBlockContainers.size === 0) {
+    return false;
+  }
+
   // Appends transactions to add the `textColor` and `backgroundColor`
   // attributes found on each `blockContainer` node to move them to the child
   // `blockContent` node.
   tr.doc.descendants((node, pos) => {
     if (
       node.type.name === "blockContainer" &&
-      targetBlockContainers[node.attrs.id]
+      targetBlockContainers.has(node.attrs.id)
     ) {
-      tr = tr.setNodeMarkup(
-        pos + 1,
-        undefined,
-        targetBlockContainers[node.attrs.id],
-      );
+      const el = tr.doc.nodeAt(pos + 1);
+      if (!el) {
+        throw new Error("No element found");
+      }
+
+      tr.setNodeMarkup(pos + 1, undefined, {
+        // preserve existing attributes
+        ...el.attrs,
+        // add the textColor and backgroundColor attributes
+        ...targetBlockContainers.get(node.attrs.id),
+      });
     }
   });
+
+  return true;
 };