feat(data-modeling): bring back collapse all under a flag COMPASS-9504

packages/compass-components/src/index.ts

@@ -244,6 +244,7 @@ export type {
   EdgeProps,
   NodeProps,
   DiagramProps,
+  FieldId,
   DiagramInstance,
   NodeField,
   NodeGlyph,

packages/compass-data-modeling/src/components/diagram-card.spec.tsx

@@ -26,6 +26,7 @@ describe('DiagramCard', () => {
                 displayPosition: [0, 0],
                 shardKey: {},
                 jsonSchema: { bsonType: 'object' },
+                isExpanded: true,
               },
             ],
             relationships: [],

packages/compass-data-modeling/src/components/diagram-editor.spec.tsx

@@ -44,13 +44,15 @@ const storageItems: MongoDBDataModelDescription[] = [
               displayPosition: [50, 50],
               shardKey: {},
               jsonSchema: { bsonType: 'object' },
+              isExpanded: true,
             },
             {
               ns: 'db1.collection2',
               indexes: [],
               displayPosition: [150, 150],
               shardKey: {},
               jsonSchema: { bsonType: 'object' },
+              isExpanded: true,
             },
           ],
           relationships: [],
@@ -78,13 +80,15 @@ const storageItems: MongoDBDataModelDescription[] = [
               displayPosition: [0, 0],
               shardKey: {},
               jsonSchema: { bsonType: 'object' },
+              isExpanded: true,
             },
             {
               ns: 'db1.collection2',
               indexes: [],
               displayPosition: [0, 0],
               shardKey: {},
               jsonSchema: { bsonType: 'object' },
+              isExpanded: true,
             },
           ],
           relationships: [],

packages/compass-data-modeling/src/components/diagram-editor.tsx

@@ -23,11 +23,13 @@ import {
   deleteRelationship,
   removeField,
   renameField,
+  toggleCollectionExpanded,
 } from '../store/diagram';
 import type {
   EdgeProps,
   NodeProps,
   DiagramProps,
+  FieldId,
 } from '@mongodb-js/compass-components';
 import {
   Banner,
@@ -58,6 +60,7 @@ import {
   relationshipToDiagramEdge,
 } from '../utils/nodes-and-edges';
 import toNS from 'mongodb-ns';
+import { usePreference } from 'compass-preferences-model/provider';
 
 const loadingContainerStyles = css({
   width: '100%',
@@ -182,6 +185,7 @@ const DiagramContent: React.FunctionComponent<{
   }) => void;
   onRelationshipDrawn: () => void;
   DiagramComponent?: typeof Diagram;
+  onToggleCollectionExpanded: (namespace: string) => void;
 }> = ({
   diagramLabel,
   database,
@@ -204,8 +208,10 @@ const DiagramContent: React.FunctionComponent<{
   onDeleteField,
   selectedItems,
   DiagramComponent = Diagram,
+  onToggleCollectionExpanded,
 }) => {
   const isDarkMode = useDarkMode();
+  const isCollapseFlagEnabled = usePreference('enableDataModelingCollapse');
   const diagram = useRef(useDiagram());
   const { openDrawer } = useDrawerActions();
   const { isDrawerOpen } = useDrawerState();
@@ -239,6 +245,7 @@ const DiagramContent: React.FunctionComponent<{
             : undefined,
         selected,
         isInRelationshipDrawingMode,
+        isExpanded: coll.isExpanded,
       });
     });
   }, [
@@ -331,9 +338,24 @@ const DiagramContent: React.FunctionComponent<{
   );
 
   const onFieldClick = useCallback(
-    (_evt: React.MouseEvent, { id: fieldPath, nodeId: namespace }) => {
+    (
+      _evt: React.MouseEvent,
+      { id, nodeId: namespace }: { id: FieldId; nodeId: string }
+    ) => {
+      // Diagramming package accepts both string ids and array of string ids for
+      // fields (to represent the field path better). While all current code in
+      // compass always uses array of strings as field id, some older saved
+      // diagrams might not. Also handling this explicitly is sort of needed
+      // anyway to convince typescript that we're doing the right thing
+      const fieldPath = Array.isArray(id)
+        ? id
+        : typeof id === 'string'
+        ? [id]
+        : undefined;
+      if (!fieldPath) {
+        return;
+      }
       _evt.stopPropagation(); // TODO(COMPASS-9659): should this be handled by the diagramming package?
-      if (!Array.isArray(fieldPath)) return; // TODO(COMPASS-9659): could be avoided with generics in the diagramming package
       onFieldSelect(namespace, fieldPath);
       openDrawer(DATA_MODELING_DRAWER_ID);
     },
@@ -398,6 +420,15 @@ const DiagramContent: React.FunctionComponent<{
     [onDiagramBackgroundClicked]
   );
 
+  const handleNodeExpandedToggle = useCallback(
+    (evt: React.MouseEvent, nodeId: string) => {
+      evt.preventDefault();
+      evt.stopPropagation();
+      onToggleCollectionExpanded(nodeId);
+    },
+    [onToggleCollectionExpanded]
+  );
+
   const diagramProps: DiagramProps = useMemo(
     () =>
       ({
@@ -414,6 +445,9 @@ const DiagramContent: React.FunctionComponent<{
         onFieldNameChange: onRenameField,
         onNodeDragStop,
         onConnect,
+        onNodeExpandToggle: isCollapseFlagEnabled
+          ? handleNodeExpandedToggle
+          : undefined,
       } satisfies DiagramProps),
     [
       isDarkMode,
@@ -429,6 +463,8 @@ const DiagramContent: React.FunctionComponent<{
       onRenameField,
       onNodeDragStop,
       onConnect,
+      handleNodeExpandedToggle,
+      isCollapseFlagEnabled,
     ]
   );
 
@@ -496,6 +532,7 @@ const ConnectedDiagramContent = connect(
     onDeleteCollection: deleteCollection,
     onDeleteRelationship: deleteRelationship,
     onDeleteField: removeField,
+    onToggleCollectionExpanded: toggleCollectionExpanded,
   }
 )(DiagramContent);
 

packages/compass-data-modeling/src/components/saved-diagrams-list.spec.tsx

@@ -30,6 +30,7 @@ const storageItems: MongoDBDataModelDescription[] = [
               displayPosition: [1, 1],
               shardKey: {},
               jsonSchema: { bsonType: 'object' },
+              isExpanded: true,
             },
           ],
           relationships: [],
@@ -57,6 +58,7 @@ const storageItems: MongoDBDataModelDescription[] = [
               displayPosition: [2, 2],
               shardKey: {},
               jsonSchema: { bsonType: 'object' },
+              isExpanded: true,
             },
           ],
           relationships: [],
@@ -84,6 +86,7 @@ const storageItems: MongoDBDataModelDescription[] = [
               displayPosition: [3, 3],
               shardKey: {},
               jsonSchema: { bsonType: 'object' },
+              isExpanded: true,
             },
           ],
           relationships: [],

packages/compass-data-modeling/src/services/data-model-storage.ts

@@ -27,6 +27,8 @@ export const RelationshipSchema = z.object({
 
 export type Relationship = z.output<typeof RelationshipSchema>;
 
+export const DEFAULT_IS_EXPANDED = true;
+
 const CollectionSchema = z.object({
   ns: z.string(),
   jsonSchema: z.custom<MongoDBJSONSchema>((value) => {
@@ -37,6 +39,7 @@ const CollectionSchema = z.object({
   shardKey: z.record(z.unknown()).optional(),
   displayPosition: z.tuple([z.number(), z.number()]),
   note: z.string().optional(),
+  isExpanded: z.boolean().default(DEFAULT_IS_EXPANDED),
 });
 
 export type DataModelCollection = z.output<typeof CollectionSchema>;
@@ -133,6 +136,10 @@ const EditSchemaVariants = z.discriminatedUnion('type', [
     field: FieldPathSchema,
     jsonSchema: z.custom<MongoDBJSONSchema>(),
   }),
+  z.object({
+    type: z.literal('ToggleExpandCollection'),
+    ns: z.string(),
+  }),
 ]);
 
 export const EditSchema = z.intersection(EditSchemaBase, EditSchemaVariants);

packages/compass-data-modeling/src/store/analysis-process.ts

@@ -4,7 +4,10 @@ import type { DataModelingThunkAction } from './reducer';
 import { analyzeDocuments, type MongoDBJSONSchema } from 'mongodb-schema';
 import { getCurrentDiagramFromState } from './diagram';
 import { UUID } from 'bson';
-import type { Relationship } from '../services/data-model-storage';
+import {
+  DEFAULT_IS_EXPANDED,
+  type Relationship,
+} from '../services/data-model-storage';
 import { applyLayout } from '@mongodb-js/compass-components';
 import {
   collectionToBaseNodeForLayout,
@@ -76,6 +79,7 @@ export type AnalysisFinishedAction = {
     ns: string;
     schema: MongoDBJSONSchema;
     position: { x: number; y: number };
+    isExpanded: boolean;
   }[];
   relations: Relationship[];
 };
@@ -154,7 +158,7 @@ async function getInitialLayout({
   collections,
   relations,
 }: {
-  collections: { ns: string; schema: MongoDBJSONSchema }[];
+  collections: { ns: string; schema: MongoDBJSONSchema; isExpanded: boolean }[];
   relations: Relationship[];
 }) {
   const hasRelations = relations.length > 0;
@@ -163,6 +167,7 @@ async function getInitialLayout({
       ns: coll.ns,
       jsonSchema: coll.schema,
       displayPosition: [0, 0],
+      isExpanded: coll.isExpanded,
     });
   });
   return await applyLayout({
@@ -250,7 +255,7 @@ export function startAnalysis(
             type: AnalysisProcessActionTypes.NAMESPACE_SCHEMA_ANALYZED,
             namespace: ns,
           });
-          return { ns, schema, sample };
+          return { ns, schema, sample, isExpanded: DEFAULT_IS_EXPANDED };
         })
       );
 

packages/compass-data-modeling/src/store/apply-edit.ts

@@ -1,8 +1,9 @@
-import type {
-  DataModelCollection,
-  Edit,
-  Relationship,
-  StaticModel,
+import {
+  DEFAULT_IS_EXPANDED,
+  type DataModelCollection,
+  type Edit,
+  type Relationship,
+  type StaticModel,
 } from '../services/data-model-storage';
 import { updateSchema } from '../utils/schema-traversal';
 import {
@@ -60,6 +61,7 @@ export function applyEdit(edit: Edit, model?: StaticModel): StaticModel {
         jsonSchema: edit.initialSchema,
         displayPosition: edit.position,
         indexes: [],
+        isExpanded: DEFAULT_IS_EXPANDED,
       };
       return {
         ...model,
@@ -271,6 +273,20 @@ export function applyEdit(edit: Edit, model?: StaticModel): StaticModel {
         }),
       };
     }
+    case 'ToggleExpandCollection': {
+      return {
+        ...model,
+        collections: model.collections.map((collection) => {
+          if (collection.ns !== edit.ns) {
+            return collection;
+          }
+          return {
+            ...collection,
+            isExpanded: !collection.isExpanded,
+          };
+        }),
+      };
+    }
     default: {
       return model;
     }

packages/compass-data-modeling/src/store/diagram.spec.ts

@@ -32,13 +32,15 @@ const model: StaticModel = {
       displayPosition: [0, 0],
       shardKey: {},
       jsonSchema: { bsonType: 'object' },
+      isExpanded: true,
     },
     {
       ns: 'db.collection2',
       indexes: [],
       displayPosition: [1, 1],
       shardKey: {},
       jsonSchema: { bsonType: 'object' },
+      isExpanded: true,
     },
   ],
   relationships: [
@@ -92,11 +94,13 @@ describe('Data Modeling store', function () {
             ns: 'db.collection1',
             schema: model.collections[0].jsonSchema,
             position: { x: 0, y: 0 },
+            isExpanded: true,
           },
           {
             ns: 'db.collection2',
             schema: model.collections[1].jsonSchema,
             position: { x: 0, y: 0 },
+            isExpanded: true,
           },
         ],
         relations: model.relationships,
@@ -160,6 +164,7 @@ describe('Data Modeling store', function () {
               displayPosition: [0, 0],
               shardKey: {},
               jsonSchema: { bsonType: 'object' },
+              isExpanded: true,
             },
           ] as StaticModel['collections'],
           relationships: [] as StaticModel['relationships'],
@@ -455,6 +460,7 @@ describe('Data Modeling store', function () {
                     field3: { bsonType: 'int' },
                   },
                 },
+                isExpanded: true,
               },
             ],
             relationships: [],
@@ -481,6 +487,7 @@ describe('Data Modeling store', function () {
                 indexes: [],
                 displayPosition: [0, 0],
                 shardKey: {},
+                isExpanded: true,
                 jsonSchema: {
                   bsonType: 'object',
                   properties: {